use std::{
    io::{self, Read, Write},
    net::{SocketAddr, TcpStream},
    os::fd::{AsRawFd, RawFd},
    time::Duration,
};

use crate::{
    crypt::{
        client::{self, Negotiation},
        config::ClientConfig,
        signature::{SignatureMethod, VerificationMethod},
        symmetric::{
            self,
            data::{InboundEncryptedMessage, OutboundPlainMessage},
            SymmetricCipher,
        },
    },
    proto::{Decode, DecodeError, Decoder, Encode, EncodeError, Encoder},
};

pub const MESSAGE_SIZE_MAX: usize = 4096;
const FRAMING_BUFFER_SIZE: usize = 8192;
const HEADER_SIZE: usize = size_of::<u64>() + size_of::<u16>();

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("I/O: {0}")]
    Io(#[from] io::Error),
    #[error("Encode error: {0}")]
    Encode(#[from] EncodeError),
    #[error("Decode error: {0}")]
    Decode(#[from] DecodeError),
    #[error("Cipher error: {0}")]
    Symmetric(#[from] symmetric::Error),
    #[error("Receive buffer too small")]
    BufferTooSmall,
}

pub struct ClientSocket {
    pub(crate) stream: TcpStream,
    pub(crate) remote: SocketAddr,
    pub(crate) buffer: [u8; MESSAGE_SIZE_MAX],
    pub(crate) recv_buf: FramingBuffer,

    // TODO: only use those for signature verification during negotiation?
    #[allow(unused)]
    pub(crate) signer: SignatureMethod,
    #[allow(unused)]
    pub(crate) verifier: VerificationMethod,
    pub(crate) symmetric: SymmetricCipher,

    pub(crate) tx_seq: u64,
}

pub enum Message<T> {
    Data(T),
    Incomplete,
    Closed,
}

pub(crate) struct FramingBuffer {
    buffer: [u8; FRAMING_BUFFER_SIZE],
    len: usize,
}

impl FramingBuffer {
    pub fn new() -> Self {
        Self {
            buffer: [0; FRAMING_BUFFER_SIZE],
            len: 0,
        }
    }

    pub fn get_mut(&mut self) -> &mut [u8] {
        &mut self.buffer[self.len..]
    }

    pub fn advance(&mut self, len: usize) {
        self.len += len;
        assert!(self.len <= self.buffer.len());
    }

    pub fn pop(&mut self, buffer: &mut [u8]) -> Option<(usize, u64)> {
        if self.len < HEADER_SIZE {
            return None;
        }
        let len = u16::from_be_bytes([self.buffer[0], self.buffer[1]]) as usize;
        if self.len < HEADER_SIZE + len {
            return None;
        }
        let mut seq = [0; 8];
        seq.copy_from_slice(&self.buffer[2..HEADER_SIZE]);
        let seq = u64::from_be_bytes(seq);
        // TODO check dst len
        buffer[..len].copy_from_slice(&self.buffer[HEADER_SIZE..len + HEADER_SIZE]);
        self.buffer.copy_within(len + HEADER_SIZE..self.len, 0);
        self.len -= len + HEADER_SIZE;
        Some((len, seq))
    }
}

impl ClientSocket {
    pub(crate) fn from_parts(
        stream: TcpStream,
        remote: SocketAddr,
        signer: SignatureMethod,
        verifier: VerificationMethod,
        symmetric: SymmetricCipher,
    ) -> Self {
        ClientSocket {
            stream,
            remote,
            signer,
            verifier,
            symmetric,
            buffer: [0; MESSAGE_SIZE_MAX],
            recv_buf: FramingBuffer::new(),
            tx_seq: 0x1234567887654321,
        }
    }

    pub fn connect(address: SocketAddr, config: ClientConfig) -> Result<Self, client::Error> {
        let stream = TcpStream::connect(address)?;
        stream.set_read_timeout(Some(Duration::from_secs(1)))?;
        stream.set_write_timeout(Some(Duration::from_secs(1)))?;
        Negotiation::new(stream, address, config).perform()
    }

    pub fn remote_address(&self) -> SocketAddr {
        self.remote
    }

    pub fn write_all<E: Encode>(&mut self, message: &E) -> Result<(), Error> {
        let mut buf = [0; MESSAGE_SIZE_MAX - 256];

        let mut encoder = Encoder::new(&mut buf[HEADER_SIZE..]);
        message.encode(&mut encoder)?;
        let (payload, _) = encoder.split_mut();

        let seq = self.tx_seq;
        self.tx_seq = self.tx_seq.wrapping_add(payload.len() as u64);

        let outbound_plain = OutboundPlainMessage::new(payload);
        let outbound_encrypted = self.symmetric.encrypter.encrypt(seq, outbound_plain)?;

        let total_len = outbound_encrypted.payload.len();
        let len: u16 = total_len.try_into().unwrap();

        buf[0..2].copy_from_slice(&len.to_be_bytes());
        buf[2..HEADER_SIZE].copy_from_slice(&seq.to_be_bytes());
        buf[HEADER_SIZE..HEADER_SIZE + outbound_encrypted.payload.len()]
            .copy_from_slice(&outbound_encrypted.payload[..]);

        self.stream.write_all(&buf[..total_len + HEADER_SIZE])?;

        Ok(())
    }

    pub fn poll_read<'de, D: Decode<'de>>(
        &mut self,
        buffer: &'de mut [u8],
    ) -> Result<Message<D>, Error> {
        if self.poll()? == 0 {
            return Ok(Message::Closed);
        }
        match self.read(buffer)? {
            Some(message) => Ok(Message::Data(message)),
            None => Ok(Message::Incomplete),
        }
    }

    pub fn poll(&mut self) -> Result<usize, Error> {
        let dst = self.recv_buf.get_mut();
        if dst.is_empty() {
            todo!()
        }
        let len = self.stream.read(dst)?;
        self.recv_buf.advance(len);
        Ok(len)
    }

    pub fn read<'de, D: Decode<'de>>(&mut self, buffer: &'de mut [u8]) -> Result<Option<D>, Error> {
        if let Some((len, seq)) = self.recv_buf.pop(&mut self.buffer) {
            if len > buffer.len() {
                return Err(Error::BufferTooSmall);
            }
            buffer[..len].copy_from_slice(&self.buffer[..len]);
            let inbound_encrypted = InboundEncryptedMessage::from(&mut buffer[..len]);
            let inbound_plain = self.symmetric.decrypter.decrypt(seq, inbound_encrypted)?;
            let mut decoder = Decoder::new(inbound_plain.payload);
            let message = D::decode(&mut decoder)?;

            Ok(Some(message))
        } else {
            // Buffer doesn't contain a full message yet
            Ok(None)
        }
    }
}

impl AsRawFd for ClientSocket {
    fn as_raw_fd(&self) -> RawFd {
        self.stream.as_raw_fd()
    }
}
