use std::convert::TryFrom;

use anyhow::{Context, Result, bail};
use blake2b_simd::Params;
use dryoc::classic::crypto_secretstream_xchacha20poly1305::{
    Header,
    Key,
    State,
    crypto_secretstream_xchacha20poly1305_init_pull,
    crypto_secretstream_xchacha20poly1305_init_push,
    crypto_secretstream_xchacha20poly1305_pull,
    crypto_secretstream_xchacha20poly1305_push,
};
use dryoc::constants::{
    CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_ABYTES,
    CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_HEADERBYTES,
    CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_KEYBYTES,
    CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_TAG_FINAL,
    CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_TAG_MESSAGE,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::TcpStream;


pub type FrameCallback<'a> = Box<dyn FnMut(Vec<u8>) -> Result<()> + Send + 'a>;

pub const MAX_FRAME_SIZE: usize = 64 * 1024 * 1024;
pub const TRANSPORT_CHUNK_SIZE: usize = 4096;
pub const CHUNK_USED_BYTES: usize = 2;
pub const CHUNK_DATA_SIZE: usize = TRANSPORT_CHUNK_SIZE - CHUNK_USED_BYTES;
pub const CIPHERTEXT_CHUNK_SIZE: usize =
    TRANSPORT_CHUNK_SIZE + CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_ABYTES;
pub const FRAME_LENGTH_BYTES: usize = 4;

const KEY_DERIVATION_PERSON: &[u8; 16] = b"vv-secretstream1";


#[derive(Debug)]
pub struct ProtocolError(pub String);


impl std::fmt::Display for ProtocolError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}


impl std::error::Error for ProtocolError {}


pub fn derive_key(shared_secret: &str) -> Result<Key> {
    let hash = Params::new()
        .hash_length(CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_KEYBYTES)
        .personal(KEY_DERIVATION_PERSON)
        .hash(shared_secret.as_bytes());

    Key::try_from(hash.as_bytes()).context("failed to convert derived bytes into secretstream key")
}


pub async fn wrap_connection_socket(
    connection_socket: TcpStream,
    shared_secret: &str,
) -> Result<SecretStreamSocket> {
    let key = derive_key(shared_secret)?;
    let (read_half, write_half) = connection_socket.into_split();
    let mut reader = BufReader::new(read_half);
    let mut writer = BufWriter::new(write_half);
    let mut push_state = State::new();
    let mut pull_state = State::new();
    let mut header = Header::default();

    crypto_secretstream_xchacha20poly1305_init_push(&mut push_state, &mut header, &key);
    writer
        .write_all(header.as_ref())
        .await
        .context("failed to write local secretstream header")?;
    writer
        .flush()
        .await
        .context("failed to flush local secretstream header")?;

    let mut peer_header_bytes = [0u8; CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_HEADERBYTES];
    reader
        .read_exact(&mut peer_header_bytes)
        .await
        .context("failed to read peer secretstream header")?;
    let peer_header =
        Header::try_from(peer_header_bytes.as_slice()).context("invalid peer header size")?;
    crypto_secretstream_xchacha20poly1305_init_pull(&mut pull_state, &peer_header, &key);

    Ok(SecretStreamSocket {
        reader,
        writer,
        push_state,
        pull_state,
        buffer: Vec::new(),
        seen_final: false,
        sent_final: false,
    })
}


pub struct SecretStreamSocket {
    reader: BufReader<OwnedReadHalf>,
    writer: BufWriter<OwnedWriteHalf>,
    push_state: State,
    pull_state: State,
    buffer: Vec<u8>,
    seen_final: bool,
    sent_final: bool,
}


impl SecretStreamSocket {
    pub async fn send_frame(&mut self, frame: &[u8]) -> Result<()> {
        if self.sent_final {
            bail!("cannot send after final chunk was sent");
        }
        if frame.len() > MAX_FRAME_SIZE {
            bail!("frame size {} exceeds limit {}", frame.len(), MAX_FRAME_SIZE);
        }

        let frame_len = u32::try_from(frame.len()).context("frame length does not fit in u32")?;
        let mut plaintext = Vec::with_capacity(FRAME_LENGTH_BYTES + frame.len());
        plaintext.extend_from_slice(&frame_len.to_be_bytes());
        plaintext.extend_from_slice(frame);

        let mut start = 0usize;
        while start < plaintext.len() {
            let end = (start + CHUNK_DATA_SIZE).min(plaintext.len());
            self.send_transport_chunk(
                &plaintext[start..end],
                CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_TAG_MESSAGE,
            )
            .await?;
            start = end;
        }

        self.writer
            .flush()
            .await
            .context("failed to flush encrypted frame")?;
        Ok(())
    }

    pub async fn recv_frame(&mut self) -> Result<Option<Vec<u8>>> {
        self.fill_plain_buffer(1).await?;
        if self.buffer.is_empty() && self.seen_final {
            return Ok(None);
        }

        let length_bytes = read_exact(self, FRAME_LENGTH_BYTES).await?;
        let frame_len = u32::from_be_bytes(length_bytes.as_slice().try_into()?) as usize;
        if frame_len > MAX_FRAME_SIZE {
            return Err(ProtocolError(format!(
                "frame size {} exceeds limit {}",
                frame_len, MAX_FRAME_SIZE
            ))
            .into());
        }

        let frame = read_exact(self, frame_len).await?;
        Ok(Some(frame))
    }

    pub async fn run_receiving_loop(&mut self, mut on_frame: FrameCallback<'_>) -> Result<()> {
        while let Some(frame) = self.recv_frame().await? {
            on_frame(frame)?;
        }
        Ok(())
    }

    pub async fn read(&mut self, size: usize) -> Result<Vec<u8>> {
        while self.buffer.len() < size && !self.seen_final {
            self.read_transport_chunk().await?;
        }

        let available = size.min(self.buffer.len());
        let data: Vec<u8> = self.buffer.drain(..available).collect();
        Ok(data)
    }

    pub async fn close(&mut self, send_final: bool) -> Result<()> {
        if send_final && !self.sent_final {
            let _ = self
                .send_transport_chunk(&[], CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_TAG_FINAL)
                .await;
            let _ = self.writer.flush().await;
            self.sent_final = true;
        }

        let _ = self.writer.shutdown().await;
        Ok(())
    }

    async fn fill_plain_buffer(&mut self, size: usize) -> Result<()> {
        while self.buffer.len() < size && !self.seen_final {
            self.read_transport_chunk().await?;
        }
        Ok(())
    }

    async fn read_transport_chunk(&mut self) -> Result<()> {
        let mut ciphertext = vec![0u8; CIPHERTEXT_CHUNK_SIZE];
        self.reader
            .read_exact(&mut ciphertext)
            .await
            .context("failed to read ciphertext chunk")?;

        let mut plaintext = vec![0u8; TRANSPORT_CHUNK_SIZE];
        let mut tag = 0u8;
        crypto_secretstream_xchacha20poly1305_pull(
            &mut self.pull_state,
            &mut plaintext,
            &mut tag,
            &ciphertext,
            None,
        )
        .context("failed to decrypt ciphertext chunk")?;

        let used = u16::from_be_bytes([plaintext[0], plaintext[1]]) as usize;
        if used > CHUNK_DATA_SIZE {
            return Err(ProtocolError(format!(
                "chunk declared {} payload bytes, limit is {}",
                used, CHUNK_DATA_SIZE
            ))
            .into());
        }

        self.buffer
            .extend_from_slice(&plaintext[CHUNK_USED_BYTES..CHUNK_USED_BYTES + used]);

        if tag == CRYPTO_SECRETSTREAM_XCHACHA20POLY1305_TAG_FINAL {
            self.seen_final = true;
        }

        Ok(())
    }

    async fn send_transport_chunk(&mut self, payload: &[u8], tag: u8) -> Result<()> {
        if payload.len() > CHUNK_DATA_SIZE {
            bail!(
                "chunk payload size {} exceeds limit {}",
                payload.len(),
                CHUNK_DATA_SIZE
            );
        }

        let mut plaintext = vec![0u8; TRANSPORT_CHUNK_SIZE];
        let used = u16::try_from(payload.len()).context("payload length does not fit in u16")?;
        plaintext[..CHUNK_USED_BYTES].copy_from_slice(&used.to_be_bytes());
        plaintext[CHUNK_USED_BYTES..CHUNK_USED_BYTES + payload.len()].copy_from_slice(payload);

        let mut ciphertext = vec![0u8; CIPHERTEXT_CHUNK_SIZE];
        crypto_secretstream_xchacha20poly1305_push(
            &mut self.push_state,
            &mut ciphertext,
            &plaintext,
            None,
            tag,
        )
        .context("failed to encrypt transport chunk")?;

        self.writer
            .write_all(&ciphertext)
            .await
            .context("failed to write ciphertext chunk")?;
        Ok(())
    }
}


async fn read_exact(stream: &mut SecretStreamSocket, size: usize) -> Result<Vec<u8>> {
    let data = stream.read(size).await?;
    if data.len() != size {
        bail!("connection closed");
    }
    Ok(data)
}
