use axum::Extension;
use axum::Json;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use uuid::Uuid;

use crate::engine::{GenerateEvent, GenerateRequest};
use crate::AppState;

#[derive(Deserialize)]
pub struct ChatRequest {
    #[serde(default)]
    pub model: Option<String>,
    pub messages: Vec<Message>,
    #[serde(default = "default_max_tokens")]
    pub max_tokens: usize,
}

#[derive(Deserialize)]
pub struct Message {
    pub role: String,
    pub content: String,
}

fn default_max_tokens() -> usize { 256 }

#[derive(Serialize)]
pub struct ModelsResponse {
    object: &'static str,
    data: Vec<ModelInfo>,
}

#[derive(Serialize)]
pub struct ModelInfo {
    id: String,
    object: &'static str,
    owned_by: &'static str,
}

pub async fn health() -> &'static str { "ok" }

pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
    Json(ModelsResponse {
        object: "list",
        data: vec![ModelInfo {
            id: state.model_name.clone(),
            object: "model",
            owned_by: "xserv",
        }],
    })
}

pub async fn chat_completions(
    Extension(state): Extension<Arc<AppState>>,
    Json(req): Json<ChatRequest>,
) -> Json<serde_json::Value> {
    let id = format!("chatcmpl-{}", Uuid::new_v4());
    let model_name = state.model_name.clone();
    let created = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs();

    // Prepare prompt tokens (MutexGuard scoped)
    let prompt = build_prompt(&req.messages);
    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);

    // Create channel and submit request (MutexGuard scoped)
    let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
    let gen_req = GenerateRequest {
        prompt_tokens,
        max_tokens: req.max_tokens,
        sender: tx,
    };
    state.engine_sender.lock().unwrap().send(gen_req).expect("engine channel closed");

    // Now await — no MutexGuards held here
    let mut content = String::new();
    let mut finish_reason = "length".to_string();
    while let Some(event) = rx.recv().await {
        match event {
            GenerateEvent::Token { text, .. } => content.push_str(&text),
            GenerateEvent::Done { finish_reason: fr } => { finish_reason = fr; break; }
        }
    }

    Json(serde_json::json!({
        "id": id,
        "object": "chat.completion",
        "created": created,
        "model": model_name,
        "choices": [{
            "index": 0,
            "message": { "role": "assistant", "content": content },
            "finish_reason": finish_reason,
        }],
        "usage": {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0
        }
    }))
}

fn build_prompt(messages: &[Message]) -> String {
    let mut prompt = String::new();
    for msg in messages {
        match msg.role.as_str() {
            "system" => { prompt.push_str(&msg.content); prompt.push('\n'); }
            "user" | "assistant" => { prompt.push_str(&msg.content); }
            _ => {}
        }
    }
    prompt
}