"""Separate Small demo; the original server files are unchanged."""
import asyncio
from contextlib import asynccontextmanager
import json
import logging
from pathlib import Path
import re
import subprocess
import tempfile
import time
from typing import Literal
import uuid

from fastapi import FastAPI, File, HTTPException, Request, UploadFile
from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse
import numpy as np
from pydantic import BaseModel, Field
import soundfile as sf

import settings
from director import Director
from prompting import build, parse_literal
from storage import AudioStore

log = logging.getLogger("small")
STATE = {"engine": None, "scorers": None, "finisher": None, "error": None, "phase": "starting", "busy": False,
         "job_phase": None, "job_started": None}
store = AudioStore()
sessions = {}
director = Director()


def initialize():
    from engine import SmallEngine, Scorers
    STATE["phase"] = "loading Small S3"
    STATE["engine"] = SmallEngine()
    if settings.FAST_INFERENCE:
        STATE["phase"] = "warming compiled GPU decoder and CUDA graphs"
        STATE["engine"].warm_up()
        STATE["phase"] = "warming chunk-search continuations and cache forks"
        STATE["engine"].warm_search()
    STATE["phase"] = "loading Parakeet and quality scorers"
    STATE["scorers"] = Scorers()
    from alignment import AudioFinisher
    STATE["phase"] = "loading Qwen word aligner"
    STATE["finisher"] = AudioFinisher()
    STATE["phase"] = "ready"


async def initialize_async():
    try:
        await asyncio.to_thread(initialize)
    except Exception as exc:
        log.exception("Model startup failed")
        STATE["error"] = f"{type(exc).__name__}: {exc}"
        STATE["phase"] = "error"


async def cleanup():
    while True:
        await asyncio.sleep(60)
        for item in store.cleanup():
            if STATE["engine"]:
                STATE["engine"].forget_reference(item.path)
        cutoff = time.time() - settings.TTL
        for sid in list(sessions):
            if sessions[sid] < cutoff and not any(x.active and x.owner == sid for x in store.items.values()):
                sessions.pop(sid, None)


@asynccontextmanager
async def lifespan(app):
    boot = asyncio.create_task(initialize_async())
    janitor = asyncio.create_task(cleanup())
    yield
    janitor.cancel()
    await asyncio.gather(janitor, return_exceptions=True)
    await boot
    await director.close()


app = FastAPI(title="Humaneness Voice Small", lifespan=lifespan)


@app.middleware("http")
async def session(request, call_next):
    sid = request.cookies.get("small_session")
    fresh = sid not in sessions
    if fresh:
        sid = uuid.uuid4().hex
    sessions[sid] = time.time()
    request.state.sid = sid
    response = await call_next(request)
    response.set_cookie("small_session", sid, httponly=True, samesite="lax", max_age=settings.TTL,
                        secure=request.headers.get("x-forwarded-proto") == "https")
    response.headers["Cache-Control"] = "no-store"
    return response


@app.get("/", response_class=HTMLResponse)
def index():
    return (settings.HERE / "index.html").read_text()


@app.get("/performance/")
def performance_report():
    path = settings.HERE / "performance/index.html"
    if not path.is_file():
        raise HTTPException(503, "Performance report is being prepared.")
    return FileResponse(path, media_type="text/html")


@app.get("/performance/{filename}")
def performance_data(filename: str):
    if filename not in {"baseline.json", "optimizations.json", "compile.json", "sweep.json", "validation.json", "latency.json", "long.json"}:
        raise HTTPException(404)
    path = settings.HERE / "performance" / filename
    if not path.is_file():
        raise HTTPException(404)
    return FileResponse(path, media_type="application/json")


@app.get("/study/")
def emotion_report():
    path = settings.HERE / "study/index.html"
    if not path.is_file():
        raise HTTPException(503, "Emotion study is being prepared.")
    return FileResponse(path, media_type="text/html")


@app.get("/search/")
def search_report():
    path = settings.HERE / "search/index.html"
    if not path.is_file():
        raise HTTPException(503, "Search report is being prepared.")
    return FileResponse(path, media_type="text/html")


@app.get("/search/{filename}")
def search_data(filename: str):
    if filename not in {"design.json", "results.json", "validation.json", "emotion-check.json"}:
        raise HTTPException(404)
    path = settings.HERE / "search" / filename
    if not path.is_file():
        raise HTTPException(404)
    return FileResponse(path, media_type="application/json")


@app.get("/search/audio/{filename}")
def search_audio(filename: str):
    if not re.fullmatch(r"(?:smoke|measure|quality|audit)-[a-z0-9_-]+\.mp3", filename):
        raise HTTPException(404)
    path = settings.HERE / "search/audio" / filename
    if not path.is_file():
        raise HTTPException(404)
    return FileResponse(path, media_type="audio/mpeg")


@app.get("/study/{filename}")
def study_data(filename: str):
    if filename not in {"design.json", "summary.json", "results.json", "scorer-environment.json"}:
        raise HTTPException(404)
    path = settings.HERE / "study" / filename
    if not path.is_file():
        raise HTTPException(404)
    return FileResponse(path, media_type="application/json")


@app.get("/study/audio/{filename}")
def study_audio(filename: str):
    if not re.fullmatch(r"[a-zA-Z0-9_.-]+\.mp3", filename):
        raise HTTPException(404)
    path = settings.HERE / "study/audio" / filename
    if not path.is_file():
        raise HTTPException(404)
    return FileResponse(path, media_type="audio/mpeg")


@app.get("/legacy.css")
def legacy_css():
    # Actual original demo styles, so changes there remain reflected here.
    original = (settings.PARENT / "index.html").read_text()
    return Response(re.search(r"<style>(.*?)</style>", original, re.S)[1], media_type="text/css")


@app.get("/api/state")
def state(request: Request):
    return {"ready": STATE["phase"] == "ready", "phase": STATE["phase"], "error": STATE["error"],
            "model": settings.REPO, "stage": settings.STAGE, "revision": settings.REVISION,
            "parameters": STATE["engine"].parameters if STATE["engine"] else None,
            "director": settings.LUNA_MODEL, "director_ready": bool(settings.luna_key()),
            "guidance_default": settings.DEFAULT_GUIDANCE,
            "temperature_default": settings.DEFAULT_TEMPERATURE,
            "format_default": settings.DEFAULT_FORMAT, "adapters": [], "ttl_seconds": settings.TTL,
            "audio_format": "mp3", "audio_channels": 1, "audio_bitrate_kbps": 96,
            "alignment_default": True,
            "alignment_available": bool(STATE["finisher"] and STATE["finisher"].available),
            "job_phase": STATE.get("job_phase"),
            "job_seconds": round(time.monotonic() - STATE["job_started"]) if STATE.get("job_started") else None,
            "busy": STATE["busy"], "references": reference_list(request.state.sid)}


def reference_list(owner):
    return [{"id": x.id, "slot": x.slot, "name": x.label,
             "url": f"/api/audio/{x.id}", "expires_at": x.touched + settings.TTL} for x in store.references(owner)]


def normalize(raw, output):
    with tempfile.NamedTemporaryFile(suffix=".upload", dir=store.root) as tmp:
        tmp.write(raw)
        tmp.flush()
        result = subprocess.run(["ffmpeg", "-nostdin", "-v", "error", "-y", "-i", tmp.name,
            "-vn", "-af", "silenceremove=start_periods=1:start_duration=0.02:start_threshold=-45dB",
            "-t", str(settings.REFERENCE_SECONDS), "-ac", "1", "-ar", "48000",
            "-c:a", "pcm_s16le", str(output)], capture_output=True, timeout=30)
        if result.returncode:
            raise ValueError("Die Audiodatei konnte nicht dekodiert werden.")
    wave, sr = sf.read(output, dtype="float32")
    if len(wave) < sr * 0.24 or not np.isfinite(wave).all() or np.max(np.abs(wave)) < 0.001:
        raise ValueError("Bitte ein hörbares Sprachsample von mindestens 0,24 Sekunden hochladen.")
    peak = float(np.max(np.abs(wave)))
    if peak > 0.95:
        wave *= 0.95 / peak
        sf.write(output, wave, sr, subtype="PCM_16")
    return len(wave) / sr


def encode_mp3(source, output):
    """Only the transport is lossy; alignment and judging use original PCM."""
    if isinstance(source, Path):
        input_args, data = ["-i", str(source)], None
    else:
        input_args = ["-f", "f32le", "-ar", "48000", "-ac", "1", "-i", "pipe:0"]
        data = np.asarray(source, dtype="<f4").reshape(-1).tobytes()
    result = subprocess.run(["ffmpeg", "-nostdin", "-v", "error", "-y", *input_args,
        "-map_metadata", "-1", "-vn", "-ac", "1", "-ar", "48000", "-c:a", "libmp3lame",
        "-b:a", "96k", "-write_xing", "1", str(output)], input=data, capture_output=True, timeout=30)
    if result.returncode:
        raise RuntimeError("Das Audio konnte nicht als Mono-MP3 kodiert werden.")


@app.post("/api/reference")
async def upload_reference(request: Request, slot: Literal["A", "B"] = "A", file: UploadFile = File(...)):
    raw = bytearray()
    try:
        while chunk := await file.read(1024 * 1024):
            raw.extend(chunk)
            if len(raw) > settings.MAX_UPLOAD:
                raise HTTPException(413, "Maximal 20 MB pro Referenzaudio.")
    finally:
        await file.close()
    if not raw:
        raise HTTPException(400, "Leere Audiodatei.")
    existing = [x for x in store.references(request.state.sid) if x.slot == slot]
    if any(x.active for x in existing):
        raise HTTPException(409, "Diese Referenz wird gerade verwendet.")
    item = store.new(request.state.sid, Path(file.filename or "reference").name[:100], "pending", slot)
    try:
        seconds = await asyncio.to_thread(normalize, raw, item.path)
        await asyncio.to_thread(encode_mp3, item.path, item.playback)
    except Exception as exc:
        store.remove(item)
        raise HTTPException(400, str(exc)) from exc
    existing = [x for x in store.references(request.state.sid) if x.slot == slot]
    if any(x.active for x in existing):
        store.remove(item)
        raise HTTPException(409, "Diese Referenz wird gerade verwendet.")
    for old in existing:
        if STATE["engine"]:
            STATE["engine"].forget_reference(old.path)
        store.remove(old)
    item.kind = "reference"
    return {"id": item.id, "seconds": round(seconds, 2), "references": reference_list(request.state.sid)}


@app.delete("/api/reference/{id_}")
async def delete_reference(id_: str, request: Request):
    try:
        item = store.get(id_, request.state.sid, touch=False)
        if item.kind != "reference":
            raise KeyError()
        store.remove(item)
        if STATE["engine"]:
            STATE["engine"].forget_reference(item.path)
    except KeyError:
        raise HTTPException(404, "Referenz nicht vorhanden.")
    except ValueError as exc:
        raise HTTPException(409, str(exc))
    return {"references": reference_list(request.state.sid)}


@app.get("/api/audio/{id_}")
def audio(id_: str, request: Request):
    try:
        item = store.get(id_, request.state.sid)
    except KeyError:
        raise HTTPException(404, "Audio abgelaufen oder nicht vorhanden.")
    return FileResponse(item.playback, media_type="audio/mpeg", filename=f"humaneness-{item.kind}.mp3")


class Turn(BaseModel):
    message: str = Field(min_length=1, max_length=5000)
    format: Literal["structured", "caption", "tags", "transcript"] = settings.DEFAULT_FORMAT
    language: Literal["en", "de", "auto"] = "auto"
    input_mode: Literal["chat", "perform", "manual"] = "chat"
    generation: Literal["single", "best", "all", "search_winner", "search_beam"] = "single"
    search_n: int = Field(default=8, ge=1, le=16)
    search_chunk_frames: Literal[10, 20] = 10
    search_seconds: float = Field(default=6.6, ge=0.8, le=60)
    n: int = Field(default=8, ge=1, le=12)
    guidance: float = Field(default=settings.DEFAULT_GUIDANCE, ge=1, le=5)
    align: bool = True
    temperature: float = Field(default=settings.DEFAULT_TEMPERATURE, ge=0, le=1.5)
    streaming: bool = False
    chunk_frames: Literal[2, 5, 10, 20] = 10
    seed: int = Field(default=777, ge=0, le=2147483647)
    seconds: float | None = Field(default=None, ge=1, le=60)
    reference_id: str | None = Field(default=None, max_length=40)
    history: list[dict] = Field(default_factory=list, max_length=32)


def language_for(text, selected, fallback="en"):
    if selected != "auto":
        return selected
    if re.search(r"[äöüß]|\b(ich|nicht|eine|einen|und|habe|mein|mir|dass|bitte|das|ist|heute)\b", text, re.I):
        return "de"
    return "de" if fallback == "de" else "en"


@app.post("/api/turn")
async def turn(body: Turn, request: Request):
    if STATE["phase"] != "ready":
        raise HTTPException(503, STATE["error"] or "Modelle werden geladen.")
    if STATE["busy"]:
        raise HTTPException(409, "Es läuft gerade eine Generierung. Bitte kurz warten.")
    reference = None
    if body.reference_id:
        try:
            reference = store.get(body.reference_id, request.state.sid)
            if reference.kind != "reference":
                raise KeyError()
        except KeyError:
            raise HTTPException(404, "Referenz abgelaufen. Bitte erneut hochladen.")
    STATE["busy"] = True
    STATE["job_phase"] = "Regie" if body.input_mode != "manual" else "Prompt vorbereiten"
    STATE["job_started"] = time.monotonic()
    if reference:
        reference.active += 1
    queue = asyncio.Queue()
    loop = asyncio.get_running_loop()
    owner = request.state.sid

    async def work():
        started = time.monotonic()
        try:
            await queue.put({"event": "status", "phase": "Regie" if body.input_mode != "manual" else "Prompt vorbereiten"})
            if body.input_mode == "manual":
                directed = parse_literal(body.message, body.format)
                llm_ms = 0
            else:
                try:
                    directed, llm_ms = await asyncio.wait_for(director.direct(
                        body.message, body.format, body.language, body.input_mode, body.history), timeout=35)
                except asyncio.TimeoutError as exc:
                    raise RuntimeError("Luna hat nach 35 Sekunden noch nicht geantwortet. Die Anfrage wurde abgebrochen; du kannst erneut senden oder einen manuellen Prompt verwenden.") from exc
            language = language_for(directed["text"], body.language, directed.get("language", "en"))
            spec = build(directed, body.format, language, body.seconds, body.message if body.input_mode == "manual" else None)
            await queue.put({"event": "direction", **spec, "llm_ms": llm_ms, "reference": reference.label if reference else None})
            STATE["job_phase"] = "Audio generieren"
            await queue.put({"event": "status", "phase": "Audio generieren"})
            n = 1 if body.generation == "single" else body.n
            def progress(done, total):
                loop.call_soon_threadsafe(queue.put_nowait, {"event": "progress", "done": done, "total": total})
            generation_started = time.monotonic()
            search_performance = None
            if body.generation.startswith("search_"):
                from beam_search import ChunkSearch
                from search_scoring import SearchScorer
                scorer = await asyncio.to_thread(SearchScorer, STATE["scorers"], spec)
                def committed(wave, detail):
                    loop.call_soon_threadsafe(queue.put_nowait, {"event": "search_chunk", "_wave": wave,
                        "step": detail["step"], "committed_until_seconds": detail["committed_until_seconds"],
                        "seconds": len(wave)/48000, "retained_rows": detail["retained_rows"]})
                wave, search_performance = await asyncio.to_thread(ChunkSearch(STATE["engine"], scorer).generate,
                    spec, body.search_n, 2 if body.generation == "search_beam" else 1,
                    body.search_chunk_frames, body.search_seconds, body.guidance, body.temperature,
                    body.seed, reference.path if reference else None, committed)
                waves = [wave]
            else:
                waves = await asyncio.to_thread(STATE["engine"].generate, spec, n, body.guidance,
                    body.seed, reference.path if reference else None, body.temperature, progress,
                    **({"streaming": True, "chunk_frames": body.chunk_frames} if body.streaming else {}))
            generation_ms = round((time.monotonic() - generation_started) * 1000)
            STATE["job_phase"] = "Sprachgrenzen & Fades"
            await queue.put({"event": "status", "phase": "Sprachgrenzen ausrichten & ein-/ausblenden"})
            alignment_started = time.monotonic()
            waves, alignments = await asyncio.to_thread(STATE["finisher"].finish, waves, spec, body.align)
            alignment_ms = round((time.monotonic() - alignment_started) * 1000)
            STATE["job_phase"] = "Takes bewerten"
            await queue.put({"event": "status", "phase": "Takes bewerten"})
            scoring_started = time.monotonic()
            if search_performance is not None:
                scores = await asyncio.to_thread(scorer.score, waves, spec["text"], [True]*len(waves))
                for i, score in enumerate(scores):
                    score["wer"] = score["prefix_wer"]
                    score["rank"] = i
            else:
                scores = await asyncio.to_thread(STATE["scorers"].score, waves, spec)
            scoring_ms = round((time.monotonic() - scoring_started) * 1000)
            candidates = []
            for i, (wave, score) in enumerate(zip(waves, scores)):
                item = store.new(owner, f"Take {i+1}", "take")
                try:
                    await asyncio.to_thread(encode_mp3, wave, item.playback)
                except Exception:
                    store.remove(item)
                    raise
                candidates.append({**score, "index": i, "url": f"/api/audio/{item.id}", "alignment": alignments[i]})
            best = min(candidates, key=lambda c: c["rank"])
            await queue.put({"event": "result", "candidates": candidates, "winner": best["index"],
                "generation": body.generation, "guidance": body.guidance, "seed": body.seed,
                "llm_ms": llm_ms, "generation_ms": generation_ms, "alignment_ms": alignment_ms, "scoring_ms": scoring_ms,
                "performance": getattr(STATE["engine"], "last_performance", None) if search_performance is None else None,
                "search_performance": search_performance,
                "total_ms": round((time.monotonic() - started) * 1000),
                "cfg_effective": body.guidance != 1 and spec["prompt"] != spec["unconditional"]})
        except Exception as exc:
            log.exception("Turn failed")
            await queue.put({"event": "error", "message": str(exc) or f"{type(exc).__name__}: Generierung fehlgeschlagen."})
        finally:
            if reference:
                reference.active -= 1
                reference.touched = time.time()
                reference.path.touch()
            STATE["busy"] = False
            STATE["job_phase"] = None
            STATE["job_started"] = None
            await queue.put(None)

    task = asyncio.create_task(work())
    # Work survives a disconnected browser, then releases its active reference.
    tasks.add(task)
    task.add_done_callback(tasks.discard)

    async def events():
        while True:
            try:
                item = await asyncio.wait_for(queue.get(), timeout=15)
            except asyncio.TimeoutError:
                yield ": keepalive\n\n"
                continue
            if item is None:
                break
            if item.get("event") == "search_chunk":
                wave = item.pop("_wave")
                chunk = store.new(owner, f"Search chunk {item['step']}", "take")
                try:
                    await asyncio.to_thread(encode_mp3, wave, chunk.playback)
                    item["url"] = f"/api/audio/{chunk.id}"
                except Exception:
                    store.remove(chunk)
                    log.exception("Search chunk encoding failed")
                    item = {"event": "status", "phase": "Live-MP3 konnte nicht kodiert werden; fertigen Take abwarten"}
            yield "data: " + json.dumps(item, ensure_ascii=False) + "\n\n"
    return StreamingResponse(events(), media_type="text/event-stream",
        headers={"X-Accel-Buffering": "no", "Cache-Control": "no-cache"})


tasks = set()


@app.post("/api/asr")
async def asr(request: Request):
    if not STATE["scorers"]:
        raise HTTPException(503, "Spracherkennung wird geladen.")
    if STATE["busy"]:
        raise HTTPException(409, "Bitte die laufende Generierung abwarten.")
    raw = bytearray()
    async for chunk in request.stream():
        raw.extend(chunk)
        if len(raw) > settings.MAX_UPLOAD:
            raise HTTPException(413, "Aufnahme zu groß.")
    if not raw:
        raise HTTPException(400, "Leere Aufnahme.")
    text = await asyncio.to_thread(STATE["scorers"].asr.transcribe, bytes(raw))
    return {"text": text}