use std::net::SocketAddr;
use std::sync::Arc;

use anyhow::{Ok, Result};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::Mutex;

use crate::connection::commands::{
    AckCommand, CommandId, HelloCommand, RejectCommand, WelcomeCommand,
};
use crate::connection::context::Context;
use crate::connection::node::Node;
use crate::constants::tcp::MAX_HANDSHAKE_SIZE;
use crate::crypto::session_key_manager::PairedSessionKey;

pub struct HandshakeHandler {
    pub local_node: Node,
    pub session_keys: Option<Arc<Mutex<PairedSessionKey>>>,
    pub on_established: Option<Arc<dyn Fn(Node, SocketAddr) + Send + Sync>>,
    pub on_rejected: Option<Arc<dyn Fn(String, SocketAddr) + Send + Sync>>,
}

impl HandshakeHandler {
    pub fn new(local_node: Node) -> Self {
        Self {
            local_node,
            session_keys: None,
            on_established: None,
            on_rejected: None,
        }
    }

    #[cfg(test)]
    pub fn new_with_node(node: Node) -> Self {
        Self {
            local_node: node,
            session_keys: None,
            on_established: None,
            on_rejected: None,
        }
    }

    #[cfg(test)]
    pub fn local_node_ref(&self) -> &Node {
        &self.local_node
    }

    pub fn with_session_keys(mut self, keys: Arc<Mutex<PairedSessionKey>>) -> Self {
        self.session_keys = Some(keys);
        self
    }

    pub fn on_established<F>(mut self, callback: F) -> Self
    where
        F: Fn(Node, SocketAddr) + Send + Sync + 'static,
    {
        self.on_established = Some(Arc::new(callback));
        self
    }

    pub fn on_rejected<F>(mut self, callback: F) -> Self
    where
        F: Fn(String, SocketAddr) + Send + Sync + 'static,
    {
        self.on_rejected = Some(Arc::new(callback));
        self
    }

    pub fn create_hello(&self, request_encryption: bool) -> HelloCommand {
        let ephemeral_public = if request_encryption && self.session_keys.is_some() {
            Some(vec![0u8; 32])
        } else {
            None
        };

        HelloCommand::new(
            self.local_node.clone(),
            ephemeral_public,
            request_encryption,
        )
    }

    pub fn create_welcome(
        &self,
        accepted: bool,
        ephemeral_public: Option<Vec<u8>>,
    ) -> WelcomeCommand {
        WelcomeCommand::new(self.local_node.clone(), accepted, ephemeral_public)
    }

    pub fn create_ack(&self, session_key_id: Option<Vec<u8>>) -> AckCommand {
        AckCommand::accepted(session_key_id)
    }

    pub fn create_reject(&self, reason: &str) -> RejectCommand {
        RejectCommand::new(reason)
    }

    pub async fn handle_server_side(
        &self,
        ctx: Arc<Mutex<Context>>,
        peer_addr: SocketAddr,
    ) -> Result<Option<Node>> {
        {
            let mut guard = ctx.lock().await;
            let reader = guard
                .reader
                .as_mut()
                .ok_or_else(|| anyhow::anyhow!("no reader"))?;
            let mut length_buf = [0u8; 4];
            reader.read_exact(&mut length_buf).await?;
            let len = u32::from_le_bytes(length_buf) as usize;
            if len > MAX_HANDSHAKE_SIZE {
                return Err(anyhow::anyhow!("handshake message too large"));
            }
            let mut data = vec![0u8; len];
            reader.read_exact(&mut data).await?;

            let id = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);

            match CommandId::from_u32(id) {
                Some(CommandId::Hello) => {
                    let hello = HelloCommand::decode(&data).map_err(|e| anyhow::anyhow!(e))?;
                    if !hello.is_valid() {
                        let reject = self.create_reject("version mismatch");
                        self.send_frame(ctx.clone(), reject.encode()).await?;
                        return Err(anyhow::anyhow!("version mismatch"));
                    }

                    if let Some(callback) = &self.on_established {
                        callback(hello.node.clone(), peer_addr);
                    }

                    let ephemeral_public =
                        if hello.request_encryption && self.session_keys.is_some() {
                            Some(vec![0u8; 32])
                        } else {
                            None
                        };

                    let welcome = self.create_welcome(true, ephemeral_public);
                    self.send_frame(ctx.clone(), welcome.encode()).await?;

                    return Ok(Some(hello.node));
                }
                Some(CommandId::Reject) => {
                    let reject = RejectCommand::decode(&data).map_err(|e| anyhow::anyhow!(e))?;
                    if let Some(callback) = &self.on_rejected {
                        callback(reject.reason.clone(), peer_addr);
                    }
                    return Err(anyhow::anyhow!("rejected: {}", reject.reason));
                }
                _ => {
                    return Err(anyhow::anyhow!("expected Hello"));
                }
            }
        }
    }

    async fn send_frame(&self, ctx: Arc<Mutex<Context>>, data: Vec<u8>) -> Result<()> {
        let mut guard = ctx.lock().await;
        let writer = guard
            .writer
            .as_mut()
            .ok_or_else(|| anyhow::anyhow!("no writer"))?;
        writer.write_all(&(data.len() as u32).to_le_bytes()).await?;
        writer.write_all(&data).await?;
        Ok(())
    }

    #[cfg(test)]
    pub async fn send_frame_test(ctx: Arc<Mutex<Context>>, data: Vec<u8>) -> Result<()> {
        Self::send_frame_internal(&ctx, data).await
    }

    #[cfg(test)]
    async fn send_frame_internal(ctx: &Arc<Mutex<Context>>, data: Vec<u8>) -> Result<()> {
        let mut guard = ctx.lock().await;
        let writer = guard
            .writer
            .as_mut()
            .ok_or_else(|| anyhow::anyhow!("no writer"))?;
        writer.write_all(&(data.len() as u32).to_le_bytes()).await?;
        writer.write_all(&data).await?;
        Ok(())
    }

    pub async fn handshake_as_client(
        &self,
        peer_addr: SocketAddr,
        request_encryption: bool,
    ) -> Result<Node> {
        let socket = tokio::net::TcpStream::connect(peer_addr).await?;
        let mut socket = socket;

        let hello = self.create_hello(request_encryption);
        let data = hello.encode();

        socket.write_all(&(data.len() as u32).to_le_bytes()).await?;
        socket.write_all(&data).await?;
        socket.flush().await?;

        let mut length_buf = [0u8; 4];
        socket.read_exact(&mut length_buf).await?;

        let len = u32::from_le_bytes(length_buf) as usize;
        let mut data = vec![0u8; len];
        socket.read_exact(&mut data).await?;

        let id = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);

        match CommandId::from_u32(id) {
            Some(CommandId::Welcome) => {
                let welcome = WelcomeCommand::decode(&data).map_err(|e| anyhow::anyhow!(e))?;
                if !welcome.accepted {
                    return Err(anyhow::anyhow!("connection rejected"));
                }

                if let Some(callback) = &self.on_established {
                    callback(welcome.node.clone(), peer_addr);
                }

                Ok(welcome.node)
            }
            Some(CommandId::Reject) => {
                let reject = RejectCommand::decode(&data).map_err(|e| anyhow::anyhow!(e))?;
                Err(anyhow::anyhow!("rejected: {}", reject.reason))
            }
            _ => Err(anyhow::anyhow!("unexpected message")),
        }
    }
}
