import base64
import os
import tempfile
from pathlib import Path
from typing import Any, Dict, List
from uuid import uuid4

import fastapi
import modal
from fastapi import Header, HTTPException
from fastapi.responses import Response

from backend.magpie_adapter import MagpieAdapter
from backend.omnivoice_adapter import OmniVoiceAdapter
from backend.synthesis_catalog import MODAL_BACKEND
from backend.types import VoiceConfig


app = modal.App("scriptorium-tts")
image = (
    modal.Image.debian_slim(python_version="3.12")
    .apt_install("git")
    .pip_install(
        "fastapi[standard]",
        "numpy>=1.26.0",
        "soundfile>=0.13.0",
        "torch>=2.8.0",
        "torchaudio>=2.8.0",
        "omnivoice>=0.1.5",
        "requests>=2.32.0",
        "huggingface_hub>=0.33.0",
        "nemo_toolkit[tts]@git+https://github.com/NVIDIA/NeMo.git@main",
        "kaldialign",
    )
    .add_local_python_source("backend")
)
jobs = modal.Dict.from_name("scriptorium-modal-jobs", create_if_missing=True)
artifacts = modal.Dict.from_name("scriptorium-modal-artifacts", create_if_missing=True)
web_app = fastapi.FastAPI()


def _check_auth(header: str | None) -> None:
    expected = os.getenv("SCRIPTORIUM_MODAL_SHARED_SECRET", "").strip()
    if not expected:
        return
    token = (header or "").removeprefix("Bearer ").strip()
    if token != expected:
        raise HTTPException(status_code=401, detail="Unauthorized")


def _adapter_for(model: str):
    if model == "magpie":
        return MagpieAdapter()
    return OmniVoiceAdapter()


def _voice_config(payload: Dict[str, Any], temp_dir: Path) -> VoiceConfig:
    data = dict(payload)
    sample_b64 = data.pop("sample_b64", None)
    voice = VoiceConfig.from_dict(data)
    if sample_b64:
        sample_path = temp_dir / "clone-sample.wav"
        sample_path.write_bytes(base64.b64decode(sample_b64))
        voice.sample_path = str(sample_path)
    return voice


@app.function(image=image, gpu="any", timeout=900)
def render_preview(payload: Dict[str, Any]) -> Dict[str, Any]:
    with tempfile.TemporaryDirectory() as tmp:
        temp_dir = Path(tmp)
        output_path = temp_dir / "preview.wav"
        voice = _voice_config(payload["voice_config"], temp_dir)
        result = _adapter_for(voice.model).synthesize(
            text=str(payload["text"]),
            output_path=output_path,
            voice_config=voice,
            diffusion_steps=int(payload.get("diffusion_steps", 32)),
            speed=float(payload.get("speed", 1.0)),
        )
        return {
            "audio_base64": base64.b64encode(output_path.read_bytes()).decode("ascii"),
            "duration_seconds": result.get("duration_seconds", 0),
            "sample_rate": result.get("sample_rate", 24000),
            "backend": MODAL_BACKEND,
            "model": voice.model,
        }


def _job_state(job_id: str) -> Dict[str, Any]:
    raw = jobs.get(
        job_id,
        {
            "status": "pending",
            "events": [],
            "cancel_requested": False,
        },
    )
    if raw is None:
        return {
            "status": "pending",
            "events": [],
            "cancel_requested": False,
        }
    if not isinstance(raw, dict):
        raise RuntimeError(f"Corrupt job state for {job_id}: expected dict, got {type(raw).__name__}")
    return dict(raw)


def _save_job_state(job_id: str, state: Dict[str, Any]) -> None:
    jobs[job_id] = state


def _append_event(job_id: str, event: Dict[str, Any]) -> None:
    state = _job_state(job_id)
    state.setdefault("events", []).append(event)
    state["status"] = event.get("type", state.get("status", "running"))
    _save_job_state(job_id, state)


@app.function(image=image, gpu="any", timeout=60 * 60)
def render_book(job_id: str, payload: Dict[str, Any]) -> None:
    with tempfile.TemporaryDirectory() as tmp:
        temp_dir = Path(tmp)
        voice = _voice_config(payload["voice_config"], temp_dir)
        chapters = [chapter for chapter in payload["chapters"] if chapter.get("included", True)]
        state = _job_state(job_id)
        state["status"] = "running"
        _save_job_state(job_id, state)
        _append_event(
            job_id,
            {
                "type": "started",
                "session_id": payload["session_id"],
                "total_chapters": len(chapters),
                "book_title": payload["book"].get("title"),
                "backend": MODAL_BACKEND,
                "model": voice.model,
            },
        )

        adapter = _adapter_for(voice.model)
        outputs: List[str] = []
        for index, chapter in enumerate(chapters):
            state = _job_state(job_id)
            if state.get("cancel_requested"):
                _append_event(
                    job_id,
                    {
                        "type": "cancelled",
                        "session_id": payload["session_id"],
                        "backend": MODAL_BACKEND,
                        "model": voice.model,
                    },
                )
                return

            chapter_id = str(chapter["id"])
            _append_event(
                job_id,
                {
                    "type": "chapter_started",
                    "session_id": payload["session_id"],
                    "chapter_id": chapter_id,
                    "chapter_title": chapter["title"],
                    "chapter_index": index,
                    "overall_progress": index / max(1, len(chapters)),
                    "backend": MODAL_BACKEND,
                    "model": voice.model,
                },
            )
            output_path = temp_dir / f"{index + 1:03d}-{chapter_id}.wav"
            result = adapter.synthesize(
                text=str(chapter["text"]),
                output_path=output_path,
                voice_config=voice,
                diffusion_steps=int(payload.get("diffusion_steps", 32)),
                speed=float(payload.get("speed", 1.0)),
            )
            filename = output_path.name
            artifacts[f"{job_id}:{filename}"] = output_path.read_bytes()
            outputs.append(filename)
            _append_event(
                job_id,
                {
                    "type": "chapter_done",
                    "session_id": payload["session_id"],
                    "chapter_id": chapter_id,
                    "duration_seconds": int(result.get("duration_seconds", 0)),
                    "overall_progress": (index + 1) / max(1, len(chapters)),
                    "filename": filename,
                    "artifact_url": f"/artifacts/{job_id}/{filename}",
                    "backend": MODAL_BACKEND,
                    "model": voice.model,
                },
            )

        _append_event(
            job_id,
            {
                "type": "completed",
                "session_id": payload["session_id"],
                "outputs": outputs,
                "backend": MODAL_BACKEND,
                "model": voice.model,
            },
        )


@web_app.post("/preview")
def preview_endpoint(
    payload: Dict[str, Any],
    authorization: str | None = Header(default=None),
) -> Dict[str, Any]:
    _check_auth(authorization)
    return render_preview.remote(payload)


@web_app.post("/renders")
def submit_render_endpoint(
    payload: Dict[str, Any],
    authorization: str | None = Header(default=None),
) -> Dict[str, str]:
    _check_auth(authorization)
    job_id = str(uuid4())
    _save_job_state(job_id, {"status": "pending", "events": [], "cancel_requested": False})
    render_book.spawn(job_id, payload)
    return {"job_id": job_id}


@web_app.get("/renders/{job_id}")
def render_status_endpoint(
    job_id: str,
    cursor: int = 0,
    authorization: str | None = Header(default=None),
) -> Dict[str, Any]:
    _check_auth(authorization)
    try:
        state = _job_state(job_id)
        events = list(state.get("events", []))
        return {
            "job_id": job_id,
            "status": state.get("status", "pending"),
            "events": events[cursor:],
        }
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Unable to fetch render status for {job_id}: {exc}") from exc


@web_app.post("/renders/{job_id}/cancel")
def cancel_render_endpoint(
    job_id: str,
    authorization: str | None = Header(default=None),
) -> Dict[str, str]:
    _check_auth(authorization)
    state = _job_state(job_id)
    state["cancel_requested"] = True
    state["status"] = "cancelled"
    _save_job_state(job_id, state)
    return {"type": "cancelled", "job_id": job_id}


@web_app.get("/artifacts/{job_id}/{filename}")
def artifact_endpoint(
    job_id: str,
    filename: str,
    authorization: str | None = Header(default=None),
):
    _check_auth(authorization)
    content = artifacts.get(f"{job_id}:{filename}")
    if content is None:
        raise HTTPException(status_code=404, detail="Artifact not found")
    return Response(content=content, media_type="audio/wav")


@app.function(image=image)
@modal.asgi_app()
def fastapi_app():
    return web_app