use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};

use anyhow::{Context, Result};
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio::time::{sleep, Duration};

use frontend_protocol::DekuBytes;
use super::api::{
    ChatCompletionRequest,
    MessageInChat,
    Request,
    RequestPayload,
    Response,
    ResponsePayload,
};
use super::talking::{SecretStreamSocket, wrap_connection_socket};

#[derive(Debug, Clone)]
pub struct MessagePiecePayload(pub String);

#[derive(Debug, Clone)]
pub enum MessagePiece {
    Piece(MessagePiecePayload), End, Cancelled, DedicatedServerDisconnected
}

#[derive(Debug)]
struct PendingChatCompletionRecord {
    response_tx: mpsc::UnboundedSender<MessagePiece>,
}

#[derive(Debug)]
enum RequestWorkItem {
    ChatCompletion {
        request: Request,
        response_tx: mpsc::UnboundedSender<MessagePiece>,
    },
    ChatCompletionCancellation {
        request_id: u64,
    },
}

#[derive(Debug)]
pub struct DedicatedAiServerConnection {
    address: String,
    port: u16,
    secret: String,
    request_tx: mpsc::UnboundedSender<RequestWorkItem>,
    next_request_id: AtomicU64,
}

pub fn connect_to_dedicated_ai_server(address: String, port: u16, secret: String)
    -> (DedicatedAiServerConnection, JoinHandle<()>) {
    let (request_tx, request_rx) = mpsc::unbounded_channel();
    let worker_address = address.clone();
    let worker_secret = secret.clone();

    let join_hand = tokio::spawn(async move {
        connection_worker_loop(worker_address, port, worker_secret, request_rx).await;
    });

    let conn_bridge = DedicatedAiServerConnection {
        address,
        port,
        secret,
        request_tx,
        next_request_id: AtomicU64::new(1),
    };
    (conn_bridge, join_hand)
}

impl DedicatedAiServerConnection {
    pub fn send_chat_completion(
        &self,
        messages: Vec<MessageInChat>,
    ) -> Result<(mpsc::UnboundedReceiver<MessagePiece>, u64)> {
        let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
        let request = Request {
            request_id,
            payload: RequestPayload::ChatCompletion(ChatCompletionRequest::new(messages)),
        };
        let (response_tx, response_rx) = mpsc::unbounded_channel();
        self.request_tx
            .send(RequestWorkItem::ChatCompletion { request, response_tx })
            .map_err(|err| anyhow::anyhow!("failed to enqueue request: {err}"))?;
        Ok((response_rx, request_id))
    }

    pub fn send_chat_completion_cancellation(&self, request_id: u64) -> Result<()> {
        self.request_tx
            .send(RequestWorkItem::ChatCompletionCancellation { request_id })
            .map_err(|err| anyhow::anyhow!("failed to enqueue cancellation: {err}"))?;
        Ok(())
    }
}

async fn connection_worker_loop(
    address: String,
    port: u16,
    secret: String,
    mut request_rx: mpsc::UnboundedReceiver<RequestWorkItem>,
) {
    // todo : fix a lot of errors
    loop {
        let socket = match TcpStream::connect((address.as_str(), port)).await {
            Ok(socket) => socket,
            Err(err) => {
                eprintln!(
                    "[dedicated-ai] failed to connect to {}:{}: {err}",
                    address, port
                );
                sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        let mut transport = match wrap_connection_socket(socket, &secret).await {
            Ok(transport) => transport,
            Err(err) => {
                eprintln!("[dedicated-ai] failed to wrap connection: {err:#}");
                sleep(Duration::from_secs(1)).await;
                continue;
            }
        };

        let mut pending: HashMap<u64, PendingChatCompletionRecord> = HashMap::new();
        let result = run_connected_loop(&mut transport, &mut request_rx, &mut pending).await;
        if let Err(err) = result {
            eprintln!("[dedicated-ai] connection error: {err:#}");
        }

        cancel_all_pending(&mut pending);
        let _ = transport.close(true).await;
        sleep(Duration::from_secs(1)).await;
    }
}

async fn run_connected_loop(
    transport: &mut SecretStreamSocket,
    request_rx: &mut mpsc::UnboundedReceiver<RequestWorkItem>,
    pending: &mut HashMap<u64, PendingChatCompletionRecord>,
) -> Result<()> {
    loop {
        tokio::select! {
            frame = transport.recv_frame() => {
                let frame = frame.context("failed to read frame")?;
                let frame = match frame {
                    Some(bytes) => bytes,
                    None => return Err(anyhow::anyhow!("connection closed by peer")),
                };
                handle_response_frame(frame, pending)?;
            }
            maybe_item = request_rx.recv() => {
                let item = match maybe_item {
                    Some(item) => item,
                    None => return Err(anyhow::anyhow!("request channel closed")), // idk why this would happen
                };
                match item {
                    RequestWorkItem::ChatCompletion { request, response_tx } => {
                        let request_id = request.request_id;
                        pending.insert(
                            request_id,
                            PendingChatCompletionRecord { response_tx },
                        );

                        let payload = request.to_bytes().context("failed to encode request")?;
                        transport.send_frame(&payload).await.context("failed to send request")?;
                    }
                    RequestWorkItem::ChatCompletionCancellation { request_id } => {
                        let request = Request {
                            request_id,
                            payload: RequestPayload::ChatCompletionCancellation,
                        };
                        let payload = request.to_bytes().context("failed to encode cancellation")?;
                        transport.send_frame(&payload).await.context("failed to send cancellation")?;
                    }
                }
            }
        }
    }
}

fn handle_response_frame(
    frame: Vec<u8>,
    pending: &mut HashMap<u64, PendingChatCompletionRecord>,
) -> Result<()> {
    let response = Response::from_bytes(&frame).context("failed to decode response")?;
    let request_id = response.request_id;
    let record = pending.get(&request_id);
    if record.is_none() {
        eprintln!("[dedicated-ai] response for unknown request_id={request_id}");
        return Ok(());
    }

    match response.payload {
        ResponsePayload::ChatCompletion(payload) => {
            let send_failed = match pending.get(&request_id) {
                Some(record) => record
                    .response_tx
                    .send(MessagePiece::Piece( MessagePiecePayload(payload.piece) ))
                    .is_err(),
                None => false,
            };
            if send_failed {
                pending.remove(&request_id);
            }
        }
        ResponsePayload::ChatCompletionEnd => {
            if let Some(record) = pending.remove(&request_id) {
                let _ = record.response_tx.send(MessagePiece::End);
            }
        }
        ResponsePayload::ChatCompletionCancellation => {
            if let Some(record) = pending.remove(&request_id) {
                let _ = record.response_tx.send(MessagePiece::Cancelled);
            }
        }
    }

    Ok(())
}

fn cancel_all_pending(pending: &mut HashMap<u64, PendingChatCompletionRecord>) {
    for (_request_id, record) in pending.drain() {
        let _ = record.response_tx.send(MessagePiece::DedicatedServerDisconnected);
    }
}
