use anyhow::Result;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::time::Instant;
use tracing::info;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::fs::File;
use std::io::Write;
use std::sync::{Arc, Mutex};
use crate::opcodes;

/// Recording file with automatic delay tracking
struct RecordingFile {
    file: File,
    last_write: Instant,
}

impl RecordingFile {
    fn new(file: File) -> Self {
        Self {
            file,
            last_write: Instant::now(),
        }
    }

    /// Write a frame with automatic preceeding delay opcode if needed
    /// If data is None, only the opcode is written
    fn maybe_save_delay_and_write(&mut self, opcode: u8, data: Option<&[u8]>) -> std::io::Result<()> {
        let elapsed_ms = self.last_write.elapsed().as_millis() as u64;

        if elapsed_ms > 1 {
            let delay_ms = elapsed_ms.min(0xFFFFFF_u64);
            let header = opcodes::write_header(opcodes::DELAY, delay_ms as usize);
            self.file.write_all(&header)?;
        }

        let len = data.map(|d| d.len()).unwrap_or(0);
        let header = opcodes::write_header(opcode, len);
        self.file.write_all(&header)?;

        if let Some(d) = data {
            self.file.write_all(d)?;
        }

        self.last_write = Instant::now();
        Ok(())
    }
}

pub async fn run(listen: SocketAddr, connect: SocketAddr, file: PathBuf) -> Result<()> {
    let listener = TcpListener::bind(&listen).await?;
    info!("Proxy listening on {}", listen);

    let (client_conn, _) = listener.accept().await?;
    info!("Accepted client connection");

    let mut file = File::create(&file)?;
    
    let header = opcodes::write_header(opcodes::NOOP, 8);
    file.write_all(&header)?;
    file.write_all(b"tcptape\x00")?;
    
    let recording_file = RecordingFile::new(file);
    let recording_file = Arc::new(Mutex::new(recording_file));

    let server_conn = TcpStream::connect(&connect).await?;
    info!("Connected to server at {}", connect);

    let (client_read, client_write) = client_conn.into_split();
    let (server_read, server_write) = server_conn.into_split();

    let file_clone = Arc::clone(&recording_file);
    let client_to_server = tokio::spawn(async move {
        let mut client_read = client_read;
        let mut server_write = server_write;
        let mut buf = [0; 4096];
        loop {
            match client_read.read(&mut buf).await {
                Ok(0) => {
                    if let Ok(mut rf) = file_clone.lock() {
                        let _ = rf.maybe_save_delay_and_write(opcodes::CLIENT_TO_SERVER_EOF, None);
                    }
                    info!("Client-to-server direction closed");
                    break;
                }
                Ok(n) => {
                    if let Ok(mut rf) = file_clone.lock() {
                        let _ = rf.maybe_save_delay_and_write(opcodes::CLIENT_TO_SERVER, Some(&buf[..n]));
                    }
                    if server_write.write_all(&buf[..n]).await.is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
    });

    let server_to_client = tokio::spawn(async move {
        let mut server_read = server_read;
        let mut client_write = client_write;
        let mut buf = [0; 4096];
        loop {
            match server_read.read(&mut buf).await {
                Ok(0) => {
                    if let Ok(mut rf) = recording_file.lock() {
                        let _ = rf.maybe_save_delay_and_write(opcodes::SERVER_TO_CLIENT_EOF, None);
                    }
                    info!("Server-to-client direction closed");
                    break;
                }
                Ok(n) => {
                    if let Ok(mut rf) = recording_file.lock() {
                        let _ = rf.maybe_save_delay_and_write(opcodes::SERVER_TO_CLIENT, Some(&buf[..n]));
                    }
                    if client_write.write_all(&buf[..n]).await.is_err() {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
    });

    let _ = tokio::join!(client_to_server, server_to_client);
    info!("Proxy connection closed");
    Ok(())
}
