#!/usr/bin/env python3
import json
import sys
import urllib.request
from datetime import datetime
from pathlib import Path

FASTAPI_SERVICES = [
    ("Base Platform", 8000),
    ("Training Center", 8001),
    ("Data Center", 8002),
    ("Tax Accelerator", 8003),
]

LLM_SERVICES = [
    ("vLLM (Main LLM)", 8100, {
        "openapi": "3.1.0",
        "info": {"title": "vLLM — Main LLM", "version": "1.0.0"},
        "servers": [{"url": "http://localhost:8100"}],
        "paths": {
            "/health": {"get": {"summary": "Health check", "tags": ["Status"], "responses": {"200": {"description": "OK"}}}},
            "/v1/models": {"get": {"summary": "List models", "tags": ["Models"], "responses": {"200": {"description": "Model list"}}}},
            "/v1/chat/completions": {"post": {"summary": "Chat completions", "tags": ["Chat"], "requestBody": {"required": True, "content": {"application/json": {"schema": {"type": "object", "required": ["model", "messages"], "properties": {"model": {"type": "string", "example": "Qwen3.5-27B-AWQ"}, "messages": {"type": "array", "items": {"type": "object", "properties": {"role": {"type": "string", "enum": ["system", "user", "assistant"]}, "content": {"type": "string"}}}}, "stream": {"type": "boolean", "default": False}, "temperature": {"type": "number", "default": 0.7}, "max_tokens": {"type": "integer", "default": 2048}}}}}}, "responses": {"200": {"description": "Completion response"}}}},
            "/v1/completions": {"post": {"summary": "Text completions", "tags": ["Completions"], "requestBody": {"required": True, "content": {"application/json": {"schema": {"type": "object", "required": ["model", "prompt"], "properties": {"model": {"type": "string"}, "prompt": {"type": "string"}, "max_tokens": {"type": "integer", "default": 256}, "temperature": {"type": "number", "default": 0.7}}}}}}, "responses": {"200": {"description": "Completion response"}}}},
        },
    }),
    ("Embedding Server", 8200, {
        "openapi": "3.1.0",
        "info": {"title": "Embedding Server", "version": "1.0.0"},
        "servers": [{"url": "http://localhost:8200"}],
        "paths": {
            "/health": {"get": {"summary": "Health check", "tags": ["Status"], "responses": {"200": {"description": "OK"}}}},
            "/v1/embeddings": {"post": {"summary": "Create embeddings", "tags": ["Embeddings"], "requestBody": {"required": True, "content": {"application/json": {"schema": {"type": "object", "required": ["input"], "properties": {"input": {"oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], "example": "你好，世界"}, "model": {"type": "string", "default": "Qwen3-Embedding-8B"}, "input_type": {"type": "string", "enum": ["document", "query"], "default": "document"}}}}}}, "responses": {"200": {"description": "Embedding vectors"}}}},
        },
    }),
    ("Rerank Server", 8300, {
        "openapi": "3.1.0",
        "info": {"title": "Rerank Server", "version": "1.0.0"},
        "servers": [{"url": "http://localhost:8300"}],
        "paths": {
            "/health": {"get": {"summary": "Health check", "tags": ["Status"], "responses": {"200": {"description": "OK"}}}},
            "/v1/rerank": {"post": {"summary": "Rerank documents", "tags": ["Rerank"], "requestBody": {"required": True, "content": {"application/json": {"schema": {"type": "object", "required": ["query", "documents"], "properties": {"query": {"type": "string", "example": "增值税税率是多少"}, "documents": {"type": "array", "items": {"type": "string"}, "example": ["增值税一般税率为13%", "企业所得税税率为25%"]}, "model": {"type": "string", "default": "Qwen3-Reranker-8B"}, "top_n": {"type": "integer", "default": 5}}}}}}, "responses": {"200": {"description": "Reranked results"}}}},
        },
    }),
    ("OCR / VL Server", 8400, {
        "openapi": "3.1.0",
        "info": {"title": "OCR / Vision-Language Server", "version": "1.0.0"},
        "servers": [{"url": "http://localhost:8400"}],
        "paths": {
            "/health": {"get": {"summary": "Health check", "tags": ["Status"], "responses": {"200": {"description": "OK"}}}},
            "/v1/chat/completions": {"post": {"summary": "Vision chat completions", "tags": ["Chat"], "requestBody": {"required": True, "content": {"application/json": {"schema": {"type": "object", "required": ["model", "messages"], "properties": {"model": {"type": "string", "example": "Qwen3-VL-8B-Instruct"}, "messages": {"type": "array", "items": {"type": "object", "properties": {"role": {"type": "string"}, "content": {"type": "string"}}}}, "max_tokens": {"type": "integer", "default": 2048}, "temperature": {"type": "number", "default": 0.7}}}}}}, "responses": {"200": {"description": "Chat response"}}}},
        },
    }),
]


def fetch_json(url, timeout=5):
    try:
        with urllib.request.urlopen(url, timeout=timeout) as r:
            return json.loads(r.read())
    except Exception:
        return None


def check_alive(port, timeout=3):
    try:
        urllib.request.urlopen(f"http://localhost:{port}/health", timeout=timeout)
        return True
    except Exception:
        return False


def build_sources():
    sources = []
    for name, port in FASTAPI_SERVICES:
        print(f"Fetching {name} (port {port})... ", end="", flush=True)
        spec = fetch_json(f"http://localhost:{port}/openapi.json")
        if spec:
            spec.setdefault("servers", [{"url": f"http://localhost:{port}"}])
            sources.append({"title": name, "port": port, "online": True, "spec": spec})
            print("OK")
        else:
            print("SKIPPED (not running)")

    for name, port, spec in LLM_SERVICES:
        online = check_alive(port)
        print(f"Checking {name} (port {port})... {'OK' if online else 'OK (offline — docs only)'}")
        sources.append({"title": name, "port": port, "online": online, "spec": spec})

    return sources


def render(sources, output):
    gen_time = datetime.now().strftime("%Y-%m-%d %H:%M")
    # Embed specs as a JS array — each spec is a separate variable to avoid one giant JSON blob
    specs_js = "const SPECS = " + json.dumps([s["spec"] for s in sources], ensure_ascii=False) + ";"
    tabs_js = "const TABS = " + json.dumps(
        [{"title": s["title"], "online": s["online"]} for s in sources], ensure_ascii=False
    ) + ";"

    html = f"""<!DOCTYPE html>
<html lang="zh">
<head>
  <meta charset="utf-8">
  <title>HelloTax AI — API Docs</title>
  <meta name="viewport" content="width=device-width,initial-scale=1">
  <style>
    *{{box-sizing:border-box;margin:0;padding:0}}
    body{{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',sans-serif;background:#0f0f13;color:#e2e8f0;height:100vh;display:flex;flex-direction:column;overflow:hidden}}
    #bar{{background:#16161e;border-bottom:1px solid #2a2a3d;padding:0 20px;height:52px;display:flex;align-items:center;gap:14px;flex-shrink:0;z-index:999}}
    .logo{{font-size:15px;font-weight:700;color:#fff;letter-spacing:-.3px;white-space:nowrap}}
    .logo span{{color:#7c6af7}}
    #tabs{{display:flex;gap:4px;overflow-x:auto;flex:1;scrollbar-width:none}}
    #tabs::-webkit-scrollbar{{display:none}}
    .tab{{padding:5px 13px;border-radius:6px;font-size:12.5px;cursor:pointer;white-space:nowrap;color:#94a3b8;border:1px solid transparent;transition:all .15s;display:flex;align-items:center;gap:6px;user-select:none}}
    .tab:hover{{color:#e2e8f0;background:#1e1e2e}}
    .tab.active{{color:#fff;background:#7c6af7;border-color:#7c6af7}}
    .dot{{width:6px;height:6px;border-radius:50%;background:#4ade80;flex-shrink:0}}
    .dot.off{{background:#4a5568}}
    #gen{{font-size:11px;color:#3d3d55;white-space:nowrap}}
    #theme-btn{{background:none;border:1px solid #2a2a3d;color:#94a3b8;padding:5px 10px;border-radius:6px;cursor:pointer;font-size:12px;white-space:nowrap;transition:all .15s}}
    #theme-btn:hover{{color:#e2e8f0;border-color:#7c6af7}}
    #scalar-mount{{flex:1;overflow:auto}}
  </style>
</head>
<body>
<div id="bar">
  <div class="logo">HelloTax <span>AI</span></div>
  <div id="tabs"></div>
  <button id="theme-btn" onclick="toggleTheme()">☀ 浅色</button>
  <div id="gen">Generated {gen_time}</div>
</div>
<div id="scalar-mount"></div>

<script src="https://cdn.jsdelivr.net/npm/@scalar/api-reference@latest/dist/browser/standalone.js" crossorigin="anonymous"></script>
<script>
{specs_js}
{tabs_js}

let currentRef = null;
let darkMode = true;
let currentIdx = 0;

function activate(idx) {{
  currentIdx = idx;
  document.querySelectorAll('.tab').forEach((t, i) => t.classList.toggle('active', i === idx));
  const mount = document.getElementById('scalar-mount');
  mount.innerHTML = '';
  if (currentRef && currentRef.destroy) currentRef.destroy();
  currentRef = Scalar.createApiReference(mount, {{
    spec: {{ content: SPECS[idx] }},
    theme: 'purple',
    darkMode: darkMode,
    hideDownloadButton: false,
    defaultHttpClient: {{ targetKey: 'javascript', clientKey: 'fetch' }},
  }});
}}

function toggleTheme() {{
  darkMode = !darkMode;
  const btn = document.getElementById('theme-btn');
  btn.textContent = darkMode ? '☀ 浅色' : '🌙 深色';
  document.getElementById('bar').style.background = darkMode ? '#16161e' : '#ffffff';
  document.getElementById('bar').style.borderBottomColor = darkMode ? '#2a2a3d' : '#e2e8f0';
  document.body.style.background = darkMode ? '#0f0f13' : '#f8f9fa';
  activate(currentIdx);
}}

const tabsEl = document.getElementById('tabs');
TABS.forEach((t, i) => {{
  const el = document.createElement('div');
  el.className = 'tab';
  el.innerHTML = `<span class="dot ${{t.online ? '' : 'off'}}"></span>${{t.title}}`;
  el.addEventListener('click', () => activate(i));
  tabsEl.appendChild(el);
}});

activate(0);
</script>
</body>
</html>"""

    Path(output).parent.mkdir(parents=True, exist_ok=True)
    Path(output).write_text(html, encoding="utf-8")
    print(f"Generated: {output}")


if __name__ == "__main__":
    output = sys.argv[1] if len(sys.argv) > 1 else "docs/api/index.html"
    sources = build_sources()
    if not sources:
        print("No sources available.")
        sys.exit(1)
    render(sources, output)