"""
Laguna S 2.1 — reasoning terminal.

A gradio.Server app: the FastAPI layer exposes one streaming endpoint at
/chat (SSE through Gradio's queue) and serves a hand-rolled HTML/CSS/JS
terminal from "/". The frontend uses Gradio's JS client to call /chat, so
it picks up queuing, SSE streaming, and concurrency control for free.

Reasoning is preserved across turns: the frontend re-sends prior assistant
turns with their `reasoning_details`, so Laguna S 2.1 continues reasoning
from where it left off on the previous call.
"""

import json
import os
import time
from typing import Iterator

from openai import OpenAI

from gradio import Server
from fastapi.responses import HTMLResponse

# ---------------------------------------------------------------------------
# Backend: OpenAI-compatible client pointed at OpenRouter.
# ---------------------------------------------------------------------------
API_KEY = os.environ.get("OPENROUTER_API_KEY") or os.environ.get("API_KEY") or ""
MODEL = os.environ.get("MODEL", "poolside/laguna-s-2.1:free")

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=API_KEY or "missing-key",
)

# ---------------------------------------------------------------------------
# ANSI palette for the terminal look. The frontend keeps this table so it
# can render lines verbatim (no markdown parsing).
# ---------------------------------------------------------------------------
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
ITALIC = "\033[3m"
GREEN = "\033[32m"
CYAN = "\033[36m"
YELLOW = "\033[33m"
MAGENTA = "\033[35m"
RED = "\033[31m"
GREY = "\033[90m"
BOLD_GREEN = BOLD + GREEN
BOLD_CYAN = BOLD + CYAN


def _extract_reasoning(reasoning_details) -> str:
    """Read text out of OpenRouter reasoning_details (list/dict/str, nested)."""
    if not reasoning_details:
        return ""
    chunks: list[str] = []

    def push(val):
        if isinstance(val, str):
            chunks.append(val)
        elif isinstance(val, list):
            for v in val:
                push(v)
        elif isinstance(val, dict):
            for k in ("summary", "text", "content", "reasoning"):
                if k in val:
                    push(val[k])

    if isinstance(reasoning_details, list):
        for p in reasoning_details:
            push(p)
    elif isinstance(reasoning_details, dict):
        push(reasoning_details)
    elif isinstance(reasoning_details, str):
        chunks.append(reasoning_details)

    return "\n".join(c for c in chunks if c).strip()


def _one_line(text: str) -> str:
    return (text or "").replace("\n", " ").strip()


def _emit(payload: dict) -> dict:
    """Strip non-JSON-safe bits before yielding (we only stream ASCII)."""
    return {"event": payload.get("event", "delta"), "text": payload.get("text", "")}


def _coerce_messages(messages):
    """Convert a JSON-decoded `messages` arg into the OpenAI list-of-dicts
    shape. The JS client sends messages as a JSON-encoded string; the Python
    api decorates parameter binding independently. Normalize here so the
    handler works for either."""
    if messages is None:
        return []
    if isinstance(messages, str):
        try:
            import json as _json
            decoded = _json.loads(messages)
            return decoded if isinstance(decoded, list) else []
        except Exception:
            return []
    if isinstance(messages, list):
        return messages
    return []


def _coerce_bool(value, default=True):
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        s = value.strip().lower()
        if s in ("false", "0", "no", "off", ""):
            return False if s != "" else default
        if s in ("true", "1", "yes", "on"):
            return True
    return bool(value) if value is not None else default


# ---------------------------------------------------------------------------
# /chat endpoint. With gradio.Server, a generator-typed function streams via
# SSE. We yield small text events (think: print-style lines the terminal
# renders as they arrive) and finish with a `done` event carrying the new
# assistant message — the frontend appends that to its in-memory history
# along with the original `reasoning_details`.
# ---------------------------------------------------------------------------
app = Server()


@app.api()
def chat(messages: list, reasoning: bool = True) -> Iterator[dict]:
    """Reasoning-aware chat endpoint. Takes `messages` (list of
    role/content/reasoning_details dicts) and an optional `reasoning` bool.
    Streams `{"event":"delta","text":...}` lines, then a final `done` event
    with the new assistant message so the frontend can persist it."""
    messages = _coerce_messages(messages)
    reasoning = _coerce_bool(reasoning, default=True)

    last_user = next(
        (m for m in reversed(messages) if m.get("role") == "user"), None
    )
    user_text = (last_user or {}).get("content", "")

    yield _emit({
        "event": "delta",
        "text": f"{GREY}$ {RESET}{GREEN}user{RESET} {DIM}»{RESET} {_one_line(user_text)}\n",
    })

    yield _emit({
        "event": "delta",
        "text": (
            f"{CYAN}laguna{RESET}{GREY}@{RESET}{YELLOW}openrouter{RESET}"
            f"{GREY}:{RESET}{DIM}~{RESET} {GREY}{'thinking…' if reasoning else 'responding…'}{RESET}\n"
        ),
    })

    try:
        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            extra_body={"reasoning": {"enabled": reasoning}},
        )
        message = response.choices[0].message
    except Exception as e:
        yield _emit({
            "event": "delta",
            "text": f"\n{BOLD}{RED}error{RESET} {DIM}»{RESET} {e}\n",
        })
        yield {"event": "done", "text": "", "assistant": None}
        return

    reasoning_text = _extract_reasoning(getattr(message, "reasoning_details", None))

    if reasoning_text:
        yield _emit({
            "event": "delta",
            "text": f"{DIM}{ITALIC}┌─ reasoning ───────────────────────{RESET}\n",
        })
        for line in reasoning_text.splitlines() or [""]:
            prefix = f"{DIM}{ITALIC}│ {RESET}" if line else f"{DIM}│{RESET} "
            yield _emit({"event": "delta", "text": prefix + line + "\n"})
            time.sleep(0.012)
        yield _emit({
            "event": "delta",
            "text": f"{DIM}{ITALIC}└───────────────────────────────────{RESET}\n",
        })

    answer = (message.content or "").strip()
    yield _emit({
        "event": "delta",
        "text": f"{BOLD_GREEN}┌─ laguna ──────────────────────────{RESET}\n",
    })
    for line in answer.splitlines() or [""]:
        prefix = f"{BOLD_GREEN}│ {RESET}" if line else f"{GREEN}│{RESET} "
        yield _emit({"event": "delta", "text": prefix + line + "\n"})
        time.sleep(0.012)
    yield _emit({
        "event": "delta",
        "text": f"{BOLD_GREEN}└───────────────────────────────────{RESET}\n",
    })
    yield _emit({"event": "delta", "text": f"{GREY}── done ──{RESET}\n"})

    # The new assistant message — keep reasoning_details so the JS client
    # can re-send it on the next turn, letting Laguna continue its reasoning.
    rd = getattr(message, "reasoning_details", None)
    rd_json = json.loads(json.dumps(rd, default=str)) if rd else None

    yield {
        "event": "done",
        "text": "",
        "assistant": {
            "role": "assistant",
            "content": message.content or "",
            "reasoning_details": rd_json,
        },
    }


# ---------------------------------------------------------------------------
# "/" serves the hand-rolled HTML terminal. It uses Gradio's JS client
# from a CDN, connects to the same origin, and calls /chat with SSE.
# ---------------------------------------------------------------------------
@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


@app.get("/healthz")
async def healthz():
    return {"ok": True, "model": MODEL}


if __name__ == "__main__":
    app.launch(show_error=True)