//! `/api/generate` — 浏览器需求 → LLM gateway → SSE 流式回吐生成的 HTML。
//!
//! 上游是 OpenAI 兼容的 gateway（gemma-4-31b-it）。我们开 `stream: true`，逐块解析
//! 上游的 `data: {...}` SSE 帧，把 delta.content 里的文本片段重新封装成我们自己的
//! 干净 SSE 事件发给前端：
//!   - `data: {"t": "<片段>"}`   —— 一段生成文本
//!   - `event: done` `data: {}`  —— 正常结束
//!   - `event: error` `data: {"message": "..."}` —— 出错
//!
//! 前端把所有 `t` 拼起来就是完整 HTML，结束后剥掉可能的 ```html fence 丢进 iframe。

use std::convert::Infallible;
use std::sync::Arc;

use axum::{
    extract::State,
    response::sse::{Event, KeepAlive, Sse},
    Json,
};
use futures::StreamExt;
use serde::Deserialize;
use serde_json::{json, Value};
use tokio_stream::wrappers::ReceiverStream;

#[derive(Clone)]
pub struct Config {
    pub gateway: String,   // http://3.135.65.204:8848/v1
    pub llm_token: String, // Bearer for LLM gateway
    pub llm_model: String, // gemma-4-31b-it
    pub max_tokens: u32,
}

impl Config {
    pub fn from_env() -> Self {
        Self {
            gateway: env_or("LLM_GATEWAY", "http://3.135.65.204:8848/v1"),
            llm_token: env_or("LLM_API_TOKEN", ""),
            llm_model: env_or("LLM_MODEL", "gemma-4-31b-it"),
            max_tokens: env_or("LLM_MAX_TOKENS", "8192")
                .parse()
                .unwrap_or(8192),
        }
    }
}

fn env_or(key: &str, fallback: &str) -> String {
    std::env::var(key).unwrap_or_else(|_| fallback.to_string())
}

#[derive(Deserialize)]
pub struct GenerateRequest {
    /// 用户这次说的需求 / 修改要求。
    pub prompt: String,
    /// 可选：当前程序的完整代码。带上它 = 在现有程序上迭代修改，否则 = 从零生成。
    #[serde(default)]
    pub current_code: Option<String>,
}

/// 系统提示词 —— 逼 gemma 吐一个自包含、纯原生 WebGL、能在沙箱 iframe 离线跑的 HTML。
pub fn system_prompt() -> String {
    "你是一个 WebGL 创意编程助手，专门为一个小朋友生成**可交互的图形小程序**。\n\
     用户用自然语言描述想要什么，你直接输出一个完整、独立、能立刻运行的 HTML 文档。\n\
     \n\
     硬性要求：\n\
     1. 只输出**一个完整的 HTML 文档**，从 `<!doctype html>` 开始到 `</html>` 结束。\n\
     2. 全部代码内联：CSS 写在 `<style>`，JS 写在 `<script>`。**绝对不要**引用任何外部\n\
        资源 / CDN / 库（不准用 three.js、不准 import、不准 <script src=...>）——必须用\n\
        **原生 WebGL / WebGL2** 手写。程序要能在没有网络的沙箱 iframe 里直接跑。\n\
     3. 用一个铺满窗口的 `<canvas>`，背景默认深色，body/html 设 margin:0、overflow:hidden。\n\
        监听 resize 让 canvas 跟着窗口和 devicePixelRatio 变。\n\
     4. **一定要可交互**：鼠标 / 触摸 / 键盘至少支持一种，让小朋友能玩起来\n\
        （拖动旋转、点击生成、按键变色…随主题而定）。视觉要鲜艳、好玩、有动画。\n\
     5. 代码要稳：着色器编译、program link 失败时不要整页崩，能力所及地容错。\n\
     6. 适当写注释（中文或英文都行），因为小朋友会看代码学习。\n\
     \n\
     只输出 HTML 代码本身，不要任何解释性文字，不要 markdown 代码围栏。"
        .to_string()
}

pub async fn handle(
    State(cfg): State<Arc<Config>>,
    Json(req): Json<GenerateRequest>,
) -> Sse<ReceiverStream<Result<Event, Infallible>>> {
    let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);

    tokio::spawn(async move {
        if let Err(msg) = stream_generation(&cfg, &req, &tx).await {
            let _ = tx
                .send(Ok(Event::default()
                    .event("error")
                    .data(json!({ "message": msg }).to_string())))
                .await;
        }
    });

    Sse::new(ReceiverStream::new(rx)).keep_alive(KeepAlive::default())
}

/// 拼用户 prompt：从零生成 vs 在现有代码上改。
fn user_content(req: &GenerateRequest) -> String {
    match req.current_code.as_deref().map(str::trim) {
        Some(code) if !code.is_empty() => format!(
            "这是当前程序的完整代码：\n\n```html\n{code}\n```\n\n\
             请在它的基础上做如下修改，并重新输出**完整**的 HTML 文档：\n{}",
            req.prompt.trim(),
            code = code,
        ),
        _ => req.prompt.trim().to_string(),
    }
}

async fn stream_generation(
    cfg: &Config,
    req: &GenerateRequest,
    tx: &tokio::sync::mpsc::Sender<Result<Event, Infallible>>,
) -> Result<(), String> {
    let body = json!({
        "model": cfg.llm_model,
        "messages": [
            { "role": "system", "content": system_prompt() },
            { "role": "user", "content": user_content(req) },
        ],
        "stream": true,
        "temperature": 0.8,
        "max_tokens": cfg.max_tokens,
    });

    let endpoint = format!("{}/chat/completions", cfg.gateway.trim_end_matches('/'));
    let client = reqwest::Client::new();
    let resp = client
        .post(&endpoint)
        .bearer_auth(&cfg.llm_token)
        .json(&body)
        .send()
        .await
        .map_err(|e| format!("调用 LLM 失败：{e}"))?;

    let status = resp.status();
    if !status.is_success() {
        let txt = resp.text().await.unwrap_or_default();
        return Err(format!("LLM 返回 {status}：{}", txt.chars().take(300).collect::<String>()));
    }

    // 逐块读上游 SSE，按行解析 `data:` 帧。
    let mut stream = resp.bytes_stream();
    let mut buf = String::new();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|e| format!("读 LLM 流出错：{e}"))?;
        buf.push_str(&String::from_utf8_lossy(&chunk));

        // 处理所有完整行（以 \n 结尾的）。
        while let Some(nl) = buf.find('\n') {
            let line: String = buf.drain(..=nl).collect();
            let line = line.trim_end_matches(['\r', '\n']);
            let payload = match line.strip_prefix("data:") {
                Some(p) => p.trim_start(),
                None => continue, // 空行 / 注释 / 其它字段，忽略
            };
            if payload == "[DONE]" {
                let _ = tx.send(Ok(Event::default().event("done").data("{}"))).await;
                return Ok(());
            }
            if let Some(text) = parse_delta(payload) {
                if !text.is_empty() {
                    let ev = Event::default().data(json!({ "t": text }).to_string());
                    if tx.send(Ok(ev)).await.is_err() {
                        return Ok(()); // 客户端断开，安静退出
                    }
                }
            }
        }
    }

    // 上游没发 [DONE] 就结束了，也算正常收尾。
    let _ = tx.send(Ok(Event::default().event("done").data("{}"))).await;
    Ok(())
}

/// 从一帧 OpenAI streaming JSON 里抠出 `choices[0].delta.content`。
/// 解析不出来就返回 None（容错：坏帧不该让整条流挂掉）。
fn parse_delta(payload: &str) -> Option<String> {
    let v: Value = serde_json::from_str(payload).ok()?;
    v.pointer("/choices/0/delta/content")
        .and_then(Value::as_str)
        .map(str::to_string)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn system_prompt_forbids_external_libs() {
        let p = system_prompt();
        assert!(p.contains("WebGL"));
        assert!(p.contains("CDN"));
        assert!(p.contains("<!doctype html>"));
    }

    #[test]
    fn parse_delta_extracts_content() {
        let frame = r#"{"choices":[{"delta":{"content":"<canvas>"}}]}"#;
        assert_eq!(parse_delta(frame).as_deref(), Some("<canvas>"));
    }

    #[test]
    fn parse_delta_handles_role_only_frame() {
        // 首帧通常只有 role，没有 content
        let frame = r#"{"choices":[{"delta":{"role":"assistant"}}]}"#;
        assert_eq!(parse_delta(frame), None);
    }

    #[test]
    fn parse_delta_handles_garbage() {
        assert_eq!(parse_delta("not json"), None);
        assert_eq!(parse_delta("[DONE]"), None);
    }

    #[test]
    fn user_content_plain_when_no_code() {
        let req = GenerateRequest { prompt: "  一个会转的立方体  ".into(), current_code: None };
        assert_eq!(user_content(&req), "一个会转的立方体");
    }

    #[test]
    fn user_content_embeds_existing_code_for_iteration() {
        let req = GenerateRequest {
            prompt: "让它转快点".into(),
            current_code: Some("<!doctype html><html></html>".into()),
        };
        let c = user_content(&req);
        assert!(c.contains("当前程序的完整代码"));
        assert!(c.contains("<!doctype html>"));
        assert!(c.contains("让它转快点"));
    }

    #[test]
    fn user_content_ignores_blank_code() {
        let req = GenerateRequest { prompt: "画个圆".into(), current_code: Some("   ".into()) };
        assert_eq!(user_content(&req), "画个圆");
    }
}
