use std::sync::Arc;

use rocket::{
    Shutdown,
    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::{select, sync::broadcast};

use crate::{
    auth::Session,
    db::{Postgres, Redis},
    llm::LlmWorker,
    messenger,
};

/// ---------- shared broadcaster ----------
pub struct ChatBroadcaster {
    buffer_size: usize,
    senders: std::sync::Mutex<std::collections::HashMap<i32, broadcast::Sender<ChatMsg>>>,
}

impl ChatBroadcaster {
    pub fn new(buffer_size: usize) -> Self {
        Self {
            buffer_size,
            senders: std::sync::Mutex::new(std::collections::HashMap::new()),
        }
    }

    /// Publish a message to the specified channel.
    pub async fn publish(&self, channel_id: i32, msg: ChatMsg) {
        let mut map = self.senders.lock().unwrap();
        let sender = map
            .entry(channel_id)
            .or_insert_with(|| broadcast::channel::<ChatMsg>(self.buffer_size).0);
        let _ = sender.send(msg);
    }

    /// Subscribe to the specified channel.
    pub fn subscribe(&self, channel_id: i32) -> broadcast::Receiver<ChatMsg> {
        let mut map = self.senders.lock().unwrap();
        let sender = map
            .entry(channel_id)
            .or_insert_with(|| broadcast::channel::<ChatMsg>(self.buffer_size).0);
        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/<channel_id>", format = "json", data = "<msg>")]
pub async fn post_message(
    mut msg: Json<ChatMsg>,
    chat: &rocket::State<Arc<ChatBroadcaster>>,
    mut postgres: Connection<Postgres>,
    mut cache: Option<Connection<Redis>>,
    session: Session,
    channel_id: i32,
) -> Result<(), String> {
    let chat = chat.inner().clone();

    let display_name = sqlx::query!(
        "SELECT display_name, username FROM users WHERE id = $1",
        session.user_id as i32
    )
    .fetch_one(&mut **postgres)
    .await
    .map(|row| row.display_name.unwrap_or(row.username))
    .unwrap_or_else(|_| "Unknown".to_string());

    msg.user_id = session.user_id;
    msg.display_name = Some(display_name);
    chat.publish(channel_id, msg.clone().into_inner()).await;

    sqlx::query!(
        "INSERT INTO messages (channel_id, user_id, content, created_at) VALUES ($1, $2, $3, $4)",
        channel_id,
        msg.user_id as i32,
        msg.text,
        OffsetDateTime::from_unix_timestamp_nanos(msg.timestamp as i128 * 1_000_000).unwrap()
    )
    .execute(&mut **postgres)
    .await
    .map_err(|_| "Failed".to_string())?;

    println!("gisfujdeghnjuisdfjngiosdfgjkosdf gnojdfsg nmodfsg");

    if let Some(ref mut cache) = cache {
        messenger::cache::insert(cache, channel_id, &msg)
            .await
            .map_err(|_| "Redis cache failed".to_string())?;
    }

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

        if let Ok(reply) = response {
            chat.publish(channel_id, reply.clone()).await;

            if let Some(ref mut cache) = cache {
                messenger::cache::insert(cache, channel_id, &reply)
                    .await
                    .ok();
            }

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

    Ok(())
}

pub async fn get_messages(
    mut db: Connection<Postgres>,
    mut redis: Connection<Redis>,
    channel_id: i32,
) -> Json<Vec<ChatMsg>> {
    if let Ok(messages) = messenger::cache::get(&mut redis, channel_id).await
        && !messages.is_empty()
    {
        return Json(messages);
    };

    if let Err(x) = messenger::cache::initialise(&mut redis, &mut db, channel_id).await {
        eprintln!("WARN: {x:?}");
    }

    if let Ok(messages) = messenger::cache::get(&mut redis, channel_id).await
        && !messages.is_empty()
    {
        return Json(messages);
    };

    let res = 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
            WHERE m.channel_id = $1
            ORDER BY m.created_at DESC LIMIT 100",
        channel_id
    )
    .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_nanos() / 1_000_000) as usize,
    })
    .collect();
    Json(res)
}

#[get("/events/<channel_id>")]
pub async fn event_stream(
    chat: &rocket::State<Arc<ChatBroadcaster>>,
    postgres: Connection<Postgres>,
    cache: Connection<Redis>,
    _session: Session,
    mut shutdown: Shutdown,
    channel_id: i32,
) -> EventStream![] {
    let mut rx = chat.subscribe(channel_id);

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

        loop {
            select!{
                // exit early on shutdown
                _ = &mut shutdown => break,

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

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