use axum::Extension;
use axum::Json;
use axum::http::StatusCode;
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use serde::{Deserialize, Serialize};
use std::convert::Infallible;
use std::path::Path;
use std::sync::Arc;
use tokio_stream::StreamExt;
use tokio_stream::wrappers::ReceiverStream;
use uuid::Uuid;

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

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

#[derive(Deserialize, Serialize, Clone)]
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,
}

// ---------------------------------------------------------------------------
// Chat Template: Jinja2 rendering via minijinja
// ---------------------------------------------------------------------------

pub struct ChatTemplate {
    source: String,
    model_type: String,
}

impl ChatTemplate {
    pub fn load(model_dir: &Path, model_type: &str) -> Self {
        // 1. Try standalone chat_template.jinja file
        let jinja_path = model_dir.join("chat_template.jinja");
        if jinja_path.exists() {
            let source = std::fs::read_to_string(&jinja_path)
                .unwrap_or_else(|e| panic!("failed to read {}: {e}", jinja_path.display()));
            eprintln!("[chat-template] loaded from {}", jinja_path.display());
            return Self {
                source,
                model_type: model_type.to_string(),
            };
        }

        // 2. Try tokenizer_config.json → chat_template field
        let tok_cfg_path = model_dir.join("tokenizer_config.json");
        if tok_cfg_path.exists() {
            if let Ok(data) = std::fs::read_to_string(&tok_cfg_path) {
                if let Ok(v) = serde_json::from_str::<serde_json::Value>(&data) {
                    if let Some(ct) = v.get("chat_template").and_then(|v| v.as_str()) {
                        eprintln!("[chat-template] loaded from tokenizer_config.json");
                        return Self {
                            source: ct.to_string(),
                            model_type: model_type.to_string(),
                        };
                    }
                }
            }
        }

        // 3. No template found — use empty source, will fall back to hardcoded
        eprintln!("[chat-template] no Jinja template found, using hardcoded fallback");
        Self {
            source: String::new(),
            model_type: model_type.to_string(),
        }
    }

    pub fn render(&self, messages: &[Message]) -> String {
        if self.source.is_empty() {
            return build_prompt_hardcoded(messages, &self.model_type);
        }

        match self.render_jinja(messages) {
            Ok(prompt) => prompt,
            Err(e) => {
                eprintln!("[chat-template] Jinja render error: {e}, falling back to hardcoded");
                build_prompt_hardcoded(messages, &self.model_type)
            }
        }
    }

    fn render_jinja(&self, messages: &[Message]) -> Result<String, minijinja::Error> {
        let mut env = minijinja::Environment::new();

        // Register custom functions the template may call.
        env.add_function("strftime_now", strftime_now);
        env.add_function("raise_exception", raise_exception);

        // Python str methods used by harmony/gpt-oss templates.
        env.add_filter("startswith", |s: String, prefix: String| -> bool {
            s.starts_with(&prefix)
        });

        env.add_template("chat", &self.source)?;
        let tmpl = env.get_template("chat")?;

        let ctx = minijinja::context! {
            messages => minijinja::Value::from_serialize(messages),
            add_generation_prompt => true,
            bos_token => "",
            eos_token => "",
        };

        tmpl.render(ctx)
    }
}

fn strftime_now(fmt: String) -> String {
    use std::time::SystemTime;
    let now = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .as_secs();
    // Only support %Y-%m-%d (the only format used by known templates)
    let days = now / 86400;
    let (y, m, d) = days_to_ymd(days);
    fmt.replace("%Y", &format!("{y:04}"))
        .replace("%m", &format!("{m:02}"))
        .replace("%d", &format!("{d:02}"))
}

fn days_to_ymd(days_since_epoch: u64) -> (u32, u32, u32) {
    // Civil calendar from days since 1970-01-01 (Rata Die algorithm)
    let z = days_since_epoch as i64 + 719468;
    let era = (if z >= 0 { z } else { z - 146096 }) / 146097;
    let doe = (z - era * 146097) as u32;
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
    let y = yoe as i64 + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if m <= 2 { y + 1 } else { y };
    (y as u32, m, d)
}

fn raise_exception(msg: String) -> Result<String, minijinja::Error> {
    Err(minijinja::Error::new(
        minijinja::ErrorKind::InvalidOperation,
        msg,
    ))
}

// ---------------------------------------------------------------------------
// Hardcoded fallback templates (for models without a Jinja template)
// ---------------------------------------------------------------------------

fn build_prompt_hardcoded(messages: &[Message], model_type: &str) -> String {
    if model_type == "gpt_oss" {
        return build_prompt_gpt_oss(messages);
    }
    // Default: Qwen3 ChatML format
    let mut prompt = String::new();
    for msg in messages {
        match msg.role.as_str() {
            "system" | "user" | "assistant" => {
                prompt.push_str("<|im_start|>");
                prompt.push_str(&msg.role);
                prompt.push('\n');
                prompt.push_str(&msg.content);
                prompt.push_str("<|im_end|>\n");
            }
            _ => {}
        }
    }
    prompt.push_str("<|im_start|>assistant\n");
    prompt.push_str("<think>\n\n</think>\n\n");
    prompt
}

fn build_prompt_gpt_oss(messages: &[Message]) -> String {
    let mut prompt = String::new();
    // Canonical harmony system message (mirrors the model's chat_template.jinja
    // build_system_message macro). A hand-rolled substitute puts gpt-oss out of
    // distribution and destabilizes channel selection. This hardcoded builder is
    // only a fallback for gpt-oss models that ship no Jinja template; the
    // gpt-oss-20b release does ship one, so the template path is normally used.
    prompt.push_str("<|start|>system<|message|>");
    prompt.push_str("You are ChatGPT, a large language model trained by OpenAI.\n");
    prompt.push_str("Knowledge cutoff: 2024-06\n");
    prompt.push_str(&format!(
        "Current date: {}\n\n",
        strftime_now("%Y-%m-%d".to_string())
    ));
    prompt.push_str("Reasoning: low\n\n");
    prompt.push_str("# Valid channels: analysis, commentary, final. Channel must be included for every message.");
    prompt.push_str("<|end|>");
    let dev_instructions: String = messages
        .iter()
        .filter(|m| m.role == "system")
        .map(|m| m.content.as_str())
        .collect::<Vec<_>>()
        .join("\n\n");
    if !dev_instructions.is_empty() {
        prompt.push_str("<|start|>developer<|message|># Instructions\n\n");
        prompt.push_str(&dev_instructions);
        prompt.push_str("<|end|>");
    }
    for msg in messages {
        match msg.role.as_str() {
            "user" => {
                prompt.push_str("<|start|>user<|message|>");
                prompt.push_str(&msg.content);
                prompt.push_str("<|end|>");
            }
            "assistant" => {
                prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
                prompt.push_str(&msg.content);
                prompt.push_str("<|end|>");
            }
            _ => {}
        }
    }
    prompt.push_str("<|start|>assistant<|channel|>final<|message|>");
    prompt
}

// ---------------------------------------------------------------------------
// HTTP handlers
// ---------------------------------------------------------------------------

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>,
) -> Response {
    if req.stream == Some(true) {
        chat_stream(state, req)
    } else {
        chat_non_stream(state, req).await
    }
}

async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
    let id = format!("chatcmpl-{}", Uuid::new_v4());
    let model_name = state.model_name.clone();
    let created = unix_timestamp();

    if let Some(response) = validate_request(&req, &model_name) {
        return response;
    }

    let prompt = state.chat_template.render(&req.messages);
    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
    let prompt_token_count = prompt_tokens.len();

    let max_seq_len = state.max_seq_len;
    if prompt_token_count >= max_seq_len {
        return bad_request(format!(
            "prompt is {} tokens, exceeds max_seq_len {}",
            prompt_token_count, max_seq_len
        ));
    }
    let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count);

    let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
    let gen_req = GenerateRequest {
        prompt_tokens,
        max_tokens,
        sampling: sampling_params(&req),
        sender: tx,
    };
    if let Err(resp) = submit_to_engine(&state, gen_req) {
        return resp;
    }

    let mut content = String::new();
    let mut completion_token_count: usize = 0;
    let mut finish_reason = "length".to_string();
    while let Some(event) = rx.recv().await {
        match event {
            GenerateEvent::Token { text, .. } => {
                completion_token_count += 1;
                content.push_str(&text);
            }
            GenerateEvent::Done { finish_reason: fr } => {
                finish_reason = fr;
                break;
            }
        }
    }

    let fr_value = match normalize_finish_reason(&finish_reason) {
        Some(s) => serde_json::Value::String(s.to_string()),
        None => serde_json::Value::Null,
    };
    Json(serde_json::json!({
        "id": id,
        "object": "chat.completion",
        "created": created,
        "model": model_name,
        "choices": [{
            "index": 0,
            "message": { "role": "assistant", "content": content },
            "finish_reason": fr_value,
        }],
        "usage": {
            "prompt_tokens": prompt_token_count,
            "completion_tokens": completion_token_count,
            "total_tokens": prompt_token_count + completion_token_count
        }
    }))
    .into_response()
}

fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
    let id = format!("chatcmpl-{}", Uuid::new_v4());
    let model_name = state.model_name.clone();
    let created = unix_timestamp();

    if let Some(response) = validate_request(&req, &model_name) {
        return response;
    }

    let prompt = state.chat_template.render(&req.messages);
    let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);

    let max_seq_len = state.max_seq_len;
    if prompt_tokens.len() >= max_seq_len {
        return bad_request(format!(
            "prompt is {} tokens, exceeds max_seq_len {}",
            prompt_tokens.len(),
            max_seq_len
        ));
    }
    let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());

    let (engine_tx, engine_rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
    let gen_req = GenerateRequest {
        prompt_tokens,
        max_tokens,
        sampling: sampling_params(&req),
        sender: engine_tx,
    };
    if let Err(resp) = submit_to_engine(&state, gen_req) {
        return resp;
    }

    // SSE event channel: engine events -> SSE events
    let (sse_tx, sse_rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);

    tokio::spawn(async move {
        let mut engine_stream = ReceiverStream::new(engine_rx);
        let mut first = true;

        while let Some(event) = engine_stream.next().await {
            match event {
                GenerateEvent::Token { text, .. } => {
                    if first {
                        // First chunk: role announcement
                        let chunk =
                            make_chunk(&id, &model_name, created, None, Some("assistant"), None);
                        let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                        first = false;
                    }
                    let chunk = make_chunk(&id, &model_name, created, Some(&text), None, None);
                    if sse_tx.send(Ok(Event::default().data(chunk))).await.is_err() {
                        return; // client disconnected
                    }
                }
                GenerateEvent::Done { finish_reason } => {
                    if first {
                        // Edge case: Done arrived with no tokens
                        let chunk =
                            make_chunk(&id, &model_name, created, None, Some("assistant"), None);
                        let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                    }
                    // Only "stop" and "length" are OpenAI-standard values. Internal
                    // codes like "error" (client-stalled from tp/pp engine) map to
                    // null so SDK clients see a clean stream close.
                    let fr = normalize_finish_reason(&finish_reason);
                    let chunk = make_chunk(&id, &model_name, created, None, None, fr);
                    let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
                    let _ = sse_tx
                        .send(Ok(Event::default().data("[DONE]".to_string())))
                        .await;
                    return;
                }
            }
        }
    });

    Sse::new(ReceiverStream::new(sse_rx))
        .keep_alive(KeepAlive::default())
        .into_response()
}

fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
    if let Some(model) = &req.model {
        if model != model_name {
            return Some(bad_request(format!(
                "model '{model}' is not loaded; available model is '{model_name}'"
            )));
        }
    }

    if req.max_tokens == 0 {
        return Some(bad_request("max_tokens must be greater than 0"));
    }

    if let Some(t) = req.temperature {
        if !t.is_finite() || t < 0.0 {
            return Some(bad_request("temperature must be a finite value >= 0"));
        }
    }
    if let Some(p) = req.top_p {
        if !p.is_finite() || !(0.0..=1.0).contains(&p) {
            return Some(bad_request("top_p must be in [0, 1]"));
        }
    }
    if let Some(k) = req.top_k {
        if k > 1_000_000 {
            return Some(bad_request("top_k must be <= 1_000_000"));
        }
    }

    None
}

/// Hand a request to the engine thread. Poison-tolerant (recovers the lock if a
/// prior handler panicked) and returns a clean 503 instead of panicking when the
/// engine thread is gone, so one dead engine doesn't cascade into every request.
fn submit_to_engine(state: &AppState, req: GenerateRequest) -> Result<(), Response> {
    let sender = state
        .engine_sender
        .lock()
        .unwrap_or_else(|e| e.into_inner());
    sender.try_send(req).map_err(|err| match err {
        std::sync::mpsc::TrySendError::Full(_) => {
            service_unavailable("inference engine is busy, retry later")
        }
        std::sync::mpsc::TrySendError::Disconnected(_) => {
            service_unavailable("inference engine is not available")
        }
    })
}

fn service_unavailable(message: impl Into<String>) -> Response {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        Json(serde_json::json!({
            "error": { "message": message.into(), "type": "server_error" }
        })),
    )
        .into_response()
}

fn bad_request(message: impl Into<String>) -> Response {
    (
        StatusCode::BAD_REQUEST,
        Json(serde_json::json!({
            "error": {
                "message": message.into(),
                "type": "invalid_request_error"
            }
        })),
    )
        .into_response()
}

fn make_chunk(
    id: &str,
    model: &str,
    created: u64,
    content: Option<&str>,
    role: Option<&str>,
    finish_reason: Option<&str>,
) -> String {
    let mut delta = serde_json::Map::new();
    if let Some(r) = role {
        delta.insert("role".into(), serde_json::Value::String(r.into()));
        // Role chunk also includes empty content per OpenAI spec
        delta.insert("content".into(), serde_json::Value::String(String::new()));
    }
    if let Some(c) = content {
        delta.insert("content".into(), serde_json::Value::String(c.into()));
    }

    let fr = match finish_reason {
        Some(r) => serde_json::Value::String(r.into()),
        None => serde_json::Value::Null,
    };

    serde_json::json!({
        "id": id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [{
            "index": 0,
            "delta": delta,
            "finish_reason": fr,
        }]
    })
    .to_string()
}

fn unix_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

fn sampling_params(req: &ChatRequest) -> SamplingParams {
    SamplingParams {
        temperature: req.temperature.unwrap_or(0.0),
        top_k: req.top_k.unwrap_or(0),
        top_p: req.top_p.unwrap_or(1.0),
    }
}

/// Map engine finish_reason strings to OpenAI-standard values. Any engine-internal
/// code (e.g. "error" from tp/pp client-stall) collapses to None so SDK clients see
/// a clean null instead of an unknown value.
fn normalize_finish_reason(fr: &str) -> Option<&'static str> {
    match fr {
        "stop" => Some("stop"),
        "length" => Some("length"),
        _ => None,
    }
}