use std::io;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

pub const PROTOCOL_VERSION: u8 = 1;

const TYPE_HELLO: u8 = 0x01;
const TYPE_LIST: u8 = 0x02;
const TYPE_STAT: u8 = 0x03;
const TYPE_HASH: u8 = 0x04;
const TYPE_GET: u8 = 0x05;
const TYPE_PUT: u8 = 0x06;
const TYPE_MKDIR: u8 = 0x07;
const TYPE_RESUME: u8 = 0x08;
const TYPE_DONE: u8 = 0x09;
const TYPE_HAVE_BLOCKS: u8 = 0x0a;
const TYPE_OPEN_DIRECT: u8 = 0x0b;
const TYPE_AUTH_HELLO: u8 = 0x0c;
const TYPE_PUT_CHUNKED: u8 = 0x0d;
const TYPE_GET_CHUNKED: u8 = 0x0e;
const TYPE_TRUNCATE: u8 = 0x0f;

const TYPE_WELCOME: u8 = 0x81;
const TYPE_OK: u8 = 0x82;
const TYPE_ERROR: u8 = 0x83;
const TYPE_DATA: u8 = 0x84;
const TYPE_STAT_RESPONSE: u8 = 0x85;
const TYPE_HASH_RESPONSE: u8 = 0x86;
const TYPE_LIST_RESPONSE: u8 = 0x87;
const TYPE_RESUME_RESPONSE: u8 = 0x88;
const TYPE_DATA_COMPRESSED: u8 = 0x89;
const TYPE_MISSING_BLOCKS: u8 = 0x8a;
const TYPE_DIRECT_READY: u8 = 0x8b;

/// Content-addressed dedup: PUT first exchanges block hashes, only missing
/// blocks go on the wire.
pub const CAP_DEDUP: u8 = 0x04;

/// Client opts out of server-side BLAKE3 hashing (Ok reply carries no
/// digest). On Linux the server additionally uses splice(2) for the
/// file→stdout path. Client must verify integrity itself if it wants to.
pub const CAP_FAST: u8 = 0x08;

/// Per-file fsync after PUT/GET. Opt-in because a 10000-file batch otherwise
/// pays ~10-20s of fdatasync barrier latency on NVMe. Server fsyncs in
/// handle_put; client fsyncs the GET dst.
pub const CAP_SYNC: u8 = 0x10;

/// Direct-TCP transport. Rejected on either side that doesn't set it, so
/// old clients/servers never attempt the rendezvous path.
pub const CAP_DIRECT_TCP: u8 = 0x20;

/// AES-256-GCM on every post-Welcome message on the direct-TCP data plane.
/// Server masks this off on transports without a session key (SSH control,
/// raw listen). Hello/Welcome itself is always plain.
pub const CAP_AEAD: u8 = 0x40;

/// Caps is an optional trailing byte after the version. A peer that doesn't
/// understand caps sends a shorter Hello/Welcome; decoders treat absence
/// as "no caps supported" for automatic v1↔v2 fallback.
pub const CAP_LZ4: u8 = 0x01;
pub const CAP_ZSTD: u8 = 0x02;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionAlgo {
    None,
    Lz4,
    Zstd,
}

impl CompressionAlgo {
    pub fn from_byte(b: u8) -> Self {
        match b {
            1 => CompressionAlgo::Lz4,
            2 => CompressionAlgo::Zstd,
            _ => CompressionAlgo::None,
        }
    }

    pub fn to_byte(self) -> u8 {
        match self {
            CompressionAlgo::None => 0,
            CompressionAlgo::Lz4 => 1,
            CompressionAlgo::Zstd => 2,
        }
    }

    /// Zstd wins if both offer it (better ratio on file content); LZ4 is
    /// the fallback when one side lacks zstd.
    pub fn negotiate(local: u8, remote: u8) -> Self {
        let both = local & remote;
        if both & CAP_ZSTD != 0 {
            CompressionAlgo::Zstd
        } else if both & CAP_LZ4 != 0 {
            CompressionAlgo::Lz4
        } else {
            CompressionAlgo::None
        }
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct ListEntry {
    pub path: String,
    pub size: u64,
    pub is_dir: bool,
}

#[derive(Debug, PartialEq, Clone)]
pub enum Message {
    Hello {
        version: u8,
        caps: u8,
    },
    List {
        path: String,
    },
    Stat {
        path: String,
    },
    Hash {
        path: String,
        offset: u64,
        limit: Option<u64>,
    },
    Get {
        path: String,
        offset: u64,
    },
    Put {
        path: String,
        size: u64,
    },
    Mkdir {
        path: String,
    },
    Resume {
        path: String,
    },
    Done,
    OpenDirectChannel,
    AuthHello {
        mac: [u8; 32],
    },
    /// One slice of a striped PUT. Server reply `Ok { hash: None }`;
    /// integrity comes from the AEAD per-frame MAC.
    PutChunked {
        path: String,
        offset: u64,
        length: u64,
    },
    GetChunked {
        path: String,
        offset: u64,
        length: u64,
    },
    /// Preamble to a striped PUT: create-or-truncate to exactly `size`.
    /// Handles "new src smaller than existing dst" and zero-byte src,
    /// which the chunk fanout can't cover.
    Truncate {
        path: String,
        size: u64,
    },

    Welcome {
        version: u8,
        caps: u8,
    },
    Ok {
        hash: Option<String>,
    },
    Error {
        message: String,
    },
    Data {
        payload: Vec<u8>,
    },
    /// `algo` matches `CompressionAlgo::to_byte`; seeing this without a
    /// negotiated compression cap is a protocol error.
    DataCompressed {
        algo: u8,
        original_size: u32,
        payload: Vec<u8>,
    },
    /// Block hashes in source order; last entry may be a shorter tail block.
    HaveBlocks {
        block_size: u32,
        hashes: Vec<[u8; 32]>,
    },
    /// Bit i (LSB-first in byte i/8) set iff the server lacks hashes[i].
    MissingBlocks {
        bits: Vec<u8>,
    },
    StatResponse {
        size: u64,
        mtime: i64,
        is_dir: bool,
    },
    HashResponse {
        hash: String,
    },
    ListResponse {
        entries: Vec<ListEntry>,
    },
    ResumeResponse {
        size: u64,
        block_hash: Option<String>,
    },
    DirectChannelReady {
        addr: String,
        session_key: [u8; 32],
    },
}

fn write_u8(buf: &mut Vec<u8>, v: u8) {
    buf.push(v);
}

fn write_u32_le(buf: &mut Vec<u8>, v: u32) {
    buf.extend_from_slice(&v.to_le_bytes());
}

fn write_u64_le(buf: &mut Vec<u8>, v: u64) {
    buf.extend_from_slice(&v.to_le_bytes());
}

fn write_i64_le(buf: &mut Vec<u8>, v: i64) {
    buf.extend_from_slice(&v.to_le_bytes());
}

fn write_string(buf: &mut Vec<u8>, s: &str) {
    let bytes = s.as_bytes();
    write_u32_le(buf, bytes.len() as u32);
    buf.extend_from_slice(bytes);
}

fn write_bytes(buf: &mut Vec<u8>, data: &[u8]) {
    write_u32_le(buf, data.len() as u32);
    buf.extend_from_slice(data);
}

fn write_opt_string(buf: &mut Vec<u8>, opt: &Option<String>) {
    match opt {
        Some(s) => {
            write_u8(buf, 1);
            write_string(buf, s);
        }
        None => write_u8(buf, 0),
    }
}

fn write_opt_u64(buf: &mut Vec<u8>, opt: &Option<u64>) {
    match opt {
        Some(v) => {
            write_u8(buf, 1);
            write_u64_le(buf, *v);
        }
        None => write_u8(buf, 0),
    }
}

fn write_list_entry(buf: &mut Vec<u8>, entry: &ListEntry) {
    write_string(buf, &entry.path);
    write_u64_le(buf, entry.size);
    write_u8(buf, entry.is_dir as u8);
}

/// Frame layout: `[4 LE payload_len][1 type][payload...]`.
pub fn encode_message(msg: &Message) -> Vec<u8> {
    let mut payload = Vec::new();

    match msg {
        Message::Hello { version, caps } => {
            write_u8(&mut payload, TYPE_HELLO);
            write_u8(&mut payload, *version);
            write_u8(&mut payload, *caps);
        }
        Message::List { path } => {
            write_u8(&mut payload, TYPE_LIST);
            write_string(&mut payload, path);
        }
        Message::Stat { path } => {
            write_u8(&mut payload, TYPE_STAT);
            write_string(&mut payload, path);
        }
        Message::Hash {
            path,
            offset,
            limit,
        } => {
            write_u8(&mut payload, TYPE_HASH);
            write_string(&mut payload, path);
            write_u64_le(&mut payload, *offset);
            write_opt_u64(&mut payload, limit);
        }
        Message::Get { path, offset } => {
            write_u8(&mut payload, TYPE_GET);
            write_string(&mut payload, path);
            write_u64_le(&mut payload, *offset);
        }
        Message::Put { path, size } => {
            write_u8(&mut payload, TYPE_PUT);
            write_string(&mut payload, path);
            write_u64_le(&mut payload, *size);
        }
        Message::Mkdir { path } => {
            write_u8(&mut payload, TYPE_MKDIR);
            write_string(&mut payload, path);
        }
        Message::Resume { path } => {
            write_u8(&mut payload, TYPE_RESUME);
            write_string(&mut payload, path);
        }
        Message::Done => {
            write_u8(&mut payload, TYPE_DONE);
        }
        Message::Welcome { version, caps } => {
            write_u8(&mut payload, TYPE_WELCOME);
            write_u8(&mut payload, *version);
            write_u8(&mut payload, *caps);
        }
        Message::Ok { hash } => {
            write_u8(&mut payload, TYPE_OK);
            write_opt_string(&mut payload, hash);
        }
        Message::Error { message } => {
            write_u8(&mut payload, TYPE_ERROR);
            write_string(&mut payload, message);
        }
        Message::Data { payload: data } => {
            write_u8(&mut payload, TYPE_DATA);
            write_bytes(&mut payload, data);
        }
        Message::DataCompressed {
            algo,
            original_size,
            payload: data,
        } => {
            write_u8(&mut payload, TYPE_DATA_COMPRESSED);
            write_u8(&mut payload, *algo);
            write_u32_le(&mut payload, *original_size);
            write_bytes(&mut payload, data);
        }
        Message::HaveBlocks { block_size, hashes } => {
            write_u8(&mut payload, TYPE_HAVE_BLOCKS);
            write_u32_le(&mut payload, *block_size);
            write_u32_le(&mut payload, hashes.len() as u32);
            for h in hashes {
                payload.extend_from_slice(h);
            }
        }
        Message::MissingBlocks { bits } => {
            write_u8(&mut payload, TYPE_MISSING_BLOCKS);
            write_bytes(&mut payload, bits);
        }
        Message::StatResponse {
            size,
            mtime,
            is_dir,
        } => {
            write_u8(&mut payload, TYPE_STAT_RESPONSE);
            write_u64_le(&mut payload, *size);
            write_i64_le(&mut payload, *mtime);
            write_u8(&mut payload, *is_dir as u8);
        }
        Message::HashResponse { hash } => {
            write_u8(&mut payload, TYPE_HASH_RESPONSE);
            write_string(&mut payload, hash);
        }
        Message::ListResponse { entries } => {
            write_u8(&mut payload, TYPE_LIST_RESPONSE);
            write_u32_le(&mut payload, entries.len() as u32);
            for entry in entries {
                write_list_entry(&mut payload, entry);
            }
        }
        Message::ResumeResponse { size, block_hash } => {
            write_u8(&mut payload, TYPE_RESUME_RESPONSE);
            write_u64_le(&mut payload, *size);
            write_opt_string(&mut payload, block_hash);
        }
        Message::OpenDirectChannel => {
            write_u8(&mut payload, TYPE_OPEN_DIRECT);
        }
        Message::AuthHello { mac } => {
            write_u8(&mut payload, TYPE_AUTH_HELLO);
            payload.extend_from_slice(mac);
        }
        Message::DirectChannelReady { addr, session_key } => {
            write_u8(&mut payload, TYPE_DIRECT_READY);
            write_string(&mut payload, addr);
            payload.extend_from_slice(session_key);
        }
        Message::PutChunked {
            path,
            offset,
            length,
        } => {
            write_u8(&mut payload, TYPE_PUT_CHUNKED);
            write_string(&mut payload, path);
            write_u64_le(&mut payload, *offset);
            write_u64_le(&mut payload, *length);
        }
        Message::GetChunked {
            path,
            offset,
            length,
        } => {
            write_u8(&mut payload, TYPE_GET_CHUNKED);
            write_string(&mut payload, path);
            write_u64_le(&mut payload, *offset);
            write_u64_le(&mut payload, *length);
        }
        Message::Truncate { path, size } => {
            write_u8(&mut payload, TYPE_TRUNCATE);
            write_string(&mut payload, path);
            write_u64_le(&mut payload, *size);
        }
    }

    let mut frame = Vec::with_capacity(4 + payload.len());
    write_u32_le(&mut frame, payload.len() as u32);
    frame.extend_from_slice(&payload);
    frame
}

struct Cursor<'a> {
    data: &'a [u8],
    pos: usize,
}

impl<'a> Cursor<'a> {
    fn new(data: &'a [u8]) -> Self {
        Self { data, pos: 0 }
    }

    fn read_u8(&mut self) -> Option<u8> {
        if self.pos >= self.data.len() {
            return None;
        }
        let v = self.data[self.pos];
        self.pos += 1;
        Some(v)
    }

    fn read_u32_le(&mut self) -> Option<u32> {
        let bytes = self.data.get(self.pos..self.pos + 4)?;
        self.pos += 4;
        Some(u32::from_le_bytes(bytes.try_into().unwrap()))
    }

    fn read_u64_le(&mut self) -> Option<u64> {
        let bytes = self.data.get(self.pos..self.pos + 8)?;
        self.pos += 8;
        Some(u64::from_le_bytes(bytes.try_into().unwrap()))
    }

    fn read_i64_le(&mut self) -> Option<i64> {
        let bytes = self.data.get(self.pos..self.pos + 8)?;
        self.pos += 8;
        Some(i64::from_le_bytes(bytes.try_into().unwrap()))
    }

    fn read_string(&mut self) -> Option<String> {
        let len = self.read_u32_le()? as usize;
        let bytes = self.data.get(self.pos..self.pos + len)?;
        self.pos += len;
        String::from_utf8(bytes.to_vec()).ok()
    }

    fn read_bytes(&mut self) -> Option<Vec<u8>> {
        let len = self.read_u32_le()? as usize;
        let bytes = self.data.get(self.pos..self.pos + len)?;
        self.pos += len;
        Some(bytes.to_vec())
    }

    fn read_fixed<const N: usize>(&mut self) -> Option<[u8; N]> {
        let slice = self.data.get(self.pos..self.pos + N)?;
        self.pos += N;
        let mut out = [0u8; N];
        out.copy_from_slice(slice);
        Some(out)
    }

    fn read_opt_string(&mut self) -> Option<Option<String>> {
        let present = self.read_u8()?;
        if present == 1 {
            Some(Some(self.read_string()?))
        } else {
            Some(None)
        }
    }

    fn read_opt_u64(&mut self) -> Option<Option<u64>> {
        let present = self.read_u8()?;
        if present == 1 {
            Some(Some(self.read_u64_le()?))
        } else {
            Some(None)
        }
    }

    fn read_list_entry(&mut self) -> Option<ListEntry> {
        let path = self.read_string()?;
        let size = self.read_u64_le()?;
        let is_dir = self.read_u8()? != 0;
        Some(ListEntry { path, size, is_dir })
    }
}

/// `None` on empty input, truncated frames, or unknown/malformed messages.
pub fn decode_message(data: &[u8]) -> Option<Message> {
    if data.is_empty() {
        return None;
    }

    let mut c = Cursor::new(data);
    let payload_len = c.read_u32_le()? as usize;
    let payload = data.get(c.pos..c.pos + payload_len)?;

    let mut p = Cursor::new(payload);
    let msg_type = p.read_u8()?;

    let msg = match msg_type {
        TYPE_HELLO => Message::Hello {
            version: p.read_u8()?,
            caps: p.read_u8().unwrap_or(0),
        },
        TYPE_LIST => Message::List {
            path: p.read_string()?,
        },
        TYPE_STAT => Message::Stat {
            path: p.read_string()?,
        },
        TYPE_HASH => Message::Hash {
            path: p.read_string()?,
            offset: p.read_u64_le()?,
            limit: p.read_opt_u64()?,
        },
        TYPE_GET => Message::Get {
            path: p.read_string()?,
            offset: p.read_u64_le()?,
        },
        TYPE_PUT => Message::Put {
            path: p.read_string()?,
            size: p.read_u64_le()?,
        },
        TYPE_MKDIR => Message::Mkdir {
            path: p.read_string()?,
        },
        TYPE_RESUME => Message::Resume {
            path: p.read_string()?,
        },
        TYPE_DONE => Message::Done,
        TYPE_WELCOME => Message::Welcome {
            version: p.read_u8()?,
            caps: p.read_u8().unwrap_or(0),
        },
        TYPE_OK => Message::Ok {
            hash: p.read_opt_string()?,
        },
        TYPE_ERROR => Message::Error {
            message: p.read_string()?,
        },
        TYPE_DATA => Message::Data {
            payload: p.read_bytes()?,
        },
        TYPE_DATA_COMPRESSED => Message::DataCompressed {
            algo: p.read_u8()?,
            original_size: p.read_u32_le()?,
            payload: p.read_bytes()?,
        },
        TYPE_HAVE_BLOCKS => {
            let block_size = p.read_u32_le()?;
            let count = p.read_u32_le()? as usize;
            let mut hashes = Vec::with_capacity(count);
            for _ in 0..count {
                let mut h = [0u8; 32];
                for byte in &mut h {
                    *byte = p.read_u8()?;
                }
                hashes.push(h);
            }
            Message::HaveBlocks { block_size, hashes }
        }
        TYPE_MISSING_BLOCKS => Message::MissingBlocks {
            bits: p.read_bytes()?,
        },
        TYPE_STAT_RESPONSE => Message::StatResponse {
            size: p.read_u64_le()?,
            mtime: p.read_i64_le()?,
            is_dir: p.read_u8()? != 0,
        },
        TYPE_HASH_RESPONSE => Message::HashResponse {
            hash: p.read_string()?,
        },
        TYPE_LIST_RESPONSE => {
            let count = p.read_u32_le()? as usize;
            let mut entries = Vec::with_capacity(count);
            for _ in 0..count {
                entries.push(p.read_list_entry()?);
            }
            Message::ListResponse { entries }
        }
        TYPE_RESUME_RESPONSE => Message::ResumeResponse {
            size: p.read_u64_le()?,
            block_hash: p.read_opt_string()?,
        },
        TYPE_OPEN_DIRECT => Message::OpenDirectChannel,
        TYPE_AUTH_HELLO => Message::AuthHello {
            mac: p.read_fixed::<32>()?,
        },
        TYPE_DIRECT_READY => Message::DirectChannelReady {
            addr: p.read_string()?,
            session_key: p.read_fixed::<32>()?,
        },
        TYPE_PUT_CHUNKED => Message::PutChunked {
            path: p.read_string()?,
            offset: p.read_u64_le()?,
            length: p.read_u64_le()?,
        },
        TYPE_GET_CHUNKED => Message::GetChunked {
            path: p.read_string()?,
            offset: p.read_u64_le()?,
            length: p.read_u64_le()?,
        },
        TYPE_TRUNCATE => Message::Truncate {
            path: p.read_string()?,
            size: p.read_u64_le()?,
        },
        _ => return None,
    };

    Some(msg)
}

pub async fn write_message<W: AsyncWriteExt + Unpin>(
    writer: &mut W,
    msg: &Message,
) -> io::Result<()> {
    // DirectChannelReady carries a session key; Zeroizing scrubs the
    // plaintext from the heap after the write (defense-in-depth on top of
    // the authenticated SSH channel the frame travels through).
    if matches!(msg, Message::DirectChannelReady { .. }) {
        let frame = zeroize::Zeroizing::new(encode_message(msg));
        writer.write_all(&frame).await
    } else {
        let frame = encode_message(msg);
        writer.write_all(&frame).await
    }
}

/// `Ok(None)` on clean EOF (zero bytes read for the length prefix).
pub async fn read_message<R: AsyncReadExt + Unpin>(reader: &mut R) -> io::Result<Option<Message>> {
    let mut len_buf = [0u8; 4];
    match reader.read_exact(&mut len_buf).await {
        Ok(_) => {}
        Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => return Ok(None),
        Err(e) => return Err(e),
    }

    let payload_len = u32::from_le_bytes(len_buf) as usize;

    const MAX_FRAME_SIZE: usize = 16 * 1024 * 1024;
    if payload_len > MAX_FRAME_SIZE {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "frame too large: {} bytes (max {})",
                payload_len, MAX_FRAME_SIZE
            ),
        ));
    }

    let mut payload = vec![0u8; payload_len];
    reader.read_exact(&mut payload).await?;

    let mut frame = Vec::with_capacity(4 + payload_len);
    frame.extend_from_slice(&len_buf);
    frame.extend_from_slice(&payload);

    decode_message(&frame)
        .map(Some)
        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "malformed protocol message"))
}
