pub(crate) use super::hasher::{Hasher, Midstate};

use super::hasher::Sink;

use std::fs::File;
use std::io::{BufWriter, Seek as _, SeekFrom, Write as _};
use std::os::unix::fs::FileExt as _;

use anyhow::Context as _;
use knot_types::ObjectFormat;

/// the most pack bytes in a side-band-64k packet: 65520 less the 4 length bytes and the band byte.
pub(crate) const PACKET_DATA: usize = 65515;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Framed {
    prefix: u64,
    pack: u64,
    format: ObjectFormat,
}

const PACKET: u64 = PACKET_DATA as u64 + 5;

impl Framed {
    pub fn of(file: &File, format: ObjectFormat) -> anyhow::Result<Self> {
        let len = file.metadata()?.len();
        let mut prefix = 0u64;
        let mut first = [0u8; 5];
        loop {
            file.read_exact_at(&mut first, prefix)
                .context("the reply ends before its pack")?;
            if first[4] == 1 {
                break;
            }
            prefix = (prefix.checked_add(packet_len(&first[..4])?)).context("a reply too long")?;
        }
        let framing = (len.checked_sub(prefix))
            .and_then(|rest| rest.checked_sub(4))
            .filter(|framing| *framing > 0)
            .context("a reply too short")?;
        let packets = framing.div_ceil(PACKET);
        let framed = Self {
            prefix,
            pack: (framing.checked_sub(5 * packets)).context("a reply too short")?,
            format,
        };
        let mut flush = [0u8; 4];
        file.read_exact_at(&mut flush, len - 4)?;
        anyhow::ensure!(
            framed.pack >= 12 + framed.hash_len()
                && framed.pack > (packets - 1) * PACKET_DATA as u64
                && &flush == b"0000",
            "the reply isn't a pack in side-band packets"
        );
        framed.check(file, 0)?;
        framed.check(file, packets - 1)?;
        Ok(framed)
    }

    fn check(&self, file: &File, at: u64) -> anyhow::Result<()> {
        let mut header = [0u8; 5];
        file.read_exact_at(&mut header, self.prefix + at * PACKET)?;
        anyhow::ensure!(
            header == self.header(at),
            "the reply isn't in whole side-band packets"
        );
        Ok(())
    }

    fn header(&self, at: u64) -> [u8; 5] {
        let data = (self.pack - at * PACKET_DATA as u64).min(PACKET_DATA as u64);
        let mut header = [1u8; 5];
        header[..4].copy_from_slice(format!("{:04x}", data + 5).as_bytes());
        header
    }

    pub fn format(&self) -> ObjectFormat {
        self.format
    }

    fn hash_len(&self) -> u64 {
        hash_len(self.format) as u64
    }

    pub fn body(&self) -> u64 {
        self.pack - self.hash_len()
    }

    fn file_at(&self, at: u64) -> u64 {
        self.prefix + at / PACKET_DATA as u64 * PACKET + 5 + at % PACKET_DATA as u64
    }

    pub fn read(&self, file: &File, range: std::ops::Range<u64>) -> anyhow::Result<Vec<u8>> {
        let mut bytes = vec![0; (range.end - range.start) as usize];
        let mut at = range.start;
        let mut rest = &mut bytes[..];
        while !rest.is_empty() {
            let packet = at / PACKET_DATA as u64;
            self.check(file, packet)?;
            let in_packet = (PACKET_DATA as u64 - at % PACKET_DATA as u64) as usize;
            let (now, later) = rest.split_at_mut(in_packet.min(rest.len()));
            file.read_exact_at(now, self.file_at(at))?;
            at += now.len() as u64;
            rest = later;
        }
        Ok(bytes)
    }

    pub fn stream_ahead(&self, file: &File, each: &mut dyn FnMut(&[u8])) -> anyhow::Result<()> {
        std::thread::scope(|scope| {
            let (send, receive) = std::sync::mpsc::sync_channel::<Chunk>(AHEAD);
            let (back, spare) = std::sync::mpsc::channel();
            let reader = scope.spawn(move || {
                self.read_chunks(file, &mut || spare.try_recv().ok(), &mut |chunk| {
                    send.send(chunk).is_ok()
                })
            });
            let mut unframe = Unframe::new(self);
            for chunk in receive {
                unframe.feed(chunk.bytes(), each)?;
                let _ = back.send(chunk);
            }
            reader.join().expect("the reply reader panicked")?;
            unframe.finish()
        })
    }

    /// starts at the page the first packet is on, which keeps every read page aligned.
    fn read_chunks(
        &self,
        file: &File,
        spare: &mut dyn FnMut() -> Option<Chunk>,
        each: &mut dyn FnMut(Chunk) -> bool,
    ) -> anyhow::Result<()> {
        let end = file.metadata()?.len().saturating_sub(4);
        let mut at = self.prefix / PAGE * PAGE;
        while at < end {
            let mut chunk = match spare() {
                Some(chunk) => chunk,
                None => Chunk {
                    pages: memmap2::MmapMut::map_anon(CHUNK)?,
                    len: 0,
                },
            };
            chunk.len = CHUNK.min((end - at) as usize);
            file.read_exact_at(&mut chunk.pages[..chunk.len], at)
                .context("the reply ends before its pack")?;
            at += chunk.len as u64;
            if !each(chunk) {
                break;
            }
        }
        Ok(())
    }
}

/// smaller in tests, which lets a chunk end partway through a packet.
const CHUNK: usize = if cfg!(test) { 160 << 10 } else { 8 << 20 };
const AHEAD: usize = 3;
const PAGE: u64 = 4096;

/// mmapped, because reads into unaligned buffers cost the kernel extra cpu.
struct Chunk {
    pages: memmap2::MmapMut,
    len: usize,
}

impl Chunk {
    fn bytes(&self) -> &[u8] {
        &self.pages[..self.len]
    }
}

struct Unframe<'a> {
    framed: &'a Framed,
    skip: usize,
    at: u64,
    split: Vec<u8>,
}

impl<'a> Unframe<'a> {
    fn new(framed: &'a Framed) -> Self {
        Self {
            framed,
            skip: (framed.prefix % PAGE) as usize,
            at: 0,
            split: Vec::new(),
        }
    }

    fn feed(&mut self, chunk: &[u8], each: &mut dyn FnMut(&[u8])) -> anyhow::Result<()> {
        let skipped = self.skip.min(chunk.len());
        self.skip -= skipped;
        let mut rest = &chunk[skipped..];
        let body = self.framed.body();
        while self.at < body && !rest.is_empty() {
            let data = (self.framed.pack - self.at).min(PACKET_DATA as u64) as usize;
            let whole = data + 5;
            let packet = if self.split.is_empty() && rest.len() >= whole {
                let (packet, after) = rest.split_at(whole);
                rest = after;
                packet
            } else {
                let taken = (whole - self.split.len()).min(rest.len());
                self.split.extend_from_slice(&rest[..taken]);
                rest = &rest[taken..];
                if self.split.len() < whole {
                    break;
                }
                &self.split[..]
            };
            anyhow::ensure!(
                packet[..5] == self.framed.header(self.at / PACKET_DATA as u64),
                "the reply isn't in whole side-band packets"
            );
            let hashed = (body - self.at).min(data as u64) as usize;
            each(&packet[5..5 + hashed]);
            self.at += data as u64;
            self.split.clear();
        }
        Ok(())
    }

    fn finish(&self) -> anyhow::Result<()> {
        anyhow::ensure!(
            self.at >= self.framed.body(),
            "the reply ends before its pack"
        );
        Ok(())
    }
}

fn packet_len(hex: &[u8]) -> anyhow::Result<u64> {
    std::str::from_utf8(hex)
        .ok()
        .and_then(|hex| u64::from_str_radix(hex, 16).ok())
        .filter(|len| *len >= 4)
        .context("a reply with a bad packet length")
}

pub(crate) enum Checksum {
    Hash,
    Given(std::sync::mpsc::Receiver<Grown>),
    Spare(Midstate),
}

pub(crate) fn extend(
    prior: &File,
    reply: &mut File,
    segment: &[u8],
    format: ObjectFormat,
    checksum: Checksum,
) -> anyhow::Result<Vec<u8>> {
    let (added, entries) = pack_entries(segment, hash_len(format))?;
    let rewrite = Rewrite {
        kept: Framed::of(prior, format)?.body(),
        dropped: 0,
        added,
        entries,
        mark: None,
        tail_mark: None,
    };
    Ok(replace_tail(prior, reply, rewrite, format, checksum)?.checksum)
}

pub(crate) struct Rewrite<'a> {
    pub kept: u64,
    pub dropped: u32,
    pub added: u32,
    pub entries: &'a [u8],
    /// the grown pack offset to take the hash state at, which lands in the entries.
    pub mark: Option<u64>,
    /// an offset in the bytes kept, where the state is taken while hashing them.
    pub tail_mark: Option<u64>,
}

#[derive(Clone)]
pub(crate) struct Grown {
    pub checksum: Vec<u8>,
    pub at_mark: Option<Midstate>,
    pub at_tail: Option<Midstate>,
}

pub(crate) fn replace_tail(
    prior: &File,
    reply: &mut File,
    rewrite: Rewrite,
    format: ObjectFormat,
    checksum: Checksum,
) -> anyhow::Result<Grown> {
    let Rewrite {
        kept,
        dropped,
        added,
        entries,
        mark,
        tail_mark,
    } = rewrite;
    let framed = Framed::of(prior, format)?;
    let body = framed.body();
    anyhow::ensure!((12..=body).contains(&kept), "a tail outside the pack");
    let header = framed.read(prior, 0..12)?;
    anyhow::ensure!(
        header.starts_with(b"PACK\0\0\0\x02"),
        "prior reply isn't a version 2 pack"
    );
    let old = u32::from_be_bytes(header[8..12].try_into()?);
    let count = (old.checked_sub(dropped))
        .and_then(|count| count.checked_add(added))
        .context("an object count out of range")?
        .to_be_bytes();
    let marked = match mark {
        Some(mark) => Some(
            (mark.checked_sub(kept))
                .filter(|at| *at <= entries.len() as u64)
                .context("a mark outside the entries")? as usize,
        ),
        None => None,
    };
    anyhow::ensure!(
        tail_mark.is_none_or(|mark| (12..=kept).contains(&mark)),
        "a tail mark outside the bytes kept"
    );
    let mut at_tail = None;
    let cut_pack = kept / PACKET_DATA as u64 * PACKET_DATA as u64;
    let mut pending = framed.read(prior, cut_pack..kept)?;
    let hasher = match checksum {
        Checksum::Given(given) => Err(given),
        Checksum::Spare(state) => {
            anyhow::ensure!(count == old.to_be_bytes(), "resumed with the count changed");
            let blocks = kept / 64;
            let mut hasher = Hasher::resume(state, blocks);
            hasher.update(&framed.read(prior, blocks * 64..kept)?);
            Ok(hasher)
        }
        Checksum::Hash => {
            let mut hasher = Hasher::new(format);
            let mut at = 0u64;
            framed.stream_ahead(prior, &mut |data| {
                let data = &data[..data.len().min((kept - at) as usize)];
                let mut rest = data;
                if at == 0 {
                    hasher.update(&rest[..8]);
                    hasher.update(&count);
                    rest = &rest[12..];
                }
                if let Some(split) = (tail_mark.filter(|_| at_tail.is_none()))
                    .and_then(|mark| mark.checked_sub(at + (data.len() - rest.len()) as u64))
                    .filter(|split| *split <= rest.len() as u64)
                {
                    hasher.update(&rest[..split as usize]);
                    at_tail = Some(hasher.state().1);
                    rest = &rest[split as usize..];
                }
                hasher.update(rest);
                at += data.len() as u64;
            })?;
            Ok(hasher)
        }
    };
    let grown = match hasher {
        Ok(mut hasher) => {
            let (before, after) = entries.split_at(marked.unwrap_or(entries.len()));
            hasher.update(before);
            let at_mark = marked.map(|_| hasher.state().1);
            hasher.update(after);
            Grown {
                checksum: hasher.finish(),
                at_mark,
                at_tail,
            }
        }
        Err(given) => given
            .recv()
            .context("the reply hashing the same pack failed")?,
    };
    let checksum = &grown.checksum;
    for (at, byte) in (8u64..12).zip(count) {
        match at.checked_sub(cut_pack) {
            Some(past) => pending[past as usize] = byte,
            None => reply
                .write_all_at(&[byte], framed.file_at(at))
                .context("raise the object count")?,
        }
    }
    let cut = framed.file_at(cut_pack) - 5;
    reply.set_len(cut).context("cut the prior tail")?;
    reply.seek(SeekFrom::Start(cut))?;
    pending.extend_from_slice(entries);
    pending.extend_from_slice(checksum);
    let mut out = BufWriter::with_capacity(1 << 20, &mut *reply);
    for chunk in pending.chunks(PACKET_DATA) {
        packet(&mut out, chunk)?;
    }
    out.write_all(b"0000")?;
    out.into_inner().map_err(|error| error.into_error())?;
    Ok(grown)
}

/// the packs must not hold an object twice or a ref delta, which this doesn't check.
pub fn frame(
    format: ObjectFormat,
    public: &[std::path::PathBuf],
    cob: &[std::path::PathBuf],
    v0: Option<(&[u8], &mut File)>,
    v2: (&[u8], &mut File),
    like: Option<(&File, usize)>,
) -> anyhow::Result<()> {
    frame_spared(format, public, cob, v0, v2, like, None).map(drop)
}

/// `spare` fillers go at the end of each reply, and their start and hash state come back.
pub(crate) fn frame_spared(
    format: ObjectFormat,
    public: &[std::path::PathBuf],
    cob: &[std::path::PathBuf],
    v0: Option<(&[u8], &mut File)>,
    v2: (&[u8], &mut File),
    like: Option<(&File, usize)>,
    spare: Option<(u32, &[u8])>,
) -> anyhow::Result<[Option<(u64, Midstate)>; 2]> {
    let count = |packs: &[std::path::PathBuf]| -> anyhow::Result<u32> {
        packs.iter().try_fold(0u32, |sum, pack| {
            sum.checked_add(objects_in(pack)?)
                .context("too many objects for one pack")
        })
    };
    let (public_count, cob_count) = (count(public)?, count(cob)?);
    let (room, fillers) = spare.unwrap_or((0, &[]));
    let same = cob.is_empty() && v0.is_some();
    let source = match (&v0, like) {
        (Some((prefix, file)), _) if prefix.len() == v2.0.len() => Some((
            &**file,
            shared_below(format, public, same.then_some(fillers), prefix.len() as u64)?,
        )),
        (None, Some((like, n))) if starts_like(like, v2.0, format, &public[..n])? => Some((
            like,
            shared_below(format, &public[..n], None, v2.0.len() as u64)?,
        )),
        _ => None,
    };
    let shared = source.map_or(0, |(_, shared)| shared);
    let mut v0 = (v0.as_ref())
        .map(|(prefix, file)| Sink::new(prefix, file, Some(Hasher::new(format)), 0))
        .transpose()?;
    let v2_file = &*v2.1;
    let mut v2 = Sink::new(v2.0, v2_file, (!same).then(|| Hasher::new(format)), shared)?;
    let too_many = || anyhow::anyhow!("too many objects for one pack");
    if let Some(v0) = &mut v0 {
        v0.pack(&pack_header(
            public_count.checked_add(room).ok_or_else(too_many)?,
        ))?;
    }
    let header = pack_header(
        (public_count.checked_add(cob_count))
            .and_then(|count| count.checked_add(room))
            .ok_or_else(too_many)?,
    );
    std::thread::scope(|scope| -> anyhow::Result<_> {
        let (send, receive) = std::sync::mpsc::sync_channel::<std::sync::Arc<Vec<u8>>>(4);
        let v2 = scope.spawn(move || -> anyhow::Result<(Sink, Option<(u64, Midstate)>)> {
            v2.pack(&header)?;
            for chunk in receive {
                v2.pack(&chunk)?;
            }
            let mark = v2.mark();
            if spare.is_some() {
                v2.pack(fillers)?;
            }
            Ok((v2, mark))
        });
        let sent = (|| -> anyhow::Result<Option<(u64, Midstate)>> {
            for pack in public {
                pack_entries_of(pack, format, &mut |chunk| {
                    if let Some(v0) = &mut v0 {
                        v0.pack(&chunk)?;
                    }
                    send.send(chunk).ok().context("the v2 reply stopped")
                })?;
            }
            let mark = v0.as_ref().and_then(Sink::mark);
            if let (Some(v0), Some(_)) = (&mut v0, spare) {
                v0.pack(fillers)?;
            }
            for pack in cob {
                pack_entries_of(pack, format, &mut |chunk| {
                    send.send(chunk).ok().context("the v2 reply stopped")
                })?;
            }
            Ok(mark)
        })();
        drop(send);
        let (mut v2, v2_mark) = v2.join().expect("the v2 reply thread panicked")?;
        let v0_mark = sent?;
        let v0_checksum = v0
            .as_mut()
            .map(|v0| v0.hasher.take().expect("v0 hashes").finish());
        let v2_checksum = match v2.hasher.take() {
            Some(hasher) => hasher.finish(),
            None => v0_checksum.clone().expect("v2 hashes without v0"),
        };
        if let (Some(v0), Some(checksum)) = (v0, v0_checksum) {
            v0.finish(&checksum)?;
        }
        let head = v2.finish(&v2_checksum)?;
        if let Some((source, shared)) = source.filter(|(_, shared)| *shared > 0) {
            clone_range(source, 0, v2_file, 0, shared).context("share a reply's bytes")?;
            v2_file.write_all_at(&head, 0).context("write a reply")?;
        }
        let marks = [v0_mark, v2_mark.or(v0_mark)];
        Ok(match spare {
            Some(_) => marks,
            None => [None, None],
        })
    })
}

/// the bytes a v2 reply shares with the v0 of `public`, rounded down to a block.
fn shared_below(
    format: ObjectFormat,
    public: &[std::path::PathBuf],
    same: Option<&[u8]>,
    prefix: u64,
) -> anyhow::Result<u64> {
    let mut common = 12u64;
    for pack in public {
        let len = std::fs::metadata(pack).with_context(|| format!("stat {}", pack.display()))?;
        common += (len.len())
            .checked_sub(12 + hash_len(format) as u64)
            .with_context(|| format!("{} is too short for a pack", pack.display()))?;
    }
    if let Some(fillers) = same {
        common += (fillers.len() + hash_len(format)) as u64;
    }
    let at = prefix + common / PACKET_DATA as u64 * PACKET;
    Ok(at / BLOCK * BLOCK)
}

fn starts_like(
    like: &File,
    prefix: &[u8],
    format: ObjectFormat,
    packs: &[std::path::PathBuf],
) -> anyhow::Result<bool> {
    let Ok(framed) = Framed::of(like, format) else {
        return Ok(false);
    };
    if framed.prefix != prefix.len() as u64 {
        return Ok(false);
    }
    let mut at = 12;
    let mut alike = true;
    for pack in packs {
        pack_entries_of(pack, format, &mut |chunk| {
            let end = at + chunk.len() as u64;
            alike = alike && end <= framed.body() && framed.read(like, at..end)? == *chunk;
            at = end;
            Ok(())
        })?;
    }
    Ok(alike)
}

/// shares extents where the filesystem can (btrfs, xfs), else copies.
/// `from` and `to` should be `BLOCK` multiples to share.
pub(crate) fn clone_range(
    source: &File,
    from: u64,
    target: &File,
    to: u64,
    len: u64,
) -> std::io::Result<()> {
    #[cfg(target_os = "linux")]
    {
        use std::os::fd::AsRawFd as _;
        let range = libc::file_clone_range {
            src_fd: source.as_raw_fd() as i64,
            src_offset: from,
            src_length: len,
            dest_offset: to,
        };
        // SAFETY: both descriptors are open for the duration of the call, and
        // `range` outlives it.
        if unsafe { libc::ioctl(target.as_raw_fd(), libc::FICLONERANGE, &range) } == 0 {
            return Ok(());
        }
    }
    let mut chunk = vec![0; 1 << 20];
    let mut at = 0;
    while at < len {
        let now = &mut chunk[..(len - at).min(1 << 20) as usize];
        source.read_exact_at(now, from + at)?;
        target.write_all_at(now, to + at)?;
        at += now.len() as u64;
    }
    Ok(())
}

/// an empty band-2 (progress) packet, which clients show nothing for.
pub(crate) const PAD: &[u8] = b"0005\x02";

/// 3277 is the inverse of 5 mod `BLOCK`, which puts the pads on a block boundary.
pub(crate) fn padding(at: u64) -> Vec<u8> {
    let short = (BLOCK - at % BLOCK) % BLOCK;
    PAD.repeat((short * 3277 % BLOCK) as usize)
}

pub use super::hasher::*;

#[cfg(any(test, feature = "bench"))]
mod benches {
    use super::*;
    use crate::bench_fixture::repo;

    const PREFIX: &[u8] = b"000dpackfile\n";

    fn packs(dir: &std::path::Path, mib: usize) -> (std::path::PathBuf, Vec<u8>) {
        let (base, main) = repo(dir, mib * 256, 4096, 4);
        let pack = std::fs::read_dir(dir.join("objects/pack"))
            .unwrap()
            .map(|entry| entry.unwrap().path())
            .find(|path| path.extension().is_some_and(|ext| ext == "pack"))
            .unwrap();
        let mut segment = crate::git_export::git_command(dir);
        segment
            .args(["pack-objects", "-q", "--stdout", "--revs"])
            .stdin(
                crate::git_export::stdin_bytes(format!("{main}\n^{base}\n").as_bytes()).unwrap(),
            );
        let segment = crate::git_export::run_cmd(segment, "git pack-objects")
            .unwrap()
            .stdout;
        (pack, segment)
    }

    #[divan::bench(args = [1, 32])]
    fn frame(bencher: divan::Bencher, mib: usize) {
        let dir = tempfile::tempdir().unwrap();
        let (pack, _) = packs(&dir.path().join("repo"), mib);
        bencher
            .with_inputs(|| tempfile::tempfile_in(dir.path()).unwrap())
            .bench_local_values(|mut file| {
                super::frame(
                    ObjectFormat::SHA1,
                    std::slice::from_ref(&pack),
                    &[],
                    None,
                    (PREFIX, &mut file),
                    None,
                )
                .unwrap();
                file
            });
    }

    #[divan::bench(args = [1, 32])]
    fn extend(bencher: divan::Bencher, mib: usize) {
        let dir = tempfile::tempdir().unwrap();
        let (pack, segment) = packs(&dir.path().join("repo"), mib);
        let mut prior = tempfile::tempfile_in(dir.path()).unwrap();
        super::frame(
            ObjectFormat::SHA1,
            &[pack],
            &[],
            None,
            (PREFIX, &mut prior),
            None,
        )
        .unwrap();
        bencher
            .with_inputs(|| {
                let reply = tempfile::tempfile_in(dir.path()).unwrap();
                reflink_or_copy(&prior, &reply).unwrap();
                reply
            })
            .bench_local_values(|mut reply| {
                let checksum = super::extend(
                    &prior,
                    &mut reply,
                    &segment,
                    ObjectFormat::SHA1,
                    Checksum::Hash,
                )
                .unwrap();
                (divan::black_box(checksum), reply)
            });
    }
}

#[cfg(test)]
#[path = "tests/reply_pack.rs"]
pub(crate) mod tests;