//! Sans-io BitTorrent session engine — the download state machine.
//!
//! The engine is **transport-free** (sans-io): it consumes typed
//! [`Event`]s (a peer connected, a peer message arrived, a block was
//! written to disk, a timer ticked) and emits typed [`Action`]s (send
//! this message, write this block, this piece verified). The daemon
//! (`enxamed`) provides the actual sockets and disk; tests drive the
//! same `on_event` with synthetic events. This is the pleme-io
//! interpreter idiom — a pure `(state, event) -> (state, actions)` core,
//! the same shape as mado's overlay FSM and the sui-spec interpreters.
//!
//! Scope (M3): the **leech** path — connect, exchange bitfields, become
//! interested, pipeline block requests rarest-first, assemble + SHA-1
//! verify pieces, announce completion. Seeding/choke-algorithm and the
//! endgame refinement layer onto this core (`theory/ENXAME.md` M4).

#![forbid(unsafe_code)]

use std::collections::{BTreeMap, HashMap};

use enxame_bitfield::Bitfield;
use enxame_metainfo::InfoDict;
use enxame_peer::{Block, PeerMessage};
use sha1::{Digest, Sha1};

/// 16 KiB — the canonical block size BitTorrent peers request in.
pub const BLOCK_LEN: u32 = 16 * 1024;
/// How many block requests to keep in flight per peer (the pipeline).
pub const PIPELINE_DEPTH: usize = 8;

/// A peer handle the engine assigns; the daemon maps it to a socket.
pub type PeerKey = u32;

/// Something happened that the engine must react to.
#[derive(Debug, Clone)]
pub enum Event {
    /// A peer connection completed its handshake and is ready.
    PeerReady(PeerKey),
    /// A decoded peer-wire message arrived from a peer.
    Message(PeerKey, PeerMessage),
    /// A peer's connection dropped.
    PeerGone(PeerKey),
    /// A block we requested was read off the wire and is ready to place
    /// (the daemon may also persist it; the engine tracks completion).
    BlockArrived { index: u32, begin: u32, data: Vec<u8> },
}

/// Something the daemon must do on the engine's behalf.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Action {
    /// Send a wire message to a peer.
    Send(PeerKey, PeerMessage),
    /// A piece passed SHA-1 verification — persist + announce `Have`.
    PieceVerified(u32),
    /// A piece failed verification — its blocks were dropped, re-pick.
    PieceCorrupt(u32),
    /// Every piece is verified — the torrent is complete.
    Complete,
    /// Drop a misbehaving peer.
    Disconnect(PeerKey),
}

/// Per-peer protocol state (the four BEP-3 booleans + their bitfield +
/// our in-flight requests to them).
#[derive(Debug, Clone)]
struct Peer {
    am_choking: bool,
    am_interested: bool,
    peer_choking: bool,
    #[allow(dead_code)] // tracked for the future choke algorithm
    peer_interested: bool,
    has: Bitfield,
    /// Blocks we've requested from this peer and not yet received.
    in_flight: Vec<Block>,
}

impl Peer {
    fn new(num_pieces: usize) -> Self {
        Self {
            am_choking: true,
            am_interested: false,
            peer_choking: true,
            peer_interested: false,
            has: Bitfield::new(num_pieces),
            in_flight: Vec::new(),
        }
    }
}

/// In-progress assembly of one piece: the blocks received so far.
#[derive(Debug, Clone)]
struct PieceAssembly {
    /// `begin` → block bytes, kept sorted so we can concatenate.
    blocks: BTreeMap<u32, Vec<u8>>,
    received: u32,
    total: u32,
}

/// The download session for one torrent.
pub struct Session {
    info: InfoDict,
    /// Pieces we have verified.
    have: Bitfield,
    /// Pieces currently being assembled.
    assembling: HashMap<u32, PieceAssembly>,
    /// Per-piece availability across connected peers (rarest-first).
    availability: Vec<u32>,
    peers: HashMap<PeerKey, Peer>,
    num_pieces: usize,
}

impl Session {
    /// Start a session for a torrent we hold no pieces of yet.
    #[must_use]
    pub fn new(info: InfoDict) -> Self {
        let num_pieces = info.piece_count();
        Self {
            have: Bitfield::new(num_pieces),
            assembling: HashMap::new(),
            availability: vec![0; num_pieces],
            peers: HashMap::new(),
            num_pieces,
            info,
        }
    }

    /// Pieces verified so far.
    #[must_use]
    pub fn have_count(&self) -> usize {
        self.have.count()
    }

    /// Total pieces.
    #[must_use]
    pub fn piece_count(&self) -> usize {
        self.num_pieces
    }

    /// Whether the whole torrent is downloaded + verified.
    #[must_use]
    pub fn is_complete(&self) -> bool {
        self.have.is_complete()
    }

    /// The byte length of piece `index` (the last piece is short).
    fn piece_len(&self, index: u32) -> u32 {
        let total = self.info.total_length();
        let plen = self.info.piece_length;
        let start = u64::from(index) * plen;
        u32::try_from((total - start).min(plen)).unwrap_or(u32::MAX)
    }

    /// Drive the state machine with one event; return the actions to
    /// perform. Pure: the only effects are the returned [`Action`]s.
    pub fn on_event(&mut self, event: Event) -> Vec<Action> {
        match event {
            Event::PeerReady(key) => {
                self.peers.entry(key).or_insert_with(|| Peer::new(self.num_pieces));
                // Greet with our bitfield (BEP-3: optional but standard).
                vec![Action::Send(key, PeerMessage::Bitfield(self.have.as_bytes().to_vec()))]
            }
            Event::PeerGone(key) => {
                if let Some(peer) = self.peers.remove(&key) {
                    for i in peer.has.iter_set() {
                        self.availability[i] = self.availability[i].saturating_sub(1);
                    }
                }
                Vec::new()
            }
            Event::Message(key, msg) => self.on_message(key, &msg),
            Event::BlockArrived { index, begin, data } => self.place_block(index, begin, data),
        }
    }

    fn on_message(&mut self, key: PeerKey, msg: &PeerMessage) -> Vec<Action> {
        let mut actions = Vec::new();
        match msg {
            PeerMessage::Choke => {
                if let Some(p) = self.peers.get_mut(&key) {
                    p.peer_choking = true;
                    p.in_flight.clear(); // choked: assume our requests are dropped
                }
            }
            PeerMessage::Unchoke => {
                if let Some(p) = self.peers.get_mut(&key) {
                    p.peer_choking = false;
                }
                actions.extend(self.fill_pipeline(key));
            }
            PeerMessage::Interested => {
                if let Some(p) = self.peers.get_mut(&key) {
                    p.peer_interested = true;
                }
            }
            PeerMessage::NotInterested => {
                if let Some(p) = self.peers.get_mut(&key) {
                    p.peer_interested = false;
                }
            }
            PeerMessage::Have(idx) => {
                if let Some(p) = self.peers.get_mut(&key) {
                    let i = *idx as usize;
                    if i < self.num_pieces && !p.has.has(i) {
                        p.has.set(i);
                        self.availability[i] += 1;
                    }
                }
                actions.extend(self.update_interest(key));
            }
            PeerMessage::Bitfield(bytes) => {
                if let Ok(bf) = Bitfield::from_bytes(bytes, self.num_pieces) {
                    if let Some(p) = self.peers.get_mut(&key) {
                        for i in p.has.iter_set() {
                            self.availability[i] = self.availability[i].saturating_sub(1);
                        }
                        for i in bf.iter_set() {
                            self.availability[i] += 1;
                        }
                        p.has = bf;
                    }
                    actions.extend(self.update_interest(key));
                } else {
                    actions.push(Action::Disconnect(key));
                }
            }
            PeerMessage::Piece { index, begin, data } => {
                // Clear the satisfied in-flight request.
                if let Some(p) = self.peers.get_mut(&key) {
                    p.in_flight.retain(|b| !(b.index == *index && b.begin == *begin));
                }
                actions.extend(self.place_block(*index, *begin, data.clone()));
                actions.extend(self.fill_pipeline(key));
            }
            // Inbound Request/Cancel/Port/Extended/KeepAlive/Unknown: the
            // leech core ignores (seeding + DHT + ext layer on top later).
            _ => {}
        }
        actions
    }

    /// Send `Interested` if the peer has a piece we lack and we aren't
    /// already interested; trigger pipeline fill if already unchoked.
    fn update_interest(&mut self, key: PeerKey) -> Vec<Action> {
        let mut actions = Vec::new();
        let wants = self
            .peers
            .get(&key)
            .is_some_and(|p| p.has.difference(&self.have).next().is_some());
        if let Some(p) = self.peers.get_mut(&key) {
            if wants && !p.am_interested {
                p.am_interested = true;
                actions.push(Action::Send(key, PeerMessage::Interested));
            }
            if !p.peer_choking {
                actions.extend(self.fill_pipeline(key));
            }
        }
        actions
    }

    /// Request blocks from `key` until its pipeline is full, rarest-first
    /// over the pieces it has that we still need.
    fn fill_pipeline(&mut self, key: PeerKey) -> Vec<Action> {
        let mut actions = Vec::new();
        // Snapshot what's requestable from this peer.
        let (peer_choking, slots, candidate_pieces) = {
            let Some(p) = self.peers.get(&key) else { return actions };
            let slots = PIPELINE_DEPTH.saturating_sub(p.in_flight.len());
            let mut pieces: Vec<u32> = (0..self.num_pieces)
                .filter(|&i| !self.have.has(i) && p.has.has(i))
                .map(|i| i as u32)
                .collect();
            // Rarest first.
            pieces.sort_by_key(|&i| self.availability[i as usize]);
            (p.peer_choking, slots, pieces)
        };
        if peer_choking || slots == 0 {
            return actions;
        }
        let mut requested = 0;
        for index in candidate_pieces {
            if requested >= slots {
                break;
            }
            for block in self.needed_blocks(index) {
                if requested >= slots {
                    break;
                }
                let already = self
                    .peers
                    .get(&key)
                    .is_some_and(|p| p.in_flight.iter().any(|b| *b == block));
                if already {
                    continue;
                }
                if let Some(p) = self.peers.get_mut(&key) {
                    p.in_flight.push(block);
                }
                actions.push(Action::Send(key, PeerMessage::Request(block)));
                requested += 1;
            }
        }
        actions
    }

    /// The blocks of `index` we still need (not yet assembled).
    fn needed_blocks(&self, index: u32) -> Vec<Block> {
        let plen = self.piece_len(index);
        let have_offsets: &BTreeMap<u32, Vec<u8>> = self
            .assembling
            .get(&index)
            .map_or(&EMPTY_BLOCKS, |a| &a.blocks);
        let mut blocks = Vec::new();
        let mut begin = 0u32;
        while begin < plen {
            let length = BLOCK_LEN.min(plen - begin);
            if !have_offsets.contains_key(&begin) {
                blocks.push(Block { index, begin, length });
            }
            begin += length;
        }
        blocks
    }

    /// Place a received block; if the piece is now whole, verify it.
    fn place_block(&mut self, index: u32, begin: u32, data: Vec<u8>) -> Vec<Action> {
        if index as usize >= self.num_pieces || self.have.has(index as usize) {
            return Vec::new(); // already have it / out of range
        }
        let plen = self.piece_len(index);
        let entry = self.assembling.entry(index).or_insert_with(|| PieceAssembly {
            blocks: BTreeMap::new(),
            received: 0,
            total: plen,
        });
        if entry.blocks.insert(begin, data).is_none() {
            entry.received += BLOCK_LEN.min(plen - begin);
        }
        if entry.received >= entry.total {
            return self.verify_piece(index);
        }
        Vec::new()
    }

    fn verify_piece(&mut self, index: u32) -> Vec<Action> {
        let Some(assembly) = self.assembling.remove(&index) else {
            return Vec::new();
        };
        let mut buf = Vec::with_capacity(assembly.total as usize);
        for (_, block) in assembly.blocks {
            buf.extend_from_slice(&block);
        }
        let digest: [u8; 20] = Sha1::digest(&buf).into();
        let expected = self.info.pieces.get(index as usize).copied().unwrap_or([0; 20]);
        if digest == expected {
            self.have.set(index as usize);
            let mut actions = vec![Action::PieceVerified(index)];
            // Announce `Have` to all peers.
            for &key in self.peers.keys() {
                actions.push(Action::Send(key, PeerMessage::Have(index)));
            }
            if self.have.is_complete() {
                actions.push(Action::Complete);
            }
            actions
        } else {
            vec![Action::PieceCorrupt(index)]
        }
    }
}

// A shared empty map so `needed_blocks` can borrow when no assembly yet.
use std::sync::LazyLock;
static EMPTY_BLOCKS: LazyLock<BTreeMap<u32, Vec<u8>>> = LazyLock::new(BTreeMap::new);

#[cfg(test)]
mod tests {
    use super::*;
    use enxame_metainfo::Layout;

    /// Build a 2-piece torrent over `data`, piece_length 8.
    fn info_for(data: &[u8], piece_length: u64) -> InfoDict {
        let pieces: Vec<[u8; 20]> = data
            .chunks(piece_length as usize)
            .map(|c| Sha1::digest(c).into())
            .collect();
        InfoDict {
            name: "t".into(),
            piece_length,
            pieces,
            layout: Layout::SingleFile { length: data.len() as u64 },
            private: false,
        }
    }

    #[test]
    fn greets_a_ready_peer_with_a_bitfield() {
        let mut s = Session::new(info_for(b"hello world!", 8));
        let actions = s.on_event(Event::PeerReady(1));
        assert!(matches!(actions.as_slice(), [Action::Send(1, PeerMessage::Bitfield(_))]));
    }

    #[test]
    fn becomes_interested_and_requests_rarest_first() {
        let mut s = Session::new(info_for(b"hello world!", 8)); // 2 pieces (8 + 4)
        s.on_event(Event::PeerReady(1));
        // Peer advertises it has both pieces.
        let mut bf = Bitfield::new(2);
        bf.set(0);
        bf.set(1);
        let acts = s.on_event(Event::Message(1, PeerMessage::Bitfield(bf.as_bytes().to_vec())));
        assert!(acts.contains(&Action::Send(1, PeerMessage::Interested)));
        // Unchoke us → we request blocks.
        let acts = s.on_event(Event::Message(1, PeerMessage::Unchoke));
        let requests: Vec<_> = acts.iter().filter(|a| matches!(a, Action::Send(_, PeerMessage::Request(_)))).collect();
        assert!(!requests.is_empty(), "expected block requests after unchoke");
    }

    #[test]
    fn assembles_and_verifies_a_piece_end_to_end() {
        let data = b"hello world!"; // 12 bytes, piece_length 8 → piece0=8, piece1=4
        let mut s = Session::new(info_for(data, 8));
        s.on_event(Event::PeerReady(1));
        let mut bf = Bitfield::new(2);
        bf.set(0);
        s.on_event(Event::Message(1, PeerMessage::Bitfield(bf.as_bytes().to_vec())));
        s.on_event(Event::Message(1, PeerMessage::Unchoke));
        // Deliver piece 0's single block (8 bytes ≤ BLOCK_LEN).
        let acts = s.on_event(Event::Message(
            1,
            PeerMessage::Piece { index: 0, begin: 0, data: data[..8].to_vec() },
        ));
        assert!(acts.contains(&Action::PieceVerified(0)));
        assert_eq!(s.have_count(), 1);
        assert!(!s.is_complete());
    }

    #[test]
    fn rejects_a_corrupt_piece() {
        let data = b"hello world!";
        let mut s = Session::new(info_for(data, 8));
        s.on_event(Event::PeerReady(1));
        let acts = s.on_event(Event::Message(
            1,
            PeerMessage::Piece { index: 0, begin: 0, data: vec![0xff; 8] }, // wrong bytes
        ));
        assert!(acts.contains(&Action::PieceCorrupt(0)));
        assert_eq!(s.have_count(), 0);
    }

    #[test]
    fn completes_when_all_pieces_verify() {
        let data = b"abcdefgh"; // one 8-byte piece
        let mut s = Session::new(info_for(data, 8));
        s.on_event(Event::PeerReady(1));
        let acts = s.on_event(Event::Message(
            1,
            PeerMessage::Piece { index: 0, begin: 0, data: data.to_vec() },
        ));
        assert!(acts.contains(&Action::Complete));
        assert!(s.is_complete());
    }
}
