"""Deterministic OpenAI-compatible mock for E2E tests (aipi stand-in).

Implements just enough of the aipi surface:

* ``GET  /v1/models``
* ``POST /v1/embeddings`` — real bag-of-words vectors (768-dim, L2-normed).
  Because similarity is *genuine token overlap*, the relevance threshold
  behaves the same way it will in production: related questions score high,
  unrelated ones score low and trigger honest deflection.
* ``POST /v1/chat/completions`` — streaming (SSE) or not. The content keys
  off markers in the system prompt:
  - user message containing ``write a long answer`` -> a ~900-word
    deterministic numbered answer (long-answers story, phase 11)
  - ``DEFLECT_MODE``  -> honest "I haven't done anything like that" answer
  - otherwise         -> upbeat answer quoting the provided document context
  - user message containing ``pretend to think slowly`` -> 3s warm-up delay
    (used by the loading-feedback story).
  - user message containing ``think out loud`` -> the answer is preceded by
    ~800 chars of deterministic ``reasoning_content`` chunks (the
    thinking-display story, phase 17).
  - user message containing ``think out loud then hesitate`` -> the
    ``think out loud`` stream, then a 4s pause before the first content
    frame (the sources-midstream story, phase 20 — a deterministic
    "leave during pure thinking" navigation window).
  - system prompt containing ``<tuning>`` (phase 15, steering notes) ->
    the composed answer ends with `` (tuning: <first note line>)`` —
    makes prompt injection observable in the UI deterministically.

  ``max_tokens`` is honored deterministically (token ≈ whitespace word),
  like a real endpoint: an answer longer than the cap is truncated. This
  is what makes the phase-11 truncation regression observable.
"""
from __future__ import annotations

import hashlib
import math
import os
import re
import signal
import threading
import time
import uuid
from typing import Any

from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

DIM = 768
TOKEN_RE = re.compile(r"[a-z0-9]+")


def embed_text(text: str) -> list[float]:
    vec = [0.0] * DIM
    for tok in TOKEN_RE.findall(text.lower()):
        idx = int(hashlib.md5(tok.encode()).hexdigest(), 16) % DIM
        vec[idx] += 1.0
    norm = math.sqrt(sum(v * v for v in vec)) or 1.0
    return [v / norm for v in vec]


def _messages(body: dict[str, Any]) -> list[dict[str, str]]:
    return body.get("messages", [])


def _system(body: dict[str, Any]) -> str:
    return " ".join(m.get("content", "") for m in _messages(body) if m.get("role") == "system")


def _user(body: dict[str, Any]) -> str:
    parts = [m.get("content", "") for m in _messages(body) if m.get("role") == "user"]
    return parts[-1] if parts else ""


def _context(body: dict[str, Any]) -> str:
    """The document context is the longest system/user message in practice."""
    msgs = _messages(body)
    return max((m.get("content", "") for m in msgs), key=len)


LONG_ANSWER_TRIGGER = "write a long answer"
#: ~920 words — comfortably past the old hard 700-token cap (where the
#: tail would be cut) yet short enough to stream in ~8s at the mock's
#: per-chunk pacing.
LONG_ANSWER_LINES = 40
LONG_ANSWER_END = "LONG-ANSWER-END"

#: Phase 17 (thinking-display story): a user message containing this
#: substring (case-insensitive) is answered with a deterministic
#: ``reasoning_content`` stream ahead of the content — same convention as
#: the other user-message triggers above. Existing E2E questions do not
#: contain the substring, so every other suite is unaffected.
THINKING_TRIGGER = "think out loud"

#: Phase 20 (sources-midstream bug): a user message containing this
#: substring (case-insensitive) gets the phase-17 thinking stream followed
#: by a multi-second pause before the FIRST content frame — the
#: navigation window for the "leave during pure thinking" scenario
#: (owner-confirmed A1.2: nothing brain-side may be persisted then).
#: Strictly longer than ``THINKING_TRIGGER``, so the phase-17 suite's
#: questions are unaffected.
SLOW_PRETOKEN_TRIGGER = "think out loud then hesitate"
PRE_CONTENT_PAUSE_S = 4.0


def long_answer() -> str:
    """~900-word deterministic walkthrough (phase 11): numbered steps plus
    a unique final line that must survive the stream untruncated."""
    lines = [
        f"{i}. Step {i}: configure node-{i} with the homelab defaults and "
        f"verify that step {i} of the long walkthrough is complete before moving on."
        for i in range(1, LONG_ANSWER_LINES + 1)
    ]
    lines.append(LONG_ANSWER_END)
    return "\n".join(lines)


#: First numbered note line of a ``<tuning>`` section (phase 15).
_TUNING_BLOCK_RE = re.compile(r"<tuning>\n(.*?)\n</tuning>", re.S)
_NOTE_LINE_RE = re.compile(r"^\d+\.\s*(.+)$")


def first_tuning_note(system: str) -> str | None:
    """The first steering note in the system prompt, or ``None``.

    The prompt numbers notes 1..N oldest-first (see
    ``app.rag.prompts.build_steering_section``); the mock echoes the first
    one into its answer so prompt injection is observable in the UI.
    """
    block = _TUNING_BLOCK_RE.search(system)
    if not block:
        return None
    for line in block.group(1).splitlines():
        m = _NOTE_LINE_RE.match(line.strip())
        if m:
            return m.group(1).strip()
    return None


def compose_answer(body: dict[str, Any]) -> str:
    system = _system(body)
    user = _user(body)
    if LONG_ANSWER_TRIGGER in user.lower():
        answer = long_answer()
    elif "DEFLECT_MODE" in system:
        answer = (
            "Ah — I haven't done anything like that, so I don't want to make stuff up! "
            "You're thinking bigger than my notes for a second. Try asking about "
            "kubernetes, backups, or deploying a new service — I know those inside out. "
            "You've got this!"
        )
    else:
        ctx = _context(body)
        snippet = ctx[:220].replace("\n", " ").strip()
        answer = (
            f"Great question — you've absolutely got this! Here's what my notes say about "
            f"“{user.strip()[:80]}”: {snippet}… That's the gist from the docs; happy to "
            "dig into any of it. (Deterministic mock answer for E2E.)"
        )
    # Steering (phase 15): when the system prompt carries <tuning>, the
    # answer ends with the first note — deterministically observable.
    note = first_tuning_note(system)
    if note:
        answer = f"{answer} (tuning: {note})"
    return answer


def compose_thinking(body: dict[str, Any]) -> str:
    """Deterministic reasoning scratchpad (thinking-display story, phase 17).

    A fixed 4-line "Step 1… Step 4" template quoting the first ~60 chars
    of the user question: unique per question, byte-stable across runs,
    ~700–900 chars total (≈ 60–75 frames at the mock's 12-char/0.02s
    pacing). The ``Step 2: Check my notes`` line fragment is what the E2E
    assertions key off.
    """
    q = _user(body).strip()[:60]
    return (
        f"Step 1: Read the question carefully — “{q}” — and figure out what kind of "
        "answer it wants (a how-to, a lookup, or a design decision) before touching "
        "the docs, so I don't over- or under-answer.\n"
        "Step 2: Check my notes for the closest match. The homelab kubernetes file "
        "is the obvious candidate, but I should also consider whether a deployments "
        "note covers the same ground better.\n"
        "Step 3: Re-read the relevant sections top to bottom so every specific — "
        "hosts, versions, ports, schedules — is exact as written rather than "
        "remembered, and note which document each fact comes from.\n"
        "Step 4: Draft the answer around those specifics, keep it tight with short "
        "paragraphs and bullets where it helps, cite the documents by path, and "
        "double-check that nothing is invented."
    )


@app.post("/__shutdown__")
def shutdown() -> dict[str, Any]:
    """Test hook (loading-feedback story): terminate this mock process to
    simulate an LLM outage. The E2E fixture restores a fresh instance on
    the same port afterwards, so the rest of the session keeps working."""

    def _die() -> None:
        time.sleep(0.1)  # let the HTTP response flush before we exit
        os.kill(os.getpid(), signal.SIGTERM)

    threading.Thread(target=_die, daemon=True).start()
    return {"status": "shutting down"}


@app.get("/v1/models")
def models() -> dict[str, Any]:
    return {
        "object": "list",
        "data": [
            {"id": "turbo", "object": "model"},
            {"id": "embed", "object": "model"},
            {"id": "lite", "object": "model"},
        ],
    }


@app.post("/v1/embeddings")
def embeddings(body: dict[str, Any]) -> dict[str, Any]:
    raw = body.get("input")
    if isinstance(raw, str):
        raw = [raw]
    inputs: list[Any] = list(raw) if isinstance(raw, list) else []
    data = [
        {"object": "embedding", "index": i, "embedding": embed_text(t)}
        for i, t in enumerate(inputs)
    ]
    return {
        "object": "list",
        "data": data,
        "model": body.get("model", "embed"),
        "usage": {"prompt_tokens": 8, "total_tokens": 8},
    }


def _sse_stream(
    answer: str, delay: float, thinking: str = "", pre_content_delay: float = 0.0
) -> Any:
    """SSE frames for one chat completion (phase 17: + reasoning).

    When ``thinking`` is non-empty its 12-char slices go out FIRST as
    ``delta.reasoning_content`` frames — same 0.02s cadence and envelope
    as the content frames, the aipi wire convention (reasoning before
    content). Without ``thinking`` the output is byte-identical to the
    content-only stream, so the other story suites are unaffected.

    ``pre_content_delay`` (phase 20) inserts a silence gap between the end
    of the thinking stream and the first content frame — the client stays
    in its pre-token "thinking" state the whole time (0.02s cadence and
    frame shapes are unchanged, so 0.0 is byte-identical to before).
    """
    model = "turbo"
    chunk_id = f"chatcmpl-{uuid.uuid4()}"
    if delay:
        time.sleep(delay)
    for piece in re.findall(r".{1,12}", thinking, re.S):
        payload = {
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": model,
            "choices": [
                {"index": 0, "delta": {"reasoning_content": piece}, "finish_reason": None}
            ],
        }
        yield f"data: {json_dumps(payload)}\n\n"
        time.sleep(0.02)
    if pre_content_delay:
        time.sleep(pre_content_delay)
    for piece in re.findall(r".{1,12}", answer, re.S):
        payload = {
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": model,
            "choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}],
        }
        yield f"data: {json_dumps(payload)}\n\n"
        time.sleep(0.02)
    yield (
        "data: "
        + json_dumps(
            {
                "id": chunk_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": model,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
            }
        )
        + "\n\n"
    )
    yield "data: [DONE]\n\n"


def json_dumps(obj: dict[str, Any]) -> str:
    import json

    return json.dumps(obj)


def _apply_max_tokens(answer: str, max_tokens: Any) -> str:
    """Deterministic stand-in for the endpoint's output cap: one token ≈
    one whitespace-separated word. Answers within the cap pass through
    byte-identical, so existing (short) answers are unaffected."""
    if not isinstance(max_tokens, int) or max_tokens <= 0:
        return answer
    words = answer.split()
    if len(words) <= max_tokens:
        return answer
    return " ".join(words[:max_tokens])


@app.post("/v1/chat/completions")
def chat_completions(body: dict[str, Any]) -> Any:
    answer = _apply_max_tokens(compose_answer(body), body.get("max_tokens"))
    delay = 3.0 if "pretend to think slowly" in _user(body) else 0.0
    user_lower = _user(body).lower()
    thinking = compose_thinking(body) if THINKING_TRIGGER in user_lower else ""
    pre_content = (
        PRE_CONTENT_PAUSE_S if SLOW_PRETOKEN_TRIGGER in user_lower else 0.0
    )

    if not body.get("stream"):
        message: dict[str, Any] = {"role": "assistant", "content": answer}
        if thinking:
            # Harmless future-proofing: the app only uses streaming, but a
            # non-streaming client that reads the field gets the reasoning.
            message["reasoning_content"] = thinking
        return {
            "id": f"chatcmpl-{uuid.uuid4()}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": body.get("model", "turbo"),
            "choices": [
                {"index": 0, "message": message, "finish_reason": "stop"}
            ],
            "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
        }

    return StreamingResponse(
        _sse_stream(answer, delay, thinking=thinking, pre_content_delay=pre_content),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )