from __future__ import annotations

import asyncio
import os
import shutil
import sqlite3
import subprocess
import tempfile
import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit, urlunsplit

import httpx
import msgpack
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field


DATA_DIR = Path(os.getenv("DATA_DIR", "./data")).resolve()
REFERENCES_DIR = DATA_DIR / "references"
GENERATIONS_DIR = DATA_DIR / "generations"
DB_PATH = DATA_DIR / "studio.sqlite3"
FISH_SPEECH_URL = os.getenv("FISH_SPEECH_URL", "http://host.docker.internal:8080/v1/tts")
FISH_SPEECH_API_KEY = os.getenv("FISH_SPEECH_API_KEY", "YOUR_API_KEY")
MAX_REFERENCE_BYTES = int(os.getenv("MAX_REFERENCE_BYTES", str(25 * 1024 * 1024)))
PAGE_SIZE = 50


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def connect_db() -> sqlite3.Connection:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    REFERENCES_DIR.mkdir(parents=True, exist_ok=True)
    GENERATIONS_DIR.mkdir(parents=True, exist_ok=True)
    with connect_db() as db:
        db.executescript(
            """
            CREATE TABLE IF NOT EXISTS references_audio (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                transcript TEXT NOT NULL,
                filename TEXT NOT NULL,
                created_at TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS generations (
                id TEXT PRIMARY KEY,
                text TEXT NOT NULL,
                reference_id TEXT NOT NULL,
                tts_url TEXT,
                status TEXT NOT NULL,
                filename TEXT,
                error TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                FOREIGN KEY(reference_id) REFERENCES references_audio(id)
            );
            CREATE INDEX IF NOT EXISTS generations_created_at_idx
                ON generations(created_at DESC);
            CREATE TABLE IF NOT EXISTS tts_endpoints (
                id TEXT PRIMARY KEY,
                url TEXT NOT NULL UNIQUE,
                created_at TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS app_settings (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            );
            """
        )
        generation_columns = {row["name"] for row in db.execute("PRAGMA table_info(generations)")}
        if "tts_url" not in generation_columns:
            db.execute("ALTER TABLE generations ADD COLUMN tts_url TEXT")
        db.execute("UPDATE generations SET tts_url=? WHERE tts_url IS NULL OR tts_url=' '", (FISH_SPEECH_URL,))
        # A running request cannot be resumed after the app process restarts.
        db.execute(
            "UPDATE generations SET status='failed', error='App restarted while this job was running', updated_at=? WHERE status='running'",
            (now_iso(),),
        )


def get_reference(reference_id: str) -> sqlite3.Row | None:
    with connect_db() as db:
        return db.execute("SELECT * FROM references_audio WHERE id=?", (reference_id,)).fetchone()


def get_generation(job_id: str) -> sqlite3.Row | None:
    with connect_db() as db:
        return db.execute("SELECT * FROM generations WHERE id=?", (job_id,)).fetchone()


def normalize_tts_url(value: str) -> str:
    raw = value.strip()
    try:
        parsed = urlsplit(raw)
        _ = parsed.port  # Validate malformed ports before the request is made.
    except ValueError as exc:
        raise HTTPException(status_code=422, detail="Enter a valid Fish Speech URL") from exc
    if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password:
        raise HTTPException(status_code=422, detail="Use an HTTP or HTTPS Fish Speech URL without embedded credentials")
    if parsed.query or parsed.fragment:
        raise HTTPException(status_code=422, detail="Fish Speech URL cannot contain a query string or fragment")
    path = parsed.path.rstrip("/")
    if path.endswith("/v1/tts"):
        tts_path = path
    else:
        tts_path = f"{path}/v1/tts" if path else "/v1/tts"
    return urlunsplit((parsed.scheme, parsed.netloc, tts_path, "", ""))


def active_tts_url() -> str:
    with connect_db() as db:
        row = db.execute("SELECT value FROM app_settings WHERE key='active_tts_url'").fetchone()
    return row["value"] if row else normalize_tts_url(FISH_SPEECH_URL)


async def probe_tts_url(value: str) -> tuple[str, str]:
    tts_url = normalize_tts_url(value)
    parsed = urlsplit(tts_url)
    endpoint_path = parsed.path
    prefix = endpoint_path[:-len("/v1/tts")] if endpoint_path.endswith("/v1/tts") else ""
    health_url = urlunsplit((parsed.scheme, parsed.netloc, f"{prefix}/v1/health", "", ""))
    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(5.0, connect=3.0)) as client:
            response = await client.get(health_url, headers={"Authorization": f"Bearer {FISH_SPEECH_API_KEY}"})
        if response.is_error:
            raise HTTPException(status_code=502, detail=f"Fish Speech health check returned HTTP {response.status_code} at {health_url}")
        health = response.json()
        if not isinstance(health, dict) or health.get("status") != "ok":
            raise HTTPException(status_code=502, detail="Server responded, but did not return a healthy Fish Speech status")
    except HTTPException:
        raise
    except (httpx.TimeoutException, httpx.HTTPError) as exc:
        raise HTTPException(status_code=502, detail=f"Could not reach Fish Speech at {tts_url}: {exc}") from exc
    except ValueError as exc:
        raise HTTPException(status_code=502, detail="Fish Speech returned an unreadable health response") from exc
    return tts_url, "Fish Speech server is reachable and healthy."


def reference_json(row: sqlite3.Row) -> dict[str, Any]:
    return {
        "id": row["id"],
        "name": row["name"],
        "transcript": row["transcript"],
        "created_at": row["created_at"],
        "audio_url": f"/api/references/{row['id']}/audio",
    }


def generation_json(row: sqlite3.Row) -> dict[str, Any]:
    return {
        "id": row["id"],
        "text": row["text"],
        "reference_id": row["reference_id"],
        "status": row["status"],
        "error": row["error"],
        "created_at": row["created_at"],
        "updated_at": row["updated_at"],
        "audio_url": f"/api/generations/{row['id']}/audio" if row["filename"] else None,
    }


class CreateGeneration(BaseModel):
    text: str = Field(min_length=1, max_length=10000)
    reference_id: str


class TtsUrlRequest(BaseModel):
    url: str = Field(min_length=1, max_length=500)


class JobQueue:
    def __init__(self) -> None:
        self.queue: asyncio.Queue[str] = asyncio.Queue()
        self.worker_task: asyncio.Task[None] | None = None
        self.http: httpx.AsyncClient | None = None

    async def start(self) -> None:
        init_db()
        self.http = httpx.AsyncClient(timeout=httpx.Timeout(1800.0, connect=10.0))
        with connect_db() as db:
            pending = db.execute("SELECT id FROM generations WHERE status='queued' ORDER BY created_at").fetchall()
        for row in pending:
            self.queue.put_nowait(row["id"])
        self.worker_task = asyncio.create_task(self.run())

    async def close(self) -> None:
        if self.worker_task:
            self.worker_task.cancel()
            try:
                await self.worker_task
            except asyncio.CancelledError:
                pass
        if self.http:
            await self.http.aclose()

    async def run(self) -> None:
        while True:
            job_id = await self.queue.get()
            try:
                await self.process(job_id)
            finally:
                self.queue.task_done()

    async def process(self, job_id: str) -> None:
        with connect_db() as db:
            changed = db.execute(
                "UPDATE generations SET status='running', updated_at=? WHERE id=? AND status='queued'",
                (now_iso(), job_id),
            ).rowcount
            if not changed:
                return
            job = db.execute("SELECT * FROM generations WHERE id=?", (job_id,)).fetchone()
            ref = db.execute("SELECT * FROM references_audio WHERE id=?", (job["reference_id"],)).fetchone()
        if not ref:
            self.fail(job_id, "Selected reference audio no longer exists")
            return

        audio_bytes = (REFERENCES_DIR / ref["filename"]).read_bytes()
        request_body = {
            "text": job["text"],
            "references": [{"audio": audio_bytes, "text": ref["transcript"]}],
            "format": "mp3",
        }
        try:
            assert self.http is not None
            response = await self.http.post(
                job["tts_url"] or active_tts_url(),
                params={"format": "msgpack"},
                content=msgpack.packb(request_body, use_bin_type=True),
                headers={
                    "Authorization": f"Bearer {FISH_SPEECH_API_KEY}",
                    "Content-Type": "application/msgpack",
                },
            )
            if response.is_error:
                detail = response.text.strip()
                self.fail(job_id, f"Fish Speech returned HTTP {response.status_code}" + (f": {detail[:500]}" if detail else ""))
                return
            if not response.content:
                self.fail(job_id, "Fish Speech returned an empty audio response")
                return
            with connect_db() as db:
                filename = f"{job_id}.mp3"
                changed = db.execute(
                    "UPDATE generations SET status='completed', filename=?, error=NULL, updated_at=? WHERE id=? AND status='running'",
                    (filename, now_iso(), job_id),
                ).rowcount
                if not changed:
                    return
                (GENERATIONS_DIR / filename).write_bytes(response.content)
        except httpx.TimeoutException:
            self.fail(job_id, "Fish Speech request timed out")
        except httpx.HTTPError as exc:
            self.fail(job_id, f"Could not reach Fish Speech: {exc}")
        except Exception as exc:
            self.fail(job_id, f"Generation failed: {exc}")

    @staticmethod
    def fail(job_id: str, message: str) -> None:
        with connect_db() as db:
            db.execute(
                "UPDATE generations SET status='failed', error=?, updated_at=? WHERE id=? AND status='running'",
                (message[:1000], now_iso(), job_id),
            )


job_queue = JobQueue()


@asynccontextmanager
async def lifespan(_: FastAPI):
    await job_queue.start()
    try:
        yield
    finally:
        await job_queue.close()


app = FastAPI(title="Fish Speech Studio", lifespan=lifespan)


@app.get("/api/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/api/tts-endpoints")
def list_tts_endpoints() -> dict[str, Any]:
    current = active_tts_url()
    with connect_db() as db:
        rows = db.execute("SELECT * FROM tts_endpoints ORDER BY created_at DESC").fetchall()
    return {
        "items": [{"id": row["id"], "url": row["url"], "created_at": row["created_at"], "active": row["url"] == current} for row in rows],
        "active_url": current,
    }


@app.post("/api/tts-endpoints/test")
async def test_tts_endpoint(request: TtsUrlRequest) -> dict[str, str | bool]:
    url, message = await probe_tts_url(request.url)
    return {"ok": True, "url": url, "message": message}


@app.post("/api/tts-endpoints", status_code=201)
async def save_tts_endpoint(request: TtsUrlRequest) -> dict[str, Any]:
    url, message = await probe_tts_url(request.url)
    endpoint_id = str(uuid.uuid4())
    with connect_db() as db:
        db.execute(
            "INSERT INTO tts_endpoints(id,url,created_at) VALUES(?,?,?) ON CONFLICT(url) DO NOTHING",
            (endpoint_id, url, now_iso()),
        )
        row = db.execute("SELECT * FROM tts_endpoints WHERE url=?", (url,)).fetchone()
        db.execute(
            "INSERT INTO app_settings(key,value) VALUES('active_tts_url',?) ON CONFLICT(key) DO UPDATE SET value=excluded.value",
            (url,),
        )
    return {"id": row["id"], "url": row["url"], "created_at": row["created_at"], "active": True, "message": message}


@app.post("/api/tts-endpoints/{endpoint_id}/activate")
def activate_tts_endpoint(endpoint_id: str) -> dict[str, Any]:
    with connect_db() as db:
        row = db.execute("SELECT * FROM tts_endpoints WHERE id=?", (endpoint_id,)).fetchone()
        if not row:
            raise HTTPException(status_code=404, detail="Saved Fish Speech URL not found")
        db.execute(
            "INSERT INTO app_settings(key,value) VALUES('active_tts_url',?) ON CONFLICT(key) DO UPDATE SET value=excluded.value",
            (row["url"],),
        )
    return {"id": row["id"], "url": row["url"], "created_at": row["created_at"], "active": True}


@app.get("/api/references")
def list_references() -> list[dict[str, Any]]:
    with connect_db() as db:
        rows = db.execute("SELECT * FROM references_audio ORDER BY created_at DESC").fetchall()
    return [reference_json(row) for row in rows]


def normalize_recording(audio: bytes, suffix: str) -> bytes:
    if not shutil.which("ffmpeg"):
        raise HTTPException(status_code=503, detail="ffmpeg is required to normalize browser recordings")
    with tempfile.TemporaryDirectory(dir=DATA_DIR) as temp_dir:
        source = Path(temp_dir) / f"recording{suffix}"
        target = Path(temp_dir) / "reference.wav"
        source.write_bytes(audio)
        result = subprocess.run(
            ["ffmpeg", "-nostdin", "-y", "-i", str(source), "-ac", "1", "-ar", "44100", "-c:a", "pcm_s16le", str(target)],
            capture_output=True,
            timeout=60,
        )
        if result.returncode != 0 or not target.exists():
            raise HTTPException(status_code=400, detail="Could not decode the browser recording")
        return target.read_bytes()


@app.post("/api/references", status_code=201)
async def create_reference(
    file: UploadFile = File(...),
    transcript: str = Form(...),
    name: str = Form(""),
    recording: bool = Form(False),
) -> dict[str, Any]:
    transcript = transcript.strip()
    if not transcript:
        raise HTTPException(status_code=422, detail="Reference transcript is required")
    audio = await file.read(MAX_REFERENCE_BYTES + 1)
    if len(audio) > MAX_REFERENCE_BYTES:
        raise HTTPException(status_code=413, detail=f"Audio must be at most {MAX_REFERENCE_BYTES // (1024 * 1024)} MiB")
    if recording:
        suffix = ".webm" if "webm" in (file.content_type or "") else ".ogg"
        audio = normalize_recording(audio, suffix)
    elif not (audio.startswith(b"RIFF") and audio[8:12] == b"WAVE"):
        raise HTTPException(status_code=415, detail="Upload a WAV file")
    if len(audio) > MAX_REFERENCE_BYTES:
        raise HTTPException(status_code=413, detail="Normalized recording is too large")
    reference_id = str(uuid.uuid4())
    filename = f"{reference_id}.wav"
    (REFERENCES_DIR / filename).write_bytes(audio)
    display_name = name.strip() or Path(file.filename or "Reference").stem or "Reference"
    created_at = now_iso()
    with connect_db() as db:
        db.execute(
            "INSERT INTO references_audio(id,name,transcript,filename,created_at) VALUES(?,?,?,?,?)",
            (reference_id, display_name[:120], transcript, filename, created_at),
        )
        row = db.execute("SELECT * FROM references_audio WHERE id=?", (reference_id,)).fetchone()
    return reference_json(row)


@app.get("/api/references/{reference_id}/audio")
def reference_audio(reference_id: str) -> FileResponse:
    row = get_reference(reference_id)
    if not row:
        raise HTTPException(status_code=404, detail="Reference not found")
    return FileResponse(REFERENCES_DIR / row["filename"], media_type="audio/wav", filename=row["filename"])


@app.post("/api/generations", status_code=202)
def create_generation(request: CreateGeneration) -> dict[str, Any]:
    text_value = request.text.strip()
    if not text_value:
        raise HTTPException(status_code=422, detail="Text to generate is required")
    if not get_reference(request.reference_id):
        raise HTTPException(status_code=404, detail="Selected reference was not found")
    job_id = str(uuid.uuid4())
    timestamp = now_iso()
    tts_url = active_tts_url()
    with connect_db() as db:
        db.execute(
            "INSERT INTO generations(id,text,reference_id,tts_url,status,created_at,updated_at) VALUES(?,?,?,?,'queued',?,?)",
            (job_id, text_value, request.reference_id, tts_url, timestamp, timestamp),
        )
    job_queue.queue.put_nowait(job_id)
    return generation_json(get_generation(job_id))


@app.get("/api/generations")
def list_generations(page: int = 1, page_size: int = PAGE_SIZE) -> dict[str, Any]:
    if page < 1:
        raise HTTPException(status_code=422, detail="Page must be at least 1")
    if page_size != PAGE_SIZE:
        raise HTTPException(status_code=422, detail=f"Page size is fixed at {PAGE_SIZE}")
    offset = (page - 1) * PAGE_SIZE
    with connect_db() as db:
        total = db.execute("SELECT COUNT(*) FROM generations").fetchone()[0]
        rows = db.execute(
            "SELECT * FROM generations ORDER BY created_at DESC, id DESC LIMIT ? OFFSET ?",
            (PAGE_SIZE, offset),
        ).fetchall()
    return {"items": [generation_json(row) for row in rows], "page": page, "page_size": PAGE_SIZE, "total": total, "pages": (total + PAGE_SIZE - 1) // PAGE_SIZE}


@app.get("/api/generations/{job_id}")
def generation_status(job_id: str) -> dict[str, Any]:
    row = get_generation(job_id)
    if not row:
        raise HTTPException(status_code=404, detail="Generation not found")
    return generation_json(row)


@app.post("/api/generations/{job_id}/cancel")
def cancel_generation(job_id: str) -> dict[str, Any]:
    row = get_generation(job_id)
    if not row:
        raise HTTPException(status_code=404, detail="Generation not found")
    if row["status"] == "cancelled":
        return generation_json(row)
    if row["status"] in {"completed", "failed"}:
        raise HTTPException(status_code=409, detail=f"Cannot cancel a {row['status']} generation")
    with connect_db() as db:
        changed = db.execute(
            "UPDATE generations SET status='cancelled', updated_at=? WHERE id=? AND status IN ('queued','running')",
            (now_iso(), job_id),
        ).rowcount
    updated = get_generation(job_id)
    if not changed and updated["status"] in {"completed", "failed"}:
        raise HTTPException(status_code=409, detail=f"Cannot cancel a {updated['status']} generation")
    return generation_json(updated)


@app.get("/api/generations/{job_id}/audio")
def generated_audio(job_id: str) -> FileResponse:
    row = get_generation(job_id)
    if not row or not row["filename"]:
        raise HTTPException(status_code=404, detail="Generated audio not found")
    return FileResponse(GENERATIONS_DIR / row["filename"], media_type="audio/mpeg", filename=row["filename"])


frontend_dist = Path(__file__).resolve().parent.parent / "frontend" / "dist"
if frontend_dist.exists():
    app.mount("/", StaticFiles(directory=frontend_dist, html=True), name="frontend")