use std::sync::Arc;

use rocket::{
    response::stream::{Event, EventStream},
    serde::json::Json,
    time::OffsetDateTime,
};
use rocket_db_pools::Connection;
use rocket_dyn_templates::{Template, context};
use serde::{Deserialize, Serialize};
use sqlx::prelude::FromRow;
use tokio::sync::broadcast;

use crate::{auth::Session, db::DbConn, llm::LlmWorker};

/// ---------- shared broadcaster ----------
pub struct ChatBroadcaster {
    sender: broadcast::Sender<ChatMsg>,
}

impl ChatBroadcaster {
    pub fn new(buffer_size: usize) -> Self {
        let (sender, _rx) = broadcast::channel::<ChatMsg>(buffer_size);
        Self { sender }
    }

    pub async fn publish(&self, msg: ChatMsg) {
        let _ = self.sender.send(msg);
    }

    pub fn subscribe(&self) -> broadcast::Receiver<ChatMsg> {
        self.sender.subscribe()
    }
}

/// ---------- Rocket routes ----------
#[derive(Debug, Serialize, Deserialize, Clone, FromRow)]
pub struct ChatMsg {
    pub display_name: Option<String>,
    pub user_id: usize,
    pub text: String,
    pub timestamp: usize,
}

#[post("/chat", format = "json", data = "<msg>")]
pub async fn post_message(
    mut msg: Json<ChatMsg>,
    chat: &rocket::State<Arc<ChatBroadcaster>>,
    mut db: Connection<DbConn>,
    session: Session,
) -> Result<(), String> {
    const CHANNEL_ID: i32 = 1;
    const LMSTUDIO_URI: &'static str = "http://127.0.0.1:1234/v1/chat/completions";

    let chat = chat.inner().clone();

    msg.user_id = session.user_id;
    chat.publish(msg.clone().into_inner()).await;

    sqlx::query!(
        "INSERT INTO messages (channel_id, user_id, content) VALUES ($1, $2, $3)",
        CHANNEL_ID,
        msg.user_id as i32,
        msg.text
    )
    .execute(&mut **db)
    .await
    .map_err(|_| "Failed".to_string())?;

    // get response
    tokio::spawn(async move {
        let response = LlmWorker::new(LMSTUDIO_URI.to_string()).query(&msg).await;

        if let Ok(message) = response {
            chat.publish(message.clone()).await;

            sqlx::query!(
                "INSERT INTO messages (channel_id, user_id, content, created_at) VALUES ($1, $2, $3, $4)",
                CHANNEL_ID,
                message.user_id as i32,
                message.text,
                OffsetDateTime::from_unix_timestamp(message.timestamp as i64).unwrap()
            )
            .execute(&mut **db)
            .await
            .map_err(|_| "Failed".to_string())
            .unwrap();
        }
    });

    Ok(())
}

#[get("/messages")]
pub async fn get_messages(mut db: Connection<DbConn>, _session: Session) -> Json<Vec<ChatMsg>> {
    Json(
        sqlx::query!(
            "SELECT u.username, u.display_name, u.id, m.content, m.created_at FROM messages m JOIN users u ON m.user_id = u.id ORDER BY m.created_at DESC LIMIT 100"
        )
        .fetch_all(&mut **db)
        .await
        .unwrap_or_else(|_| Vec::new())
        .into_iter()
        .rev()
        .map(|msg| ChatMsg {
            display_name: Some(msg.display_name.unwrap_or(msg.username)),
            user_id: msg.id as usize,
            text: msg.content,
            timestamp: msg.created_at.unwrap().unix_timestamp() as usize,
        })
        .collect(),
    )
    // Json(vec![])
}

#[get("/events")]
pub async fn event_stream(
    chat: &rocket::State<Arc<ChatBroadcaster>>,
    db: Connection<DbConn>,
    ag: Session,
) -> EventStream![] {
    let mut rx = chat.subscribe();

    EventStream! {
        // Initialize the stream with the last 100 messages
        for msg in get_messages(db, ag).await.0 {
            yield Event::json(&msg);
        }

        loop {
            match rx.recv().await {
                Ok(msg) => yield Event::json(&msg),
                Err(broadcast::error::RecvError::Lagged(_)) => {
                    yield Event::comment("RecvError::Lagged");
                }
                Err(broadcast::error::RecvError::Closed) => break,
            }
        }
    }
}

#[get("/")]
pub async fn chat_page(session: Session) -> Template {
    Template::render("chat", context!(user_id: session.user_id))
}
