pub mod response;
pub use response::Response;

use crate::Peer;
use cyphernet::addr::{HostName, PartialAddr};
use rand::RngExt;
use std::{
    error::Error,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket},
};

pub trait Server {
    fn announce(
        &self,
        info_hash: &[u8; 20],
        tracker: &SocketAddr,
        peers_buffer: Option<&mut Vec<Peer>>,
    ) -> Result<Response, Box<dyn Error>>;
}

impl Server for UdpSocket {
    fn announce(
        &self,
        id20: &[u8; 20],
        tracker: &SocketAddr,
        peers_buffer: Option<&mut Vec<Peer>>,
    ) -> Result<Response, Box<dyn Error>> {
        // Setup environment...
        let mut rng = rand::rng();
        let transaction_id: u32 = rng.random();

        let mut c = Vec::with_capacity(16);
        c.extend_from_slice(&0x41727101980u64.to_be_bytes()); // magic const
        c.extend_from_slice(&0u32.to_be_bytes());
        c.extend_from_slice(&transaction_id.to_be_bytes());

        // Initial connection...
        self.send_to(&c, tracker)?;

        let mut buf = [0u8; 2048];
        let (amt, _) = self.recv_from(&mut buf)?;

        if amt < 16 {
            return Err("Tracker response too short".into());
        }

        let action = u32::from_be_bytes(buf[0..4].try_into()?);
        let res_transaction_id = u32::from_be_bytes(buf[4..8].try_into()?);

        if res_transaction_id != transaction_id {
            return Err("Transaction ID missmatch".into());
        }

        if action == 3 {
            return Err(
                format!("Tracker error: `{}`", String::from_utf8_lossy(&buf[8..amt])).into(),
            );
        }

        let connection_id = u64::from_be_bytes(buf[8..16].try_into()?);

        // Sending announce request...
        let mut a = Vec::with_capacity(98);
        let announce_trans_id: u32 = rng.random();
        let mut peer_id = [0u8; 20];
        rng.fill(&mut peer_id);

        a.extend_from_slice(&connection_id.to_be_bytes());
        a.extend_from_slice(&1u32.to_be_bytes()); // action = 1 (announce)
        a.extend_from_slice(&announce_trans_id.to_be_bytes());
        a.extend_from_slice(id20);
        a.extend_from_slice(&peer_id);
        a.extend_from_slice(&0u64.to_be_bytes()); // downloaded
        a.extend_from_slice(&1024u64.to_be_bytes()); // left
        a.extend_from_slice(&0u64.to_be_bytes()); // uploaded
        a.extend_from_slice(&0u32.to_be_bytes()); // event = 0 (none)
        a.extend_from_slice(&0u32.to_be_bytes()); // IP address = 0 (default)
        a.extend_from_slice(&0u32.to_be_bytes()); // key
        a.extend_from_slice(&(-1i32).to_be_bytes()); // num want = -1
        a.extend_from_slice(&self.local_addr()?.port().to_be_bytes());

        self.send_to(&a, tracker)?;

        let (amt, _) = self.recv_from(&mut buf)?;

        let res_action = u32::from_be_bytes(buf[0..4].try_into()?);
        let res_announce_trans_id = u32::from_be_bytes(buf[4..8].try_into()?);

        if res_announce_trans_id != announce_trans_id {
            return Err("Announce Transaction ID does not match.".into());
        }

        if res_action == 3 {
            return Err(format!(
                "Tracker announce declined: {}",
                String::from_utf8_lossy(&buf[8..amt])
            )
            .into());
        }

        if amt < 20 {
            return Err("Unexpected prefix len".into());
        }

        // Extract peers info...
        let interval = u32::from_be_bytes(buf[8..12].try_into()?);
        let leechers = u32::from_be_bytes(buf[12..16].try_into()?);
        let seeders = u32::from_be_bytes(buf[16..20].try_into()?);

        let peers_data = &buf[20..amt];

        if let Some(peers) = peers_buffer {
            if tracker.is_ipv6() {
                for chunk in peers_data.chunks_exact(18) {
                    let ip_bytes: [u8; 16] = chunk[0..16].try_into()?;
                    peers.push(PartialAddr {
                        host: HostName::Ip(IpAddr::V6(Ipv6Addr::from(ip_bytes))),
                        port: Some(u16::from_be_bytes(chunk[16..18].try_into()?)),
                    })
                }
            } else {
                for chunk in peers_data.chunks_exact(6) {
                    let ip_bytes: [u8; 4] = chunk[0..4].try_into()?;
                    peers.push(PartialAddr {
                        host: HostName::Ip(IpAddr::V4(Ipv4Addr::from(ip_bytes))),
                        port: Some(u16::from_be_bytes(chunk[4..6].try_into()?)),
                    })
                }
            }
        }

        // Done.
        Ok(Response {
            interval,
            leechers,
            seeders,
        })
    }
}
