from __future__ import annotations

import gc
import hashlib
import hmac
import html
import json
import os
import re
import secrets
import shutil
import sqlite3
import subprocess
import threading
import time
import uuid
import zipfile
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from pathlib import Path, PurePosixPath
from typing import TYPE_CHECKING, Annotated, Any, Iterator
from xml.etree import ElementTree

import imageio_ffmpeg
import soundfile as sf
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, UploadFile
from fastapi.responses import FileResponse
from num2words import num2words

if TYPE_CHECKING:
    from qwen_tts import Qwen3TTSModel


ROOT = Path(os.environ.get("AUDIOBOOK_SERVICE_ROOT", "/home/sevenhill/apps/aletheia-audiobook"))
SERVICE_ROOT = ROOT / "service"
DATA_ROOT = ROOT / "audiobook-jobs"
TOKEN_FILE = SERVICE_ROOT / "service-token.txt"
DATABASE = DATA_ROOT / "jobs.db3"
PERFORMANCE_LOG = DATA_ROOT / "performance.jsonl"
MODEL_ID = os.environ.get("QWEN_TTS_MODEL_ID", "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
MAX_UPLOAD_BYTES = 200 * 1024 * 1024
GPU_IDS = tuple(
    int(value.strip())
    for value in os.environ.get("QWEN_TTS_GPU_IDS", "0,1,2").split(",")
    if value.strip()
)
SPEAKER = os.environ.get("QWEN_TTS_SPEAKER", "Ryan")
VOICE_INSTRUCTION = (
    "Read in a calm, clear, natural audiobook style with steady pacing and distinct diction."
)

SERVICE_ROOT.mkdir(parents=True, exist_ok=True)
DATA_ROOT.mkdir(parents=True, exist_ok=True)
OLLAMA_PAUSE_MARKER = SERVICE_ROOT / "ollama-paused-for-audiobook"
if not TOKEN_FILE.exists():
    TOKEN_FILE.write_text(secrets.token_urlsafe(48), encoding="ascii")
API_TOKEN = TOKEN_FILE.read_text(encoding="ascii").strip()


def pause_ollama_for_audiobook() -> None:
    """Release every CMP GPU before loading the three Qwen workers."""
    subprocess.run(["systemctl", "--user", "stop", "ollama.service"], check=True)
    subprocess.run(["systemctl", "--user", "mask", "--runtime", "ollama.service"], check=True)
    OLLAMA_PAUSE_MARKER.write_text("1", encoding="ascii")


def restore_ollama_after_audiobook() -> None:
    if not OLLAMA_PAUSE_MARKER.exists():
        return
    subprocess.run(["systemctl", "--user", "unmask", "--runtime", "ollama.service"], check=True)
    subprocess.run(["systemctl", "--user", "start", "ollama.service"], check=True)
    OLLAMA_PAUSE_MARKER.unlink(missing_ok=True)


# The worker terminates after every job to release CUDA/Triton contexts. The
# next lightweight receiver process restores Ollama before accepting new work.
restore_ollama_after_audiobook()

app = FastAPI(title="Aletheia CMP Audiobook Service", version="1.1.0")
queue_condition = threading.Condition()
queued_jobs: list[str] = []
worker_state = {"activeJobId": None, "modelLoaded": False}


@contextmanager
def database() -> Iterator[sqlite3.Connection]:
    connection = sqlite3.connect(DATABASE, timeout=30)
    connection.row_factory = sqlite3.Row
    try:
        yield connection
        connection.commit()
    finally:
        connection.close()


def initialize_database() -> None:
    with database() as connection:
        connection.execute(
            """
            CREATE TABLE IF NOT EXISTS jobs (
                id TEXT PRIMARY KEY,
                title TEXT NOT NULL,
                author TEXT NOT NULL,
                source_name TEXT NOT NULL,
                input_path TEXT NOT NULL,
                output_path TEXT NOT NULL,
                status TEXT NOT NULL,
                stage TEXT NOT NULL,
                processed_characters INTEGER NOT NULL DEFAULT 0,
                total_characters INTEGER NOT NULL DEFAULT 0,
                current_chapter TEXT,
                duration_ms INTEGER NOT NULL DEFAULT 0,
                chapters_json TEXT NOT NULL DEFAULT '[]',
                error_message TEXT,
                cancel_requested INTEGER NOT NULL DEFAULT 0,
                created_at INTEGER NOT NULL,
                updated_at INTEGER NOT NULL
            )
            """
        )
        connection.execute(
            "UPDATE jobs SET status='queued', stage='queued' "
            "WHERE status IN ('extracting','synthesizing','encoding')"
        )
        rows = connection.execute(
            "SELECT id FROM jobs WHERE status='queued' ORDER BY created_at"
        ).fetchall()
    queued_jobs.extend(row["id"] for row in rows)


def require_token(authorization: Annotated[str | None, Header()] = None) -> None:
    expected = f"Bearer {API_TOKEN}"
    if not authorization or not hmac.compare_digest(authorization, expected):
        raise HTTPException(status_code=401, detail="Invalid API token")


def now_ms() -> int:
    return int(time.time() * 1000)


def get_job(job_id: str) -> sqlite3.Row:
    with database() as connection:
        row = connection.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
    if row is None:
        raise HTTPException(status_code=404, detail="Audiobook job not found")
    return row


def job_payload(row: sqlite3.Row) -> dict:
    total = row["total_characters"]
    processed = row["processed_characters"]
    progress = min(100, int(processed * 100 / total)) if total else 0
    if row["status"] == "ready":
        progress = 100
    return {
        "id": row["id"],
        "status": row["status"],
        "stage": row["stage"],
        "progress": progress,
        "processedCharacters": processed,
        "totalCharacters": total,
        "currentChapter": row["current_chapter"],
        "durationMs": row["duration_ms"],
        "chapters": json.loads(row["chapters_json"] or "[]"),
        "error": row["error_message"],
        "createdAt": row["created_at"],
        "updatedAt": row["updated_at"],
    }


@app.get("/health")
def health() -> dict:
    return {
        "status": "ok",
        "workerMode": "lazy-model",
        "activeJobId": worker_state["activeJobId"],
        "modelLoaded": worker_state["modelLoaded"],
    }


@app.post("/v1/audiobooks", dependencies=[Depends(require_token)])
async def create_audiobook(
    book: Annotated[UploadFile, File()],
    title: Annotated[str, Form()],
    author: Annotated[str, Form()] = "",
) -> dict:
    extension = Path(book.filename or "book.epub").suffix.lower()
    if extension not in {".epub", ".fb2"}:
        raise HTTPException(status_code=400, detail="Only EPUB and FB2 are supported")
    job_id = uuid.uuid4().hex
    job_root = DATA_ROOT / job_id
    job_root.mkdir(parents=True)
    input_path = job_root / f"source{extension}"
    received = 0
    try:
        with input_path.open("wb") as output:
            while chunk := await book.read(1024 * 1024):
                received += len(chunk)
                if received > MAX_UPLOAD_BYTES:
                    raise HTTPException(status_code=413, detail="Book file is too large")
                output.write(chunk)
    except Exception:
        shutil.rmtree(job_root, ignore_errors=True)
        raise
    timestamp = now_ms()
    with database() as connection:
        connection.execute(
            """
            INSERT INTO jobs (
                id,title,author,source_name,input_path,output_path,status,stage,created_at,updated_at
            ) VALUES (?,?,?,?,?,?,?,?,?,?)
            """,
            (
                job_id,
                title.strip() or Path(book.filename or "Книга").stem,
                author.strip(),
                book.filename or input_path.name,
                str(input_path),
                str(job_root / "audiobook.m4a"),
                "queued",
                "queued",
                timestamp,
                timestamp,
            ),
        )
    with queue_condition:
        queued_jobs.append(job_id)
        queue_condition.notify()
    return job_payload(get_job(job_id))


@app.get("/v1/audiobooks/{job_id}", dependencies=[Depends(require_token)])
def audiobook_status(job_id: str) -> dict:
    return job_payload(get_job(job_id))


@app.get("/v1/audiobooks/{job_id}/file", dependencies=[Depends(require_token)])
def audiobook_file(job_id: str) -> FileResponse:
    row = get_job(job_id)
    output = Path(row["output_path"])
    if row["status"] != "ready" or not output.is_file():
        raise HTTPException(status_code=409, detail="Audiobook is not ready")
    safe_title = re.sub(r"[^0-9A-Za-zА-Яа-яЁё._ -]+", "_", row["title"]).strip() or "audiobook"
    return FileResponse(output, media_type="audio/mp4", filename=f"{safe_title}.m4a")


@app.delete("/v1/audiobooks/{job_id}", dependencies=[Depends(require_token)])
def cancel_or_delete(job_id: str) -> dict:
    row = get_job(job_id)
    if row["status"] == "queued":
        with queue_condition:
            queued_jobs[:] = [queued_id for queued_id in queued_jobs if queued_id != job_id]
        delete_job(job_id, Path(row["input_path"]).parent)
        return {"status": "deleted"}
    if row["status"] in {"extracting", "synthesizing", "encoding"}:
        with database() as connection:
            connection.execute(
                "UPDATE jobs SET cancel_requested=1, stage='cancelling', updated_at=? WHERE id=?",
                (now_ms(), job_id),
            )
        return {"status": "cancelling"}
    delete_job(job_id, Path(row["input_path"]).parent)
    return {"status": "deleted"}


def delete_job(job_id: str, job_root: Path) -> None:
    shutil.rmtree(job_root, ignore_errors=True)
    with database() as connection:
        connection.execute("DELETE FROM jobs WHERE id=?", (job_id,))


def finish_cancellation(job_id: str, job_root: Path) -> None:
    update_job(job_id, status="cancelled", stage="cancelled")
    delete_job(job_id, job_root)


def local_name(tag: str) -> str:
    return tag.rsplit("}", 1)[-1].lower()


def clean_text(value: str) -> str:
    value = html.unescape(re.sub(r"<[^>]+>", " ", value))
    value = value.replace("\u00a0", " ")
    return re.sub(r"\s+", " ", value).strip()


def extract_epub(path: Path) -> list[tuple[str, str]]:
    sections: list[tuple[str, str]] = []
    with zipfile.ZipFile(path) as archive:
        container = ElementTree.fromstring(archive.read("META-INF/container.xml"))
        rootfile = next(node for node in container.iter() if local_name(node.tag) == "rootfile")
        opf_path = rootfile.attrib["full-path"]
        opf = ElementTree.fromstring(archive.read(opf_path))
        manifest = {
            node.attrib["id"]: node.attrib["href"]
            for node in opf.iter()
            if local_name(node.tag) == "item" and "id" in node.attrib and "href" in node.attrib
        }
        spine = [
            node.attrib["idref"]
            for node in opf.iter()
            if local_name(node.tag) == "itemref" and "idref" in node.attrib
        ]
        base = PurePosixPath(opf_path).parent
        for index, item_id in enumerate(spine, start=1):
            href = manifest.get(item_id)
            if not href:
                continue
            entry = str(base / href.split("#", 1)[0])
            raw = archive.read(entry).decode("utf-8", errors="replace")
            heading_match = re.search(r"<h[1-3]\b[^>]*>(.*?)</h[1-3]>", raw, re.I | re.S)
            title = clean_text(heading_match.group(1)) if heading_match else f"Раздел {index}"
            paragraphs = [clean_text(match) for match in re.findall(r"<p\b[^>]*>(.*?)</p>", raw, re.I | re.S)]
            text = "\n".join(paragraph for paragraph in paragraphs if paragraph)
            if text:
                sections.append((title, text))
    return sections


def element_text_without_nested_sections(element: ElementTree.Element) -> str:
    parts: list[str] = []
    if element.text:
        parts.append(element.text)
    for child in element:
        if local_name(child.tag) != "section":
            parts.append(" ".join(child.itertext()))
        if child.tail:
            parts.append(child.tail)
    return clean_text(" ".join(parts))


def extract_fb2(path: Path) -> list[tuple[str, str]]:
    root = ElementTree.parse(path).getroot()
    sections: list[tuple[str, str]] = []
    for section in root.iter():
        if local_name(section.tag) != "section":
            continue
        title_element = next((child for child in section if local_name(child.tag) == "title"), None)
        title = clean_text(" ".join(title_element.itertext())) if title_element is not None else ""
        text = element_text_without_nested_sections(section)
        if text:
            sections.append((title or f"Раздел {len(sections) + 1}", text))
    return sections


def normalize_numbers(text: str) -> str:
    def replace(match: re.Match[str]) -> str:
        value = match.group(0)
        try:
            return num2words(int(value), lang="ru")
        except (ValueError, OverflowError):
            return value

    return re.sub(r"(?<![\w+])\d{1,12}(?!\w)", replace, text)


def split_chunks(sections: list[tuple[str, str]], limit: int = 520) -> list[dict]:
    chunks: list[dict] = []
    for chapter, text in sections:
        sentences = re.split(r"(?<=[.!?…])\s+|\n+", text)
        current = ""
        for sentence in sentences:
            sentence = normalize_numbers(sentence.strip())
            if not sentence or not re.search(r"[А-Яа-яЁё]", sentence):
                continue
            if len(sentence) > limit:
                pieces = re.split(r"(?<=[,;:])\s+", sentence)
            else:
                pieces = [sentence]
            for piece in pieces:
                if current and len(current) + 1 + len(piece) > limit:
                    chunks.append({"chapter": chapter, "text": current})
                    current = piece
                else:
                    current = f"{current} {piece}".strip()
        if current:
            chunks.append({"chapter": chapter, "text": current})
    return chunks


def update_job(job_id: str, **values: object) -> None:
    values["updated_at"] = now_ms()
    assignments = ",".join(f"{key}=?" for key in values)
    with database() as connection:
        connection.execute(
            f"UPDATE jobs SET {assignments} WHERE id=?",
            (*values.values(), job_id),
        )


def cancellation_requested(job_id: str) -> bool:
    with database() as connection:
        row = connection.execute("SELECT cancel_requested FROM jobs WHERE id=?", (job_id,)).fetchone()
    return row is None or bool(row["cancel_requested"])


def load_model(device_id: int) -> "Qwen3TTSModel":
    import torch
    from qwen_tts import Qwen3TTSModel

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is unavailable")
    torch.set_float32_matmul_precision("high")
    model = Qwen3TTSModel.from_pretrained(
        MODEL_ID,
        device_map=f"cuda:{device_id}",
        dtype=torch.bfloat16,
        attn_implementation="sdpa",
    )
    return model


def load_models() -> list["Qwen3TTSModel"]:
    if not GPU_IDS:
        raise RuntimeError("No Qwen GPU IDs are configured")
    models = [load_model(device_id) for device_id in GPU_IDS]
    worker_state["modelLoaded"] = True
    return models


def generate_custom_voice_with_metrics(
    model: Any,
    batch_chunks: list[dict],
    device_id: int,
) -> tuple[list[Any], int]:
    """Measure Qwen talker and codec stages without changing generated audio."""
    import torch

    timings: dict[str, float] = {}
    original_generate = model.model.generate
    speech_tokenizer = model.model.speech_tokenizer
    original_decode = speech_tokenizer.decode

    def synchronize() -> None:
        if torch.cuda.is_available():
            torch.cuda.synchronize(device_id)

    def timed_generate(*args: Any, **kwargs: Any) -> Any:
        synchronize()
        started = time.perf_counter()
        try:
            return original_generate(*args, **kwargs)
        finally:
            synchronize()
            timings["talkerGenerateSeconds"] = time.perf_counter() - started

    def timed_decode(*args: Any, **kwargs: Any) -> Any:
        synchronize()
        started = time.perf_counter()
        try:
            return original_decode(*args, **kwargs)
        finally:
            synchronize()
            timings["codecDecodeSeconds"] = time.perf_counter() - started

    model.model.generate = timed_generate
    speech_tokenizer.decode = timed_decode
    started = time.perf_counter()
    try:
        wavs, sample_rate = model.generate_custom_voice(
            text=[chunk["text"] for chunk in batch_chunks],
            language=["Russian"] * len(batch_chunks),
            speaker=[SPEAKER] * len(batch_chunks),
            instruct=[VOICE_INSTRUCTION] * len(batch_chunks),
        )
        synchronize()
    finally:
        model.model.generate = original_generate
        speech_tokenizer.decode = original_decode

    total_seconds = time.perf_counter() - started
    talker_seconds = timings.get("talkerGenerateSeconds", 0.0)
    codec_seconds = timings.get("codecDecodeSeconds", 0.0)
    record = {
        "timestampMs": now_ms(),
        "batchSize": len(batch_chunks),
        "characters": sum(len(chunk["text"]) for chunk in batch_chunks),
        "totalSeconds": round(total_seconds, 4),
        "talkerGenerateSeconds": round(talker_seconds, 4),
        "codecDecodeSeconds": round(codec_seconds, 4),
        "wrapperSeconds": round(max(0.0, total_seconds - talker_seconds - codec_seconds), 4),
        "cudaPeakAllocatedMiB": round(
            torch.cuda.max_memory_allocated(device_id) / (1024 * 1024), 1
        ),
    }
    with PERFORMANCE_LOG.open("a", encoding="utf-8") as output:
        output.write(json.dumps(record, ensure_ascii=False) + "\n")
    return wavs, sample_rate


def unload_models(models: list[Any]) -> None:
    if models:
        models.clear()
        import torch

        gc.collect()
        if torch.cuda.is_available():
            for device_id in GPU_IDS:
                torch.cuda.empty_cache()
    else:
        gc.collect()
    worker_state["modelLoaded"] = False


def encode_m4a(chunk_files: list[Path], output: Path, job_root: Path) -> None:
    concat = job_root / "concat.txt"
    concat.write_text(
        "\n".join(f"file '{path.as_posix()}'" for path in chunk_files),
        encoding="utf-8",
    )
    command = [
        imageio_ffmpeg.get_ffmpeg_exe(),
        "-y",
        "-f",
        "concat",
        "-safe",
        "0",
        "-i",
        str(concat),
        "-vn",
        "-c:a",
        "aac",
        "-b:a",
        "96k",
        "-movflags",
        "+faststart",
        str(output),
    ]
    completed = subprocess.run(command, capture_output=True, text=True, encoding="utf-8", errors="replace")
    if completed.returncode != 0:
        raise RuntimeError(completed.stderr[-2000:] or "FFmpeg failed")


def run_job(job_id: str) -> None:
    row = get_job(job_id)
    job_root = Path(row["input_path"]).parent
    chunks_root = job_root / "chunks"
    chunks_root.mkdir(exist_ok=True)
    models: list[Any] = []
    try:
        if cancellation_requested(job_id):
            finish_cancellation(job_id, job_root)
            return
        update_job(job_id, status="extracting", stage="extracting", error_message=None)
        input_path = Path(row["input_path"])
        sections = extract_epub(input_path) if input_path.suffix.lower() == ".epub" else extract_fb2(input_path)
        chunks = split_chunks(sections)
        if not chunks:
            raise RuntimeError("В книге не найден русский текст для озвучивания")
        total = sum(len(chunk["text"]) for chunk in chunks)
        manifest_path = job_root / "chunks.json"
        manifest_path.write_text(json.dumps(chunks, ensure_ascii=False, indent=2), encoding="utf-8")
        update_job(job_id, status="synthesizing", stage="loading_model", total_characters=total)
        if cancellation_requested(job_id):
            finish_cancellation(job_id, job_root)
            return
        pause_ollama_for_audiobook()
        models = load_models()
        import torch

        processed = 0
        duration_ms = 0
        chapter_markers: list[dict] = []
        last_chapter: str | None = None
        chunk_files: list[Path] = []
        index = 0
        while index < len(chunks):
            if cancellation_requested(job_id):
                finish_cancellation(job_id, job_root)
                return
            output = chunks_root / f"{index:06d}.wav"
            completed_indices: list[int]
            if output.is_file():
                completed_indices = [index]
            else:
                batch_indices: list[int] = []
                cursor = index
                while cursor < len(chunks) and len(batch_indices) < len(models):
                    candidate = chunks_root / f"{cursor:06d}.wav"
                    if candidate.is_file():
                        break
                    batch_indices.append(cursor)
                    cursor += 1
                with ThreadPoolExecutor(max_workers=len(batch_indices)) as executor:
                    futures = [
                        executor.submit(
                            generate_custom_voice_with_metrics,
                            model,
                            [chunks[batch_index]],
                            device_id,
                        )
                        for batch_index, model, device_id in zip(
                            batch_indices,
                            models[: len(batch_indices)],
                            GPU_IDS[: len(batch_indices)],
                            strict=True,
                        )
                    ]
                    rendered = [future.result() for future in futures]
                if cancellation_requested(job_id):
                    finish_cancellation(job_id, job_root)
                    return
                for batch_index, (wavs, sample_rate) in zip(batch_indices, rendered, strict=True):
                    sf.write(
                        chunks_root / f"{batch_index:06d}.wav",
                        wavs[0],
                        sample_rate,
                        subtype="PCM_16",
                    )
                completed_indices = batch_indices

            for completed_index in completed_indices:
                completed_chunk = chunks[completed_index]
                completed_output = chunks_root / f"{completed_index:06d}.wav"
                if completed_chunk["chapter"] != last_chapter:
                    chapter_markers.append({"title": completed_chunk["chapter"], "startMs": duration_ms})
                    last_chapter = completed_chunk["chapter"]
                info = sf.info(completed_output)
                duration_ms += round(info.frames * 1000 / info.samplerate)
                chunk_files.append(completed_output)
                processed += len(completed_chunk["text"])
                update_job(
                    job_id,
                    status="synthesizing",
                    stage="synthesizing",
                    processed_characters=processed,
                    current_chapter=completed_chunk["chapter"],
                    duration_ms=duration_ms,
                    chapters_json=json.dumps(chapter_markers, ensure_ascii=False),
                )
            index = completed_indices[-1] + 1

        unload_models(models)
        models = []
        if cancellation_requested(job_id):
            finish_cancellation(job_id, job_root)
            return
        update_job(job_id, status="encoding", stage="encoding", current_chapter=None)
        output = Path(row["output_path"])
        encode_m4a(chunk_files, output, job_root)
        if cancellation_requested(job_id):
            finish_cancellation(job_id, job_root)
            return
        if not output.is_file() or output.stat().st_size == 0:
            raise RuntimeError("Итоговый M4A-файл не создан")
        update_job(
            job_id,
            status="ready",
            stage="ready",
            processed_characters=total,
            current_chapter=None,
            chapters_json=json.dumps(chapter_markers, ensure_ascii=False),
        )
    except Exception as error:
        update_job(
            job_id,
            status="failed",
            stage="failed",
            error_message=f"{type(error).__name__}: {error}",
        )
    finally:
        unload_models(models)


def worker_loop() -> None:
    while True:
        with queue_condition:
            while not queued_jobs:
                queue_condition.wait()
            job_id = queued_jobs.pop(0)
        worker_state["activeJobId"] = job_id
        try:
            try:
                run_job(job_id)
            except HTTPException as error:
                if error.status_code != 404:
                    raise
        finally:
            worker_state["activeJobId"] = None
            # Qwen/Triton can retain a CUDA context after Python references are
            # released. Let systemd start a clean idle receiver so CMP VRAM is
            # fully available between audiobook requests. Unfinished jobs are
            # restored from SQLite by initialize_database() after restart.
            os._exit(0)


initialize_database()
threading.Thread(target=worker_loop, name="audiobook-worker", daemon=True).start()