#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Minimalist Qwen-HTML WebUI Server (:8791)
- English-only minimal interface
- Backend options: Llama.cpp API Server or Direct Model Path
- Real-time streaming for Reasoning Trace (<think>) & Rendered HTML (<div class="ai-view">)
"""

import os, re
import json
import time
import argparse
import subprocess
import threading
from pathlib import Path
from typing import List, Dict, Optional, Union, Any
import httpx
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

# --- Self-contained layout: everything resolved relative to this file, env-overridable ---
BASE_DIR = Path(__file__).resolve().parent
STATIC_DIR = str(BASE_DIR / "static")
CONFIG_FILE = str(BASE_DIR / "config.json")
HISTORY_FILE = str(BASE_DIR / "chat_history.jsonl")

PORT = int(os.environ.get("WEBUI_PORT", "8791"))
DEFAULT_API_URL = os.environ.get("API_URL", "http://127.0.0.1:18085")
DEFAULT_MODEL_PATH = os.environ.get("MODEL_PATH", str(BASE_DIR.parent))
DEFAULT_MODEL_ALIAS = os.environ.get("MODEL_ALIAS", "qwen-3.8-27b")

# Direct (llama.cpp GGUF) backend — only used by POST /api/model/load
LLAMA_SERVER_BIN = os.environ.get(
    "LLAMA_SERVER_BIN", "/mnt/sdb/ubuntu/llama.cpp-clef/build/bin/llama-server"
)
LLAMA_SCOPE = os.environ.get("LLAMA_SCOPE", "llama-server-q8")
LLAMA_BACKEND_PORT = int(os.environ.get("LLAMA_BACKEND_PORT", "18085"))

DEFAULT_CONFIG = {
    "backend_mode": "api",  # 'api' = talk to vLLM/OpenAI backend; 'direct' = managed llama.cpp server
    "api_url": DEFAULT_API_URL,
    "model_path": DEFAULT_MODEL_PATH,
    "model_alias": DEFAULT_MODEL_ALIAS,
    "context_size": 81920,
    "gpu_layers": 99,
    "do_sample": True,
    "temperature": 1.0,
    "top_p": 0.95,
    "top_k": 20,
    "min_p": 0.0,
    "presence_penalty": 1.5,
    "repetition_penalty": 1.0,
    "max_tokens": 8192,
    "enable_thinking": True,
    "theme_preset": "default",
    "custom_css": "",
    "system_prompt": "",
}

config_lock = threading.Lock()

def load_config() -> dict:
    cfg = dict(DEFAULT_CONFIG)
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                cfg.update(json.load(f))
        except Exception:
            pass
    # Env (set by run_webui.sh, resolved relative to this folder) wins over a
    # stale config.json so a copied folder never points at the old model.
    for key, env in (("api_url", "API_URL"),
                     ("model_path", "MODEL_PATH"),
                     ("model_alias", "MODEL_ALIAS")):
        if os.environ.get(env):
            cfg[key] = os.environ[env]
    return cfg

def save_config(cfg: dict):
    with config_lock:
        with open(CONFIG_FILE, "w", encoding="utf-8") as f:
            json.dump(cfg, f, indent=2, ensure_ascii=False)

current_config = load_config()

history_lock = threading.Lock()

def log_chat(record: dict):
    with history_lock:
        with open(HISTORY_FILE, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

app = FastAPI(title="Qwen-HTML WebUI")
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

@app.get("/")
async def root():
    return FileResponse(os.path.join(STATIC_DIR, "index.html"))

@app.get("/api/config")
async def get_config_endpoint():
    return current_config

class UpdateConfigRequest(BaseModel):
    backend_mode: Optional[str] = None
    api_url: Optional[str] = None
    model_path: Optional[str] = None
    context_size: Optional[int] = None
    gpu_layers: Optional[int] = None
    do_sample: Optional[bool] = None
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    top_k: Optional[int] = None
    min_p: Optional[float] = None
    presence_penalty: Optional[float] = None
    repetition_penalty: Optional[float] = None
    max_tokens: Optional[int] = None
    enable_thinking: Optional[bool] = None
    theme_preset: Optional[str] = None
    custom_css: Optional[str] = None
    system_prompt: Optional[str] = None

@app.post("/api/config")
async def update_config_endpoint(req: UpdateConfigRequest):
    updates = req.dict(exclude_unset=True)
    current_config.update(updates)
    save_config(current_config)
    return {"status": "ok", "config": current_config}

@app.get("/api/status")
async def status_endpoint():
    api_url = current_config.get("api_url", DEFAULT_API_URL)
    is_alive = False
    model_info = "Unknown"
    latency_ms = None
    
    t0 = time.time()
    try:
        async with httpx.AsyncClient(timeout=2.0) as client:
            resp = await client.get(f"{api_url}/health")
            if resp.status_code == 200:
                is_alive = True
                latency_ms = round((time.time() - t0) * 1000, 1)
                m_resp = await client.get(f"{api_url}/v1/models")
                if m_resp.status_code == 200:
                    models = m_resp.json().get("data", [])
                    if models:
                        model_info = models[0].get("id", DEFAULT_MODEL_ALIAS)
    except Exception:
        pass

    return {
        "status": "ready" if is_alive else "disconnected",
        "backend_mode": current_config.get("backend_mode", "api"),
        "api_url": api_url,
        "model_path": current_config.get("model_path"),
        "model_info": model_info,
        "latency_ms": latency_ms,
        "enable_thinking": current_config.get("enable_thinking", True)
    }

class LoadModelRequest(BaseModel):
    model_path: str
    context_size: Optional[int] = 65536
    gpu_layers: Optional[int] = 99

@app.post("/api/model/load")
async def load_model_endpoint(req: LoadModelRequest):
    model_path = req.model_path.strip()

    # 支持三种写法：
    #   1) 本地路径            /path/to/model.gguf          -> -m
    #   2) HuggingFace 仓库     user/repo[:quant]            -> -hf（国内可 export HF_ENDPOINT=https://hf-mirror.com）
    #   3) 直链                 https://.../model.gguf       -> -mu
    if model_path.startswith(("http://", "https://")):
        source, model_arg = "url", ["-mu", model_path]
    elif re.fullmatch(r"[\w.-]+/[\w.-]+(:\w+)?", model_path) and not os.path.exists(model_path):
        source, model_arg = "hf", ["-hf", model_path]
    elif os.path.exists(model_path):
        source, model_arg = "local", ["-m", model_path]
    else:
        return JSONResponse(status_code=400, content={
            "error": f"Not found: {model_path}（既不是本地文件，也不是 user/repo 形式的 HF 仓库名）"})
    
    ctx = req.context_size or DEFAULT_CONFIG["context_size"]
    ngl = req.gpu_layers if req.gpu_layers is not None else 99

    # Restart managed llama-server scope
    subprocess.run(["systemctl", "--user", "stop", f"{LLAMA_SCOPE}.scope"], check=False)
    time.sleep(1.5)

    start_cmd = [
        "systemd-run", "--user", "--scope", f"--unit={LLAMA_SCOPE}",
        LLAMA_SERVER_BIN,
        "--host", "0.0.0.0",
        "--port", str(LLAMA_BACKEND_PORT),
        *model_arg,
        "-c", str(ctx),
        "-ngl", str(ngl),
        "-fa", "on",
        "--alias", DEFAULT_MODEL_ALIAS,
        "-t", "8"
    ]
    # Vision projector: sibling of the chosen model, else any well-known sibling dir
    mmproj_dirs = [os.path.dirname(model_path)] if source == "local" else []
    for extra in (current_config.get("mmproj_path"), DEFAULT_MODEL_PATH):
        if extra and extra not in mmproj_dirs:
            mmproj_dirs.append(extra if os.path.isdir(extra) else os.path.dirname(extra) or ".")
    for d in mmproj_dirs:
        candidate = os.path.join(d, "mmproj-Qwen3.5-4B-BF16.gguf")
        if os.path.exists(candidate):
            start_cmd.extend(["--mmproj", candidate])
            break

    proc = subprocess.Popen(start_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    
    # Poll for health check
    ready = False
    for _ in range(25):
        time.sleep(0.5)
        try:
            with httpx.Client(timeout=1.0) as client:
                res = client.get(f"http://127.0.0.1:{LLAMA_BACKEND_PORT}/health")
                if res.status_code == 200:
                    ready = True
                    break
        except Exception:
            pass

    if ready:
        current_config["model_path"] = model_path
        current_config["context_size"] = ctx
        current_config["gpu_layers"] = ngl
        current_config["backend_mode"] = "direct"
        save_config(current_config)
        return {"status": "ok", "message": f"Successfully loaded [{source}] {model_path}"}
    else:
        return JSONResponse(status_code=500, content={"error": "Failed to start llama-server with the specified model."})

class ChatMessage(BaseModel):
    role: str
    content: Union[str, List[Any], Dict[str, Any]]

class ChatStreamRequest(BaseModel):
    messages: List[ChatMessage]
    enable_thinking: Optional[bool] = None
    reasoning_effort: Optional[str] = None
    do_sample: Optional[bool] = None
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    top_k: Optional[int] = None
    min_p: Optional[float] = None
    presence_penalty: Optional[float] = None
    repetition_penalty: Optional[float] = None
    max_tokens: Optional[int] = None
    system_prompt: Optional[str] = None

@app.post("/api/chat/stream")
async def chat_stream_endpoint(req: ChatStreamRequest, request: Request):
    client_ip = request.client.host if request.client else "unknown"
    t0 = time.time()
    
    enable_thinking = req.enable_thinking if req.enable_thinking is not None else current_config.get("enable_thinking", True)
    do_sample = req.do_sample if req.do_sample is not None else current_config.get("do_sample", True)
    temp = req.temperature if req.temperature is not None else current_config.get("temperature", 1.0)
    if not do_sample:
        temp = 0.0
    top_p = req.top_p if req.top_p is not None else current_config.get("top_p", 0.95)
    top_k = req.top_k if req.top_k is not None else current_config.get("top_k", 20)
    min_p = req.min_p if req.min_p is not None else current_config.get("min_p", 0.0)
    pres_penalty = req.presence_penalty if req.presence_penalty is not None else current_config.get("presence_penalty", 1.5)
    rep_penalty = req.repetition_penalty if req.repetition_penalty is not None else current_config.get("repetition_penalty", 1.0)
    max_tokens = req.max_tokens if req.max_tokens is not None else current_config.get("max_tokens", 8192)
    sys_prompt = req.system_prompt if req.system_prompt is not None else current_config.get("system_prompt", "")
    effort = req.reasoning_effort if req.reasoning_effort is not None else current_config.get("reasoning_effort")
    if effort in ("low", "medium"):
        effort = "low"
    elif effort in ("high", "xhigh"):
        effort = "high"
    else:
        effort = "low"
    
    messages_payload = []
    if sys_prompt and sys_prompt.strip():
        messages_payload.append({"role": "system", "content": sys_prompt.strip()})
    
    for msg in req.messages:
        messages_payload.append({"role": msg.role, "content": msg.content})

    api_url = current_config.get("api_url", DEFAULT_API_URL)

    # Resolve the model id: use configured alias, but fall back to whatever
    # the backend actually serves (aliases drift when the user edits
    # config.json; a stale name would 404 every chat).
    model_alias = current_config.get("model_alias") or ""
    async with httpx.AsyncClient(timeout=3.0) as probe:
        try:
            r = await probe.get(f"{api_url}/v1/models")
            if r.status_code == 200:
                served = [m.get("id", "") for m in r.json().get("data", [])]
                if served and model_alias not in served:
                    model_alias = served[0]
        except Exception:
            pass
    if not model_alias:
        model_alias = "qwen-3.8-27b"

    payload = {
        "model": model_alias,
        "messages": messages_payload,
        "temperature": temp,
        "top_p": top_p,
        "max_tokens": max_tokens,
        "stream": True,
        "stop": ["<｜end▁of▁sentence｜>", "<|im_end|>", "</div>\n</div>\n</div>\n</div>"],
        "chat_template_kwargs": {"enable_thinking": enable_thinking,
                                 **({"reasoning_effort": effort} if effort else {})},
    }
    if top_k is not None and top_k > 0:
        payload["top_k"] = top_k
    if min_p is not None and min_p > 0:
        payload["min_p"] = min_p
    if pres_penalty is not None:
        payload["presence_penalty"] = pres_penalty
    if rep_penalty is not None and rep_penalty > 0:
        payload["repetition_penalty"] = rep_penalty

    async def event_generator():
        thought_chunks = []
        content_chunks = []
        token_count = 0
        
        in_thinking_mode = False
        
        try:
            async with httpx.AsyncClient(timeout=180.0) as client:
                async with client.stream(
                    "POST",
                    f"{api_url}/v1/chat/completions",
                    json=payload,
                    headers={"Content-Type": "application/json"}
                ) as response:
                    if response.status_code != 200:
                        err_text = await response.aread()
                        yield f"data: {json.dumps({'type': 'error', 'error': err_text.decode('utf-8')}, ensure_ascii=False)}\n\n"
                        return

                    async for line in response.aiter_lines():
                        if not line or not line.startswith("data: "):
                            continue
                        data_str = line[6:].strip()
                        if data_str == "[DONE]":
                            break
                        try:
                            chunk_json = json.loads(data_str)
                            choice = chunk_json.get("choices", [{}])[0]
                            delta = choice.get("delta", {})
                            
                            # Reasoning trace
                            reasoning = delta.get("reasoning_content") or delta.get("reasoning")
                            if reasoning:
                                thought_chunks.append(reasoning)
                                token_count += 1
                                out_data = json.dumps({"type": "thought", "content": reasoning}, ensure_ascii=False)
                                yield f"data: {out_data}\n\n"
                            
                            # Main content (HTML)
                            if "content" in delta and delta["content"]:
                                text = delta["content"]
                                token_count += 1
                                
                                # Intercept and route raw <think> tags if present in content
                                cur = text
                                while cur:
                                    if not in_thinking_mode:
                                        if "<think>" in cur:
                                            before, cur = cur.split("<think>", 1)
                                            if before:
                                                content_chunks.append(before)
                                                yield f"data: {json.dumps({'type': 'content', 'content': before}, ensure_ascii=False)}\n\n"
                                            in_thinking_mode = True
                                        else:
                                            content_chunks.append(cur)
                                            yield f"data: {json.dumps({'type': 'content', 'content': cur}, ensure_ascii=False)}\n\n"
                                            break
                                    else:
                                        if "</think>" in cur:
                                            thought_part, cur = cur.split("</think>", 1)
                                            if thought_part:
                                                thought_chunks.append(thought_part)
                                                yield f"data: {json.dumps({'type': 'thought', 'content': thought_part}, ensure_ascii=False)}\n\n"
                                            in_thinking_mode = False
                                        else:
                                            thought_chunks.append(cur)
                                            yield f"data: {json.dumps({'type': 'thought', 'content': cur}, ensure_ascii=False)}\n\n"
                                            break
                        except Exception:
                            continue

            elapsed = round(time.time() - t0, 2)
            tok_per_sec = round(token_count / elapsed, 1) if elapsed > 0 else 0
            
            full_thought = "".join(thought_chunks).strip()
            full_content = "".join(content_chunks).strip()
            
            done_payload = json.dumps({
                "type": "done",
                "tokens": token_count,
                "elapsed": elapsed,
                "tok_per_sec": tok_per_sec,
            }, ensure_ascii=False)
            yield f"data: {done_payload}\n\n"
            
            # Log chat
            log_chat({
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
                "client_ip": client_ip,
                "messages": [m.dict() for m in req.messages],
                "tokens": token_count,
                "elapsed": elapsed,
                "tok_per_sec": tok_per_sec,
                "enable_thinking": enable_thinking,
                "thought_len": len(full_thought),
                "content_len": len(full_content)
            })

        except Exception as ex:
            err_data = json.dumps({"type": "error", "error": str(ex)}, ensure_ascii=False)
            yield f"data: {err_data}\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")

if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Qwen-HTML WebUI")
    ap.add_argument("--port", type=int, default=PORT, help=f"listen port (default {PORT}, env WEBUI_PORT)")
    args = ap.parse_args()
    print(f"Starting Qwen-HTML WebUI on http://0.0.0.0:{args.port}  (api={DEFAULT_API_URL})")
    uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="info")