use std::{
    collections::HashMap,
    io::{self, Read, Write},
    net::{SocketAddr, TcpListener, TcpStream},
    os::fd::{AsRawFd, RawFd},
};

use cross::io::Poll;
use x25519_dalek::{EphemeralSecret, PublicKey};

use crate::{
    crypt::{
        sig_algo_name, signature::fingerprint_sha256, stream::ClientSocket, symmetric::data::Iv,
        ServerHello, ServerNegotiationMessage, V1_CLIENT_RANDOM_LEN,
    },
    proto::{Decode, DecodeError, Decoder, Encode, EncodeError, Encoder},
};

use super::{
    config::ServerConfig,
    signature::{SignatureMethod, VerificationMethod},
    symmetric::{self, SymmetricCipher},
    ClientNegotiationMessage,
};

pub struct ServerSocket {
    poll: Poll,
    listener: TcpListener,
    pending: HashMap<RawFd, PendingClient>,
    config: ServerConfig,
}

struct PendingClient {
    stream: TcpStream,
    address: SocketAddr,
    state: NegotiationState,

    signer: Option<SignatureMethod>,
    verifier: Option<VerificationMethod>,
    symmetric: Option<SymmetricCipher>,

    buffer: [u8; 512],
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("I/O error: {0}")]
    Io(#[from] io::Error),
    #[error("Encode error: {0}")]
    Encode(#[from] EncodeError),
    #[error("Decode error: {0}")]
    Decode(#[from] DecodeError),
    #[error("Unexpected client message")]
    UnexpectedClientMessage,

    #[error("Symmetric cipher error: {0}")]
    Symmetric(#[from] symmetric::Error),

    #[error("Unsupported signature algorithm")]
    NoSignerForAlgorithm,
    #[error("Client public key rejected")]
    ClientKeyRejected,
}

enum NegotiationOutcome {
    Pending,
    Accepted,
}

#[derive(Clone, Copy)]
enum NegotiationState {
    None,
    Hello,
    StartKex(u8, u8, [u8; V1_CLIENT_RANDOM_LEN]),
    DHExchange,
}

impl ServerSocket {
    pub fn bind(addr: SocketAddr, config: ServerConfig) -> Result<Self, Error> {
        let mut poll = Poll::new()?;
        let listener = TcpListener::bind(addr)?;
        poll.add(&listener)?;
        Ok(Self {
            listener,
            poll,
            config,
            pending: HashMap::new(),
        })
    }

    fn accept(&mut self) -> Result<(), Error> {
        let (stream, address) = self.listener.accept()?;
        log::info!("{address}: connect");
        self.poll.add(&stream)?;
        self.pending.insert(
            stream.as_raw_fd(),
            PendingClient {
                stream,
                address,
                state: NegotiationState::None,
                buffer: [0; 512],

                signer: None,
                verifier: None,
                symmetric: None,
            },
        );
        Ok(())
    }

    pub fn poll(&mut self) -> Result<Option<ClientSocket>, Error> {
        let fd = self.poll.wait(None)?.unwrap();

        match fd {
            fd if fd == self.listener.as_raw_fd() => {
                self.accept()?;
                Ok(None)
            }
            fd if let Some(pending) = self.pending.get_mut(&fd) => {
                let address = pending.address;
                match pending.update(&mut self.config) {
                    Ok(NegotiationOutcome::Pending) => Ok(None),
                    Ok(NegotiationOutcome::Accepted) => {
                        self.poll.remove(&fd)?;
                        let client = self.pending.remove(&fd).unwrap();
                        Ok(Some(ClientSocket::from_parts(
                            client.stream,
                            address,
                            client.signer.unwrap(),
                            client.verifier.unwrap(),
                            client.symmetric.unwrap(),
                        )))
                    }
                    Err(error) => {
                        log::error!("{address}: {error}");
                        self.poll.remove(&fd)?;
                        self.pending.remove(&fd);
                        Ok(None)
                    }
                }
            }
            _ => unreachable!(),
        }
    }
}

impl PendingClient {
    fn update(&mut self, config: &mut ServerConfig) -> Result<NegotiationOutcome, Error> {
        let address = self.address;
        let state = self.state;
        let message = self.recv()?;

        match (message, state) {
            // None -> Hello
            (ClientNegotiationMessage::Hello { protocol }, NegotiationState::None) => {
                log::debug!("{address}: Hello v{protocol}");

                let kex_algos = (config.offer_kex_algorithms)();
                let sig_algos = (config.offer_sig_algorhtms)();
                let symmetric_ciphersuites = (config.offer_ciphersuites)();

                let hello = ServerHello {
                    kex_algos,
                    sig_algos,
                    symmetric_ciphersuites,
                };
                self.send(&ServerNegotiationMessage::Hello(hello))?;
                self.state = NegotiationState::Hello;
                Ok(NegotiationOutcome::Pending)
            }
            // Hello -> StartKex
            (
                ClientNegotiationMessage::StartKex {
                    kex_algorithm,
                    sig_algorithm,
                    ciphersuite,
                    client_random,
                    key_data,
                },
                NegotiationState::Hello,
            ) => {
                log::debug!("{address}: StartKex {{ ... }}");
                log::debug!("{address}: client random: {client_random:02x?}");

                // TODO entropy test on client random

                let sig_algorithm_name = sig_algo_name(sig_algorithm).unwrap_or("???");
                let their_fingerprint = fingerprint_sha256(sig_algorithm_name, key_data);
                let client_random = *client_random;
                log::info!("{address}: their fingerprint {their_fingerprint}");

                let verifier = match config
                    .signature_keystore
                    .accept_public_key(sig_algorithm, key_data)
                {
                    Some(signer) => signer,
                    None => return Err(Error::ClientKeyRejected),
                };

                let signer = match config
                    .signature_keystore
                    .signer_for_algorithm(sig_algorithm)
                {
                    Some(signer) => signer,
                    None => return Err(Error::NoSignerForAlgorithm),
                };
                let public_key = signer.verifying_key_bytes();

                self.send(&ServerNegotiationMessage::StartKex {
                    sig_algorithm,
                    key_data: &public_key,
                })?;
                self.signer = Some(signer);
                self.verifier = Some(verifier);
                self.state = NegotiationState::StartKex(ciphersuite, kex_algorithm, client_random);
                Ok(NegotiationOutcome::Pending)
            }
            (
                ClientNegotiationMessage::DHPublicKey(key_data),
                NegotiationState::StartKex(ciphersuite, _kex_algorithm, client_random),
            ) => {
                log::debug!("{address}: DHPublicKey {{ ... }}");
                let (ciphersuite_algo, ciphersuite_name) = symmetric::ciphersuite(ciphersuite)
                    .ok_or(Error::Symmetric(symmetric::Error::InvalidCiphersuite(
                        ciphersuite,
                    )))?;
                let key_size = ciphersuite_algo.key_size();

                log::debug!("{address}: selected ciphersuite {ciphersuite_name}");

                let mut rng = rand::thread_rng();
                let mut their_public = [0; 32];
                their_public.copy_from_slice(key_data);
                let their_public = PublicKey::from(their_public);
                let ephemeral = EphemeralSecret::random_from_rng(&mut rng);
                let public = PublicKey::from(&ephemeral);
                let shared = ephemeral.diffie_hellman(&their_public);

                let iv = Iv::from(client_random);
                let symmetric = SymmetricCipher::from_ciphersuite(
                    ciphersuite_algo,
                    &shared.as_bytes()[..key_size],
                    iv,
                )?;

                self.send(&ServerNegotiationMessage::DHPublicKey(public.as_bytes()))?;
                self.state = NegotiationState::DHExchange;
                self.symmetric = Some(symmetric);
                Ok(NegotiationOutcome::Pending)
            }
            (ClientNegotiationMessage::Agreed, NegotiationState::DHExchange) => {
                log::debug!("{address}: Agreed");
                self.send(&ServerNegotiationMessage::Agreed)?;
                Ok(NegotiationOutcome::Accepted)
            }
            (_, _) => Err(Error::UnexpectedClientMessage),
        }
    }

    fn recv(&mut self) -> Result<ClientNegotiationMessage, Error> {
        self.stream
            .read_exact(&mut self.buffer[..size_of::<u16>()])?;
        let len = u16::from_le_bytes([self.buffer[0], self.buffer[1]]) as usize;
        if len >= self.buffer.len() {
            todo!()
        }
        self.stream.read_exact(&mut self.buffer[..len])?;
        let mut decoder = Decoder::new(&self.buffer[..len]);
        let message = ClientNegotiationMessage::decode(&mut decoder)?;
        Ok(message)
    }

    fn send(&mut self, message: &ServerNegotiationMessage) -> Result<(), Error> {
        let mut encoder = Encoder::new(&mut self.buffer[2..]);
        message.encode(&mut encoder)?;
        let payload = encoder.get();
        let len = payload.len() as u16;
        self.buffer[..2].copy_from_slice(&len.to_le_bytes());
        self.stream.write_all(&self.buffer[..len as usize + 2])?;
        Ok(())
    }
}

impl AsRawFd for ServerSocket {
    fn as_raw_fd(&self) -> RawFd {
        self.poll.as_raw_fd()
    }
}
