use crate::core::cas;
use crate::core::compress;
use crate::core::framing::Framing;
use crate::core::protocol::{
    self, CompressionAlgo, ListEntry, Message, CAP_AEAD, CAP_DEDUP, CAP_DIRECT_TCP, CAP_FAST,
    CAP_LZ4, CAP_SYNC, CAP_ZSTD, PROTOCOL_VERSION,
};
use crate::core::protocol_aead::Direction;
use anyhow::{bail, Result};
use std::path::{Path, PathBuf};
use tokio::fs;
use tokio::io::{self, AsyncWriteExt};

const CHUNK_SIZE: usize = 4 * 1024 * 1024;

/// Advertised capability mask. The client picks by intersecting with
/// its own; CAP_FAST is always offered since --fast is client-side opt-in.
const SERVER_CAPS: u8 =
    CAP_LZ4 | CAP_ZSTD | CAP_DEDUP | CAP_FAST | CAP_SYNC | CAP_DIRECT_TCP | CAP_AEAD;

/// Default root is `$HOME`; pass `--root /` to opt out of the sandbox.
fn resolve_root(arg: Option<PathBuf>) -> Result<PathBuf> {
    let raw = match arg {
        Some(p) => p,
        None => directories::UserDirs::new()
            .map(|u| u.home_dir().to_path_buf())
            .ok_or_else(|| anyhow::anyhow!("no $HOME to use as default --root"))?,
    };
    std::fs::create_dir_all(&raw)?;
    Ok(std::fs::canonicalize(&raw)?)
}

/// Reject null bytes and literal `..` components, then canonicalize
/// (splicing the tail if the path doesn't yet exist) and require a
/// lexical prefix match against `root`. The `..` check is belt-and-
/// suspenders; canonicalize catches symlink tricks too.
fn validate_path(raw: &str, root: &Path) -> Result<PathBuf> {
    if raw.contains('\0') {
        bail!("path contains null byte");
    }
    let path = Path::new(raw);

    for component in path.components() {
        if matches!(component, std::path::Component::ParentDir) {
            bail!("path contains '..'");
        }
    }

    let canonical = canonicalize_with_ancestor(path)?;
    if !canonical.starts_with(root) {
        bail!(
            "path {} escapes server root {}",
            canonical.display(),
            root.display()
        );
    }
    Ok(canonical)
}

/// Only follows symlinks on the existing prefix, so a symlink escaping
/// `root` is still caught by the caller's `starts_with` check.
fn canonicalize_with_ancestor(path: &Path) -> Result<PathBuf> {
    if path.exists() {
        return Ok(std::fs::canonicalize(path)?);
    }
    let mut ancestor = path.to_path_buf();
    let mut tail: Vec<std::ffi::OsString> = Vec::new();
    while !ancestor.exists() {
        match ancestor.file_name() {
            Some(n) => tail.push(n.to_os_string()),
            None => break,
        }
        if !ancestor.pop() {
            break;
        }
    }
    let mut out = if ancestor.as_os_str().is_empty() {
        std::env::current_dir()?
    } else if ancestor.exists() {
        std::fs::canonicalize(&ancestor)?
    } else {
        return Err(anyhow::anyhow!(
            "no existing ancestor for {}",
            path.display()
        ));
    };
    for seg in tail.iter().rev() {
        out.push(seg);
    }
    Ok(out)
}

pub async fn run(root: Option<PathBuf>) -> Result<()> {
    let root = resolve_root(root)?;
    let mut stdin = io::stdin();
    let mut stdout = io::stdout();
    // stdin/stdout is the sshd-invoked transport:
    // - allow_splice: writer is real STDOUT_FILENO
    // - allow_direct_tcp: peer is already SSH-authenticated
    // - direct_tcp_key=None: no session key, CAP_AEAD masked off
    run_session(&mut stdin, &mut stdout, &root, true, true, None).await
}

pub async fn run_listen(root: Option<PathBuf>, addr: std::net::SocketAddr) -> Result<()> {
    // Non-loopback bind has no peer auth; gate behind an explicit env
    // var so a copy-pasted dev command can't open the filesystem to
    // the LAN by accident.
    if !addr.ip().is_loopback() {
        if std::env::var("BCMR_UNSAFE_LAN_LISTEN").is_ok_and(|v| v == "1") {
            eprintln!(
                "bcmr serve: WARNING — binding {addr} on a non-loopback address with no peer \
                 authentication. Anyone who can reach this port can read and write files under \
                 the --root jail."
            );
        } else {
            bail!(
                "bcmr serve --listen refuses non-loopback address {addr} without \
                 BCMR_UNSAFE_LAN_LISTEN=1 (no peer auth on this transport yet)"
            );
        }
    }
    let root = resolve_root(root)?;
    let listener = tokio::net::TcpListener::bind(addr).await?;
    let bound = listener.local_addr()?;
    // INVARIANT: test harnesses parse this line to pick up the
    // kernel-assigned port — don't change the format.
    use tokio::io::AsyncWriteExt as _;
    let mut stdout = io::stdout();
    stdout
        .write_all(format!("LISTENING {bound}\n").as_bytes())
        .await?;
    stdout.flush().await?;

    loop {
        let (stream, _peer) = listener.accept().await?;
        let root = root.clone();
        tokio::spawn(async move {
            let (mut reader, mut writer) = stream.into_split();
            // TCP peer can't splice (fd != STDOUT_FILENO) and must not
            // escalate to another direct-TCP channel (recursive
            // rendezvous = amplification / port exhaustion).
            let _ = run_session(&mut reader, &mut writer, &root, false, false, None).await;
        });
    }
}

async fn run_session<R, W>(
    reader: &mut R,
    writer: &mut W,
    root: &Path,
    allow_splice: bool,
    allow_direct_tcp: bool,
    direct_tcp_key: Option<&[u8; 32]>,
) -> Result<()>
where
    R: tokio::io::AsyncRead + Unpin,
    W: tokio::io::AsyncWrite + Unpin,
{
    // Closing CAP_DIRECT_TCP on non-TTY transports closes the recursive-
    // rendezvous amplification vector. CAP_AEAD needs a session key.
    let mut offered_caps = SERVER_CAPS;
    if !allow_direct_tcp {
        offered_caps &= !CAP_DIRECT_TCP;
    }
    if direct_tcp_key.is_none() {
        offered_caps &= !CAP_AEAD;
    }
    let mut rendezvous_tasks = RendezvousTasks::new();

    let mut framing = Framing::plain();

    let effective_caps = match framing.read_message(reader).await? {
        Some(Message::Hello { version, caps }) => {
            if version != PROTOCOL_VERSION {
                framing
                    .write_message(
                        writer,
                        &Message::Error {
                            message: format!(
                                "protocol version mismatch: client={version} server={PROTOCOL_VERSION}"
                            ),
                        },
                    )
                    .await?;
                writer.flush().await?;
                return Ok(());
            }
            offered_caps & caps
        }
        Some(other) => {
            framing
                .write_message(
                    writer,
                    &Message::Error {
                        message: format!("expected Hello, got {other:?}"),
                    },
                )
                .await?;
            writer.flush().await?;
            return Ok(());
        }
        None => return Ok(()),
    };

    // Passing the same mask on both sides: client + server have already
    // intersected caps, so a single-sided call agrees on algo here.
    let algo = CompressionAlgo::negotiate(effective_caps, effective_caps);
    let fast = (effective_caps & CAP_FAST) != 0;
    let sync = (effective_caps & CAP_SYNC) != 0;
    let direct_tcp = (effective_caps & CAP_DIRECT_TCP) != 0;
    let aead = (effective_caps & CAP_AEAD) != 0;

    // Downgrade guard: a MITM stripping CAP_AEAD from Hello/Welcome
    // would otherwise run the rendezvous session in plain framing.
    if direct_tcp_key.is_some() && !aead {
        framing
            .write_message(
                writer,
                &Message::Error {
                    message: "direct-TCP transport requires CAP_AEAD; refusing plain session"
                        .to_string(),
                },
            )
            .await?;
        writer.flush().await?;
        return Ok(());
    }

    framing
        .write_message(
            writer,
            &Message::Welcome {
                version: PROTOCOL_VERSION,
                caps: effective_caps,
            },
        )
        .await?;
    writer.flush().await?;

    // INVARIANT: Welcome is the plain→AEAD switchover boundary.
    // Both sides MUST flip together here.
    if aead {
        let key = direct_tcp_key
            .expect("cap mask guarantees direct_tcp_key is Some when CAP_AEAD negotiated");
        framing =
            Framing::aead_from_key(key, Direction::ServerToClient, Direction::ClientToServer)?;
    }

    loop {
        let msg = match framing.read_message(reader).await? {
            Some(m) => m,
            None => break,
        };

        // Get/GetChunked write Data+Ok to `writer` themselves; the
        // dispatch loop must NOT send a response for those arms.
        let response = match msg {
            Message::Get { path, offset } => {
                match validate_path(&path, root) {
                    Ok(p) => {
                        if let Err(e) = handle_get(
                            p.to_str().unwrap_or(&path),
                            offset,
                            algo,
                            fast,
                            allow_splice,
                            writer,
                            &mut framing,
                        )
                        .await
                        {
                            eprintln!("serve: handler error: {e}");
                            framing
                                .write_message(
                                    writer,
                                    &Message::Error {
                                        message: e.to_string(),
                                    },
                                )
                                .await?;
                        }
                    }
                    Err(e) => {
                        framing
                            .write_message(
                                writer,
                                &Message::Error {
                                    message: e.to_string(),
                                },
                            )
                            .await?;
                    }
                }
                writer.flush().await?;
                continue;
            }
            Message::Stat { path } => match validate_path(&path, root) {
                Ok(p) => handle_stat(p.to_str().unwrap_or(&path)).await,
                Err(e) => Err(e),
            },
            Message::List { path } => match validate_path(&path, root) {
                Ok(p) => handle_list(p.to_str().unwrap_or(&path)).await,
                Err(e) => Err(e),
            },
            Message::Hash {
                path,
                offset,
                limit,
            } => match validate_path(&path, root) {
                Ok(p) => handle_hash(p.to_str().unwrap_or(&path), offset, limit).await,
                Err(e) => Err(e),
            },
            Message::Put { path, size } => match validate_path(&path, root) {
                Ok(p) => {
                    handle_put(
                        p.to_str().unwrap_or(&path),
                        size,
                        sync,
                        writer,
                        reader,
                        &mut framing,
                    )
                    .await
                }
                Err(e) => Err(e),
            },
            Message::PutChunked {
                path,
                offset,
                length,
            } => match validate_path(&path, root) {
                Ok(p) => {
                    handle_put_chunked(
                        p.to_str().unwrap_or(&path),
                        offset,
                        length,
                        sync,
                        reader,
                        &mut framing,
                    )
                    .await
                }
                Err(e) => Err(e),
            },
            Message::GetChunked {
                path,
                offset,
                length,
            } => {
                match validate_path(&path, root) {
                    Ok(p) => {
                        if let Err(e) = handle_get_chunked(
                            p.to_str().unwrap_or(&path),
                            offset,
                            length,
                            algo,
                            writer,
                            &mut framing,
                        )
                        .await
                        {
                            eprintln!("serve: handler error: {e}");
                            framing
                                .write_message(
                                    writer,
                                    &Message::Error {
                                        message: e.to_string(),
                                    },
                                )
                                .await?;
                        }
                    }
                    Err(e) => {
                        framing
                            .write_message(
                                writer,
                                &Message::Error {
                                    message: e.to_string(),
                                },
                            )
                            .await?;
                    }
                }
                writer.flush().await?;
                continue;
            }
            Message::Mkdir { path } => match validate_path(&path, root) {
                Ok(p) => handle_mkdir(p.to_str().unwrap_or(&path)).await,
                Err(e) => Err(e),
            },
            Message::Truncate { path, size } => match validate_path(&path, root) {
                Ok(p) => handle_truncate(p.to_str().unwrap_or(&path), size).await,
                Err(e) => Err(e),
            },
            Message::Resume { path } => match validate_path(&path, root) {
                Ok(p) => handle_resume(p.to_str().unwrap_or(&path)).await,
                Err(e) => Err(e),
            },
            Message::OpenDirectChannel => {
                if !direct_tcp {
                    Err(anyhow::anyhow!(
                        "CAP_DIRECT_TCP not negotiated on this session"
                    ))
                } else if rendezvous_tasks.len() >= MAX_RENDEZVOUS_PER_SESSION {
                    Err(anyhow::anyhow!(
                        "too many concurrent direct-TCP rendezvous requests \
                         on this session (limit {MAX_RENDEZVOUS_PER_SESSION})"
                    ))
                } else {
                    match handle_open_direct_channel(root.to_path_buf()) {
                        Ok((msg, handle)) => {
                            rendezvous_tasks.push(handle);
                            Ok(msg)
                        }
                        Err(e) => Err(e),
                    }
                }
            }
            other => Err(anyhow::anyhow!("unexpected message: {other:?}")),
        };

        let reply = match response {
            Ok(msg) => msg,
            Err(e) => {
                eprintln!("serve: handler error: {e}");
                Message::Error {
                    message: e.to_string(),
                }
            }
        };

        framing.write_message(writer, &reply).await?;
        writer.flush().await?;
    }

    rendezvous_tasks.drain_gracefully().await;
    Ok(())
}

async fn handle_stat(path: &str) -> Result<Message> {
    let meta = fs::metadata(path).await?;
    let mtime = meta
        .modified()
        .ok()
        .and_then(|t| {
            t.duration_since(std::time::UNIX_EPOCH)
                .ok()
                .map(|d| d.as_secs() as i64)
        })
        .unwrap_or(0);
    Ok(Message::StatResponse {
        size: meta.len(),
        mtime,
        is_dir: meta.is_dir(),
    })
}

async fn handle_list(path: &str) -> Result<Message> {
    let base = Path::new(path).to_path_buf();
    let entries = tokio::task::spawn_blocking(move || -> Result<Vec<ListEntry>> {
        let mut out = Vec::new();
        for entry in walkdir::WalkDir::new(&base).min_depth(1) {
            let entry = entry?;
            let rel = entry
                .path()
                .strip_prefix(&base)
                .unwrap_or(entry.path())
                .to_string_lossy()
                .into_owned();
            let meta = entry.metadata()?;
            out.push(ListEntry {
                path: rel,
                size: meta.len(),
                is_dir: meta.is_dir(),
            });
        }
        Ok(out)
    })
    .await??;

    Ok(Message::ListResponse { entries })
}

async fn handle_hash(path: &str, offset: u64, limit: Option<u64>) -> Result<Message> {
    let path = path.to_owned();
    let hash = tokio::task::spawn_blocking(move || -> Result<String> {
        use std::io::Read;
        let mut file = std::fs::File::open(&path)?;
        if offset > 0 {
            use std::io::Seek;
            file.seek(std::io::SeekFrom::Start(offset))?;
        }
        let mut hasher = blake3::Hasher::new();
        let mut remaining = limit;
        let mut buf = vec![0u8; CHUNK_SIZE];
        loop {
            let to_read = match remaining {
                Some(0) => break,
                Some(r) => r.min(buf.len() as u64) as usize,
                None => buf.len(),
            };
            let n = file.read(&mut buf[..to_read])?;
            if n == 0 {
                break;
            }
            hasher.update(&buf[..n]);
            if let Some(r) = remaining.as_mut() {
                *r -= n as u64;
            }
        }
        Ok(hasher.finalize().to_hex().to_string())
    })
    .await??;

    Ok(Message::HashResponse { hash })
}

/// Writes Data messages + terminal Ok directly to `out`; returns
/// `Result<()>` so the dispatch loop doesn't double-send.
async fn handle_get<W>(
    path: &str,
    offset: u64,
    algo: CompressionAlgo,
    fast: bool,
    allow_splice: bool,
    out: &mut W,
    framing: &mut Framing,
) -> Result<()>
where
    W: tokio::io::AsyncWrite + Unpin,
{
    // splice(2) bypasses userspace, so it can't run with compression
    // (no bytes to feed the encoder) or AEAD (no bytes to encrypt).
    // `allow_splice` guards against non-STDOUT_FILENO writers (TCP
    // sockets) — the splice path hardcodes that fd. CAP_FAST without
    // splice still wins from skipping the inline hash.
    #[cfg(target_os = "linux")]
    {
        if fast && algo == CompressionAlgo::None && allow_splice && !framing.is_aead() {
            return handle_get_splice_linux(path, offset, out).await;
        }
    }
    let _ = allow_splice;

    use tokio::io::AsyncReadExt;

    let mut file = fs::File::open(path).await?;
    if offset > 0 {
        use tokio::io::AsyncSeekExt;
        file.seek(std::io::SeekFrom::Start(offset)).await?;
    }

    let mut hasher = (!fast).then(blake3::Hasher::new);
    let mut buf = vec![0u8; CHUNK_SIZE];
    loop {
        let n = file.read(&mut buf).await?;
        if n == 0 {
            break;
        }
        if let Some(h) = hasher.as_mut() {
            h.update(&buf[..n]);
        }
        let frame = compress::encode_block(algo, buf[..n].to_vec());
        framing.write_message(out, &frame).await?;
    }

    let hash = hasher.map(|h| h.finalize().to_hex().to_string());
    framing.write_message(out, &Message::Ok { hash }).await?;
    Ok(())
}

/// Linux zero-copy GET: file → pipe → stdout via splice(2).
///
/// The entire loop runs inside ONE spawn_blocking and writes frame
/// headers directly with write(2). An earlier implementation that used
/// one spawn_blocking per 4 MiB chunk + async stdout for headers was
/// *slower* than the buffered path (~256 thread bounces per GiB).
///
/// F_SETPIPE_SZ is advisory; on distros where `/proc/sys/fs/pipe-max-
/// size` is 1 MiB each 4 MiB chunk takes 4 splice rounds, which is
/// fine — the syscall cost is negligible.
#[cfg(target_os = "linux")]
async fn handle_get_splice_linux<W>(path: &str, offset: u64, out: &mut W) -> Result<()>
where
    W: tokio::io::AsyncWrite + Unpin,
{
    use std::io::Seek;
    use std::os::unix::io::{AsRawFd, FromRawFd, OwnedFd};
    use tokio::io::AsyncWriteExt;

    // Flush the async stdout buffer so raw write(2) in the blocking
    // loop doesn't race with it.
    out.flush().await?;

    let path = path.to_owned();
    let result = tokio::task::spawn_blocking(move || -> std::io::Result<()> {
        let mut std_file = std::fs::File::open(&path)?;
        if offset > 0 {
            std_file.seek(std::io::SeekFrom::Start(offset))?;
        }
        let total_size = std_file.metadata()?.len() - offset;

        let mut fds = [0i32; 2];
        if unsafe { libc::pipe2(fds.as_mut_ptr(), libc::O_CLOEXEC) } < 0 {
            return Err(std::io::Error::last_os_error());
        }
        let pipe_r = unsafe { OwnedFd::from_raw_fd(fds[0]) };
        let pipe_w = unsafe { OwnedFd::from_raw_fd(fds[1]) };

        let _requested_sz = unsafe {
            libc::fcntl(
                pipe_w.as_raw_fd(),
                libc::F_SETPIPE_SZ,
                CHUNK_SIZE as libc::c_int,
            )
        };

        let file_fd = std_file.as_raw_fd();
        let pipe_w_fd = pipe_w.as_raw_fd();
        let pipe_r_fd = pipe_r.as_raw_fd();
        let stdout_fd = libc::STDOUT_FILENO;

        let mut remaining = total_size;
        while remaining > 0 {
            let chunk = remaining.min(CHUNK_SIZE as u64) as usize;

            // Wire frame: [u32 payload_len][u8 TYPE_DATA=0x84][u32 chunk_len].
            // payload_len = 1 + 4 + chunk.
            let mut header = [0u8; 9];
            header[0..4].copy_from_slice(&((5 + chunk) as u32).to_le_bytes());
            header[4] = 0x84;
            header[5..9].copy_from_slice(&(chunk as u32).to_le_bytes());
            write_all_fd(stdout_fd, &header)?;

            splice_n(file_fd, pipe_w_fd, pipe_r_fd, stdout_fd, chunk)?;
            remaining -= chunk as u64;
        }

        // Ok { hash: None } = [u32 payload_len=2][u8 TYPE_OK=0x82][u8 present=0].
        let mut ok_frame = [0u8; 6];
        ok_frame[0..4].copy_from_slice(&2u32.to_le_bytes());
        ok_frame[4] = 0x82;
        ok_frame[5] = 0;
        write_all_fd(stdout_fd, &ok_frame)?;

        drop(pipe_r);
        drop(pipe_w);
        drop(std_file);
        Ok(())
    })
    .await?;
    result.map_err(Into::into)
}

/// Loops because splice may move fewer bytes than requested (pipe
/// smaller than `n`). Standard file→pipe→sock pairing.
#[cfg(target_os = "linux")]
fn splice_n(
    file_fd: i32,
    pipe_w_fd: i32,
    pipe_r_fd: i32,
    stdout_fd: i32,
    n: usize,
) -> std::io::Result<()> {
    let mut moved = 0usize;
    while moved < n {
        let want = n - moved;
        let got = unsafe {
            libc::splice(
                file_fd,
                std::ptr::null_mut(),
                pipe_w_fd,
                std::ptr::null_mut(),
                want,
                libc::SPLICE_F_MOVE | libc::SPLICE_F_MORE,
            )
        };
        if got < 0 {
            return Err(std::io::Error::last_os_error());
        }
        if got == 0 {
            break;
        }
        let mut drained = 0usize;
        while drained < got as usize {
            let n2 = unsafe {
                libc::splice(
                    pipe_r_fd,
                    std::ptr::null_mut(),
                    stdout_fd,
                    std::ptr::null_mut(),
                    got as usize - drained,
                    libc::SPLICE_F_MOVE | libc::SPLICE_F_MORE,
                )
            };
            if n2 < 0 {
                return Err(std::io::Error::last_os_error());
            }
            if n2 == 0 {
                break;
            }
            drained += n2 as usize;
        }
        moved += drained;
    }
    Ok(())
}

/// Write::write_all on a raw fd. Loops on short writes / EINTR.
#[cfg(target_os = "linux")]
fn write_all_fd(fd: i32, mut buf: &[u8]) -> std::io::Result<()> {
    while !buf.is_empty() {
        let n = unsafe { libc::write(fd, buf.as_ptr() as *const _, buf.len()) };
        if n < 0 {
            let e = std::io::Error::last_os_error();
            if e.kind() == std::io::ErrorKind::Interrupted {
                continue;
            }
            return Err(e);
        }
        if n == 0 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::WriteZero,
                "write returned 0",
            ));
        }
        buf = &buf[n as usize..];
    }
    Ok(())
}

/// When the client opens with HaveBlocks (CAP_DEDUP negotiated), blocks
/// already in the local CAS are served from cache without the wire
/// transfer; new blocks are written to the file AND deposited in the
/// CAS. The returned composite hash covers the whole file regardless.
async fn handle_put<W, R>(
    path: &str,
    declared_size: u64,
    sync: bool,
    out: &mut W,
    reader: &mut R,
    framing: &mut Framing,
) -> Result<Message>
where
    W: tokio::io::AsyncWrite + Unpin,
    R: tokio::io::AsyncRead + Unpin,
{
    ensure_parent_dir(path).await?;

    let mut file = fs::File::create(path).await?;
    let mut hasher = blake3::Hasher::new();
    let mut written: u64 = 0;

    // HaveBlocks = dedup open; anything else = straight-to-Data.
    let first = framing.read_message(reader).await?;
    let mut dedup_state: Option<DedupState> = None;
    let mut next: Option<Message> = first;

    if let Some(Message::HaveBlocks { hashes, .. }) = next {
        // Evict before writing so this PUT doesn't push further past
        // the cap. Non-fatal; worst case the disk fills up later.
        if let Some(cap) = cas::cap_bytes() {
            let _ = tokio::task::spawn_blocking(move || cas::evict_to_cap(cap)).await;
        }

        let mut bits = vec![0u8; hashes.len().div_ceil(8)];
        for (i, h) in hashes.iter().enumerate() {
            if !cas::has(h) {
                bits[i / 8] |= 1 << (i % 8);
            }
        }
        framing
            .write_message(out, &Message::MissingBlocks { bits: bits.clone() })
            .await?;
        out.flush().await?;
        dedup_state = Some(DedupState {
            hashes,
            bits,
            cursor: 0,
        });
        next = None;
    }

    let mut msg = next;
    loop {
        let m = match msg.take() {
            Some(m) => m,
            None => match framing.read_message(reader).await? {
                Some(m) => m,
                None => break,
            },
        };
        match m {
            Message::Data { payload } => {
                enforce_write_bound(written, payload.len(), declared_size)?;
                consume_block(&payload, &mut file, &mut hasher, dedup_state.as_mut()).await?;
                written += payload.len() as u64;
            }
            Message::DataCompressed {
                algo,
                original_size,
                payload,
            } => {
                let decoded = compress::decode_block(algo, original_size, &payload)?;
                enforce_write_bound(written, decoded.len(), declared_size)?;
                consume_block(&decoded, &mut file, &mut hasher, dedup_state.as_mut()).await?;
                written += decoded.len() as u64;
            }
            Message::Done => break,
            other => return Err(anyhow::anyhow!("put: unexpected message {other:?}")),
        }
    }

    // Serve trailing cache-hit blocks *after* draining the wire so
    // writes preserve source order.
    if let Some(state) = dedup_state.as_mut() {
        flush_remaining_cas_blocks(state, &mut file, &mut hasher, &mut written, declared_size)
            .await?;
    }

    file.flush().await?;
    if sync {
        file.sync_all().await?;
    }
    drop(file);

    let hash = hasher.finalize().to_hex().to_string();
    Ok(Message::Ok { hash: Some(hash) })
}

struct DedupState {
    hashes: Vec<[u8; 32]>,
    bits: Vec<u8>,
    cursor: usize,
}

impl DedupState {
    fn is_missing(&self, idx: usize) -> bool {
        (self.bits.get(idx / 8).copied().unwrap_or(0) >> (idx % 8)) & 1 == 1
    }
}

async fn consume_block(
    block: &[u8],
    file: &mut tokio::fs::File,
    hasher: &mut blake3::Hasher,
    dedup: Option<&mut DedupState>,
) -> Result<()> {
    if let Some(state) = dedup {
        while state.cursor < state.hashes.len() && !state.is_missing(state.cursor) {
            let cached = cas::read(&state.hashes[state.cursor])?;
            hasher.update(&cached);
            file.write_all(&cached).await?;
            state.cursor += 1;
        }
        if state.cursor < state.hashes.len() && state.is_missing(state.cursor) {
            let mut h = [0u8; 32];
            h.copy_from_slice(blake3::hash(block).as_bytes());
            // Best-effort CAS write; serving the file matters more.
            let _ = cas::write(&h, block);
            state.cursor += 1;
        }
    }
    hasher.update(block);
    file.write_all(block).await?;
    Ok(())
}

async fn flush_remaining_cas_blocks(
    state: &mut DedupState,
    file: &mut tokio::fs::File,
    hasher: &mut blake3::Hasher,
    written: &mut u64,
    declared_size: u64,
) -> Result<()> {
    while state.cursor < state.hashes.len() {
        if state.is_missing(state.cursor) {
            return Err(anyhow::anyhow!(
                "client said block {} was missing but never sent it",
                state.cursor
            ));
        }
        let cached = cas::read(&state.hashes[state.cursor])?;
        enforce_write_bound(*written, cached.len(), declared_size)?;
        hasher.update(&cached);
        file.write_all(&cached).await?;
        *written += cached.len() as u64;
        state.cursor += 1;
    }
    Ok(())
}

/// Without this a malicious client could declare `size: 100` then
/// stream TB of Data frames.
fn enforce_write_bound(written: u64, incoming: usize, declared: u64) -> Result<()> {
    if written + incoming as u64 > declared {
        bail!(
            "put: client would write {} bytes past the declared size of {} \
             (already wrote {})",
            written + incoming as u64 - declared,
            declared,
            written
        );
    }
    Ok(())
}

/// No truncate: parallel serve processes from the same pool write
/// disjoint ranges without clobbering. Integrity comes from the AEAD
/// per-frame MAC (or client-side `--verify` rehash), not per-chunk
/// server-side hashes.
async fn handle_put_chunked<R>(
    path: &str,
    offset: u64,
    length: u64,
    sync: bool,
    reader: &mut R,
    framing: &mut Framing,
) -> Result<Message>
where
    R: tokio::io::AsyncRead + Unpin,
{
    use tokio::io::{AsyncSeekExt, AsyncWriteExt as _};

    ensure_parent_dir(path).await?;
    let mut file = fs::OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(false)
        .open(path)
        .await?;
    file.seek(std::io::SeekFrom::Start(offset)).await?;

    let mut written: u64 = 0;
    loop {
        if written == length {
            break;
        }
        match framing.read_message(reader).await? {
            Some(Message::Data { payload }) => {
                if written + payload.len() as u64 > length {
                    bail!(
                        "put_chunked: client sent {} bytes past the declared {}",
                        written + payload.len() as u64 - length,
                        length
                    );
                }
                file.write_all(&payload).await?;
                written += payload.len() as u64;
            }
            Some(Message::DataCompressed {
                algo,
                original_size,
                payload,
            }) => {
                let decoded = compress::decode_block(algo, original_size, &payload)?;
                if written + decoded.len() as u64 > length {
                    bail!(
                        "put_chunked: client sent {} bytes past the declared {}",
                        written + decoded.len() as u64 - length,
                        length
                    );
                }
                file.write_all(&decoded).await?;
                written += decoded.len() as u64;
            }
            Some(Message::Done) => break,
            Some(other) => bail!("put_chunked: unexpected message {other:?}"),
            None => bail!("put_chunked: client closed connection before Done"),
        }
    }
    if written != length {
        bail!(
            "put_chunked: declared {} bytes, received {}",
            length,
            written
        );
    }
    if sync {
        file.sync_all().await?;
    }
    Ok(Message::Ok { hash: None })
}

async fn handle_get_chunked<W>(
    path: &str,
    offset: u64,
    length: u64,
    algo: CompressionAlgo,
    out: &mut W,
    framing: &mut Framing,
) -> Result<()>
where
    W: tokio::io::AsyncWrite + Unpin,
{
    use tokio::io::{AsyncReadExt, AsyncSeekExt};

    let mut file = fs::File::open(path).await?;
    file.seek(std::io::SeekFrom::Start(offset)).await?;

    let mut remaining = length;
    let mut buf = vec![0u8; CHUNK_SIZE];
    while remaining > 0 {
        let want = remaining.min(CHUNK_SIZE as u64) as usize;
        let n = file.read(&mut buf[..want]).await?;
        if n == 0 {
            bail!(
                "get_chunked: unexpected EOF at offset {} (still needed {} bytes)",
                offset + length - remaining,
                remaining
            );
        }
        let frame = compress::encode_block(algo, buf[..n].to_vec());
        framing.write_message(out, &frame).await?;
        remaining -= n as u64;
    }
    framing
        .write_message(out, &Message::Ok { hash: None })
        .await?;
    Ok(())
}

async fn handle_mkdir(path: &str) -> Result<Message> {
    fs::create_dir_all(path).await?;
    Ok(Message::Ok { hash: None })
}

async fn ensure_parent_dir(path: &str) -> Result<()> {
    if let Some(parent) = Path::new(path).parent() {
        if !parent.as_os_str().is_empty() {
            fs::create_dir_all(parent).await?;
        }
    }
    Ok(())
}

async fn handle_truncate(path: &str, size: u64) -> Result<Message> {
    ensure_parent_dir(path).await?;
    let f = fs::OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .open(path)
        .await?;
    f.set_len(size).await?;
    Ok(Message::Ok { hash: None })
}

/// Overridable via `BCMR_RENDEZVOUS_TIMEOUT_SECS` for tests.
const RENDEZVOUS_ACCEPT_TIMEOUT_SECS: u64 = 30;

/// Per-connection read budget, sub-leased from the accept deadline so
/// hostile `connect+stall` attempts can't drain it.
const AUTH_HELLO_READ_TIMEOUT_SECS: u64 = 2;

const MAX_RENDEZVOUS_PER_SESSION: usize = 16;

fn rendezvous_accept_timeout() -> std::time::Duration {
    let secs = std::env::var("BCMR_RENDEZVOUS_TIMEOUT_SECS")
        .ok()
        .and_then(|s| s.parse::<u64>().ok())
        .unwrap_or(RENDEZVOUS_ACCEPT_TIMEOUT_SECS);
    std::time::Duration::from_secs(secs)
}

/// Happy path: `drain_gracefully` awaits running data sessions to
/// completion. Error path: Drop aborts orphans so a `?`-propagated
/// error doesn't leak listeners past the control session's lifetime.
struct RendezvousTasks {
    handles: Vec<tokio::task::JoinHandle<()>>,
}

impl RendezvousTasks {
    fn new() -> Self {
        Self {
            handles: Vec::new(),
        }
    }

    fn push(&mut self, h: tokio::task::JoinHandle<()>) {
        self.handles.push(h);
    }

    fn len(&self) -> usize {
        self.handles.len()
    }

    async fn drain_gracefully(mut self) {
        for h in self.handles.drain(..) {
            let _ = h.await;
        }
    }
}

impl Drop for RendezvousTasks {
    fn drop(&mut self) {
        for h in &self.handles {
            h.abort();
        }
    }
}

/// Domain-separation tag for the AuthHello MAC — must match the
/// client side.
const AUTH_HELLO_TAG: &[u8] = b"bcmr-direct-v1";

/// Bind on the SSH-facing interface (from `$SSH_CONNECTION.server_ip`)
/// so the listener is only exposed on the network the SSH session
/// already traversed — not 0.0.0.0 across every local interface.
fn rendezvous_bind_ip() -> std::net::IpAddr {
    use std::net::{IpAddr, Ipv4Addr};
    let parsed = std::env::var("SSH_CONNECTION").ok().and_then(|v| {
        let parts: Vec<&str> = v.split_whitespace().collect();
        if parts.len() == 4 {
            parts[2].parse::<IpAddr>().ok()
        } else {
            None
        }
    });
    parsed.unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST))
}

fn handle_open_direct_channel(root: PathBuf) -> Result<(Message, tokio::task::JoinHandle<()>)> {
    use ring::rand::{SecureRandom, SystemRandom};
    use zeroize::Zeroizing;

    // INVARIANT: sync body — adding `.await` here re-introduces the
    // type-level recursion run_session → this → run_direct_session
    // → run_session that won't compile. Use spawn_blocking or
    // Box::pin if you need async bind work.
    let bind_addr = std::net::SocketAddr::new(rendezvous_bind_ip(), 0);
    let std_listener = std::net::TcpListener::bind(bind_addr)?;
    std_listener.set_nonblocking(true)?;
    let addr = std_listener.local_addr()?;
    let listener = tokio::net::TcpListener::from_std(std_listener)?;

    let rng = SystemRandom::new();
    let mut session_key: Zeroizing<[u8; 32]> = Zeroizing::new([0u8; 32]);
    rng.fill(session_key.as_mut())
        .map_err(|_| anyhow::anyhow!("ring::rand failed to produce session key"))?;
    let key_out = *session_key;

    let handle = tokio::spawn(run_rendezvous(listener, session_key, root));

    Ok((
        Message::DirectChannelReady {
            addr: addr.to_string(),
            session_key: key_out,
        },
        handle,
    ))
}

async fn run_rendezvous(
    listener: tokio::net::TcpListener,
    session_key: zeroize::Zeroizing<[u8; 32]>,
    root: PathBuf,
) {
    // Loop-accept-verify until either an authenticated peer dials in
    // or the outer deadline elapses — a squatter that flunks AuthHello
    // must not consume the single rendezvous slot.
    let deadline = tokio::time::Instant::now() + rendezvous_accept_timeout();
    let authed = loop {
        let remaining = match deadline.checked_duration_since(tokio::time::Instant::now()) {
            Some(d) if !d.is_zero() => d,
            _ => return,
        };
        let mut stream = match tokio::time::timeout(remaining, listener.accept()).await {
            Ok(Ok((stream, _peer))) => stream,
            Ok(Err(e)) => {
                eprintln!("serve: direct-tcp accept failed: {e}");
                return;
            }
            Err(_) => return,
        };
        if verify_auth_hello(&mut stream, &session_key).await {
            break stream;
        }
    };
    drop(listener);

    if let Err(e) = run_direct_session(authed, &session_key, &root).await {
        eprintln!("serve: direct-tcp session error: {e}");
    }
}

async fn verify_auth_hello(stream: &mut tokio::net::TcpStream, session_key: &[u8; 32]) -> bool {
    let read_budget = std::time::Duration::from_secs(AUTH_HELLO_READ_TIMEOUT_SECS);
    let mac = match tokio::time::timeout(read_budget, protocol::read_message(stream)).await {
        Ok(Ok(Some(Message::AuthHello { mac }))) => mac,
        _ => return false,
    };
    // blake3::Hash::PartialEq is constant-time.
    blake3::Hash::from(mac) == blake3::keyed_hash(session_key, AUTH_HELLO_TAG)
}

async fn run_direct_session(
    stream: tokio::net::TcpStream,
    session_key: &[u8; 32],
    root: &Path,
) -> Result<()> {
    let (mut reader, mut writer) = stream.into_split();

    // Box::pin breaks the type-level recursion between run_session
    // and run_direct_session. Flags: no splice (fd != STDOUT_FILENO),
    // no nested direct-TCP, session_key present so CAP_AEAD is offered.
    Box::pin(run_session(
        &mut reader,
        &mut writer,
        root,
        false,
        false,
        Some(session_key),
    ))
    .await
}

async fn handle_resume(path: &str) -> Result<Message> {
    let meta = match fs::metadata(path).await {
        Ok(m) => m,
        Err(_) => {
            return Ok(Message::ResumeResponse {
                size: 0,
                block_hash: None,
            });
        }
    };

    let size = meta.len();
    if size < CHUNK_SIZE as u64 {
        return Ok(Message::ResumeResponse {
            size,
            block_hash: None,
        });
    }

    // Hash the last complete 4 MiB block.
    let block_start = (size / CHUNK_SIZE as u64 - 1) * CHUNK_SIZE as u64;
    let path = path.to_owned();
    let block_hash = tokio::task::spawn_blocking(move || -> Result<String> {
        use std::io::{Read, Seek};
        let mut file = std::fs::File::open(&path)?;
        file.seek(std::io::SeekFrom::Start(block_start))?;
        let mut buf = vec![0u8; CHUNK_SIZE];
        let mut hasher = blake3::Hasher::new();
        let n = file.read(&mut buf)?;
        hasher.update(&buf[..n]);
        Ok(hasher.finalize().to_hex().to_string())
    })
    .await??;

    Ok(Message::ResumeResponse {
        size,
        block_hash: Some(block_hash),
    })
}
