use std::io::SeekFrom;

use tokio::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt};

use crate::chacha::{ChaChaConfigv1, ChaChaHeaderv1};

pub struct ChaChaWriterAsync<W: AsyncWrite + AsyncSeek + Unpin + Send> {
	inner: W,
	header: ChaChaHeaderv1,

	encryption_key: [u8; 32],
	buffer: Vec<u8>,
	plaintext_bytes_written: u64,
}

impl<W: AsyncWrite + AsyncSeek + Unpin + Send> ChaChaWriterAsync<W> {
	pub async fn new(mut inner: W, encryption_key: [u8; 32]) -> Result<Self, std::io::Error> {
		let header = ChaChaHeaderv1 {
			config: ChaChaConfigv1::default(),
			plaintext_size: 0,
		};
		inner.write_all(&serialize_header(header)?).await?;

		Ok(Self {
			inner,
			header,
			encryption_key,
			buffer: Vec::new(),
			plaintext_bytes_written: 0,
		})
	}

	pub async fn write(&mut self, buf: &[u8]) -> Result<(), std::io::Error> {
		self.buffer.extend_from_slice(buf);
		self.plaintext_bytes_written += buf.len() as u64;

		let chunk_size = self.header.config.chunk_size as usize;
		while self.buffer.len() >= chunk_size {
			let encrypted = encrypt_chunk(&self.encryption_key, &self.buffer[..chunk_size])?;
			self.inner.write_all(&encrypted).await?;
			self.buffer.drain(..chunk_size);
		}

		Ok(())
	}

	/// Encrypt and write any buffered plaintext, patch the header with the
	/// final `plaintext_size`, then return the inner writer.
	pub async fn finish(mut self) -> Result<W, std::io::Error> {
		if !self.buffer.is_empty() {
			let encrypted = encrypt_chunk(&self.encryption_key, &self.buffer)?;
			self.inner.write_all(&encrypted).await?;
		}

		self.inner.seek(SeekFrom::Start(0)).await?;
		let header_bytes = serialize_header(ChaChaHeaderv1 {
			config: self.header.config,
			plaintext_size: self.plaintext_bytes_written,
		})?;
		self.inner.write_all(&header_bytes).await?;

		Ok(self.inner)
	}
}

fn encrypt_chunk(key: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>, std::io::Error> {
	use chacha20poly1305::{
		XChaCha20Poly1305,
		aead::{Aead, AeadCore, KeyInit, OsRng},
	};

	let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);
	let cipher = XChaCha20Poly1305::new(chacha20poly1305::Key::from_slice(key));
	let ciphertext = cipher
		.encrypt(&nonce, plaintext)
		.map_err(|_| std::io::Error::other("encryption failed"))?;

	let mut output = Vec::with_capacity(nonce.len() + ciphertext.len());
	output.extend_from_slice(&nonce);
	output.extend_from_slice(&ciphertext);
	Ok(output)
}

fn serialize_header(header: ChaChaHeaderv1) -> Result<Vec<u8>, std::io::Error> {
	use binrw::BinWriterExt;
	use std::io::Cursor;

	let mut buf = Cursor::new(Vec::new());
	buf.write_le(&header).map_err(std::io::Error::other)?;
	Ok(buf.into_inner())
}