from __future__ import annotations

import sys
from device_auth import DeviceStore
import html
import json
import os
import re
import shutil
import sqlite3
import subprocess
import threading
import time
import uuid
import zipfile
from contextlib import contextmanager
from pathlib import Path, PurePosixPath
from typing import TYPE_CHECKING, Annotated, Any, Iterator
from xml.etree import ElementTree

from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, UploadFile
from fastapi.responses import FileResponse
from num2words import num2words

ROOT = Path(os.environ.get("AUDIOBOOK_SERVICE_ROOT", "/home/sevenhill/apps/aletheia-audiobook"))
SERVICE_ROOT = ROOT / "service"
DATA_ROOT = ROOT / "audiobook-jobs"
DATABASE = DATA_ROOT / "jobs.db3"
MAX_UPLOAD_BYTES = 200 * 1024 * 1024
SERVICE_ROOT.mkdir(parents=True, exist_ok=True)
DATA_ROOT.mkdir(parents=True, exist_ok=True)

app = FastAPI(title="Aletheia CMP Audiobook Service", version="2.0.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
            )
            """
        )
        if "owner_id" not in {row[1] for row in connection.execute("PRAGMA table_info(jobs)")}:
            connection.execute("ALTER TABLE jobs ADD COLUMN owner_id TEXT")
        connection.execute("CREATE INDEX IF NOT EXISTS idx_jobs_owner ON jobs(owner_id)")
        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)


devices = DeviceStore(DATABASE)


def require_token(authorization: Annotated[str | None, Header()] = None) -> str:
    try:
        if not authorization or not authorization.startswith("Bearer "):
            raise PermissionError()
        return devices.authenticate(authorization[7:])
    except PermissionError:
        raise HTTPException(status_code=401, detail="Invalid device credential")


@app.post("/v1/devices/enroll")
def enroll_device(payload: dict) -> dict:
    try:
        return devices.enroll(payload.get("code", ""), payload.get("label", "Android"))
    except PermissionError:
        raise HTTPException(status_code=403, detail="Invalid or expired pairing code")


def owned_job(job_id: str, owner: str) -> sqlite3.Row:
    row = get_job(job_id)
    if row["owner_id"] != owner:
        raise HTTPException(status_code=404, detail="Audiobook job not found")
    return row


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:
    from voicestudio_client import check_available
    try:
        check_available()
    except Exception:
        raise HTTPException(status_code=503, detail="VoiceStudio unavailable")
    return {
        "status": "ok",
        "workerMode": "voicestudio",
        "modelLoaded": worker_state["modelLoaded"],
    }


@app.post("/v1/audiobooks")
async def create_audiobook(
    book: Annotated[UploadFile, File()],
    owner: Annotated[str, Depends(require_token)],
    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,owner_id
            ) 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,
                owner,
            ),
        )
    with queue_condition:
        queued_jobs.append(job_id)
        queue_condition.notify()
    return job_payload(get_job(job_id))


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


@app.get("/v1/audiobooks/{job_id}/file")
def audiobook_file(job_id: str, owner: Annotated[str, Depends(require_token)]) -> FileResponse:
    row = owned_job(job_id, owner)
    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}")
def cancel_or_delete(job_id: str, owner: Annotated[str, Depends(require_token)]) -> dict:
    row = owned_job(job_id, owner)
    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 run_job(job_id: str) -> None:
    from voicestudio_client import render_audiobook, split_sections
    row = get_job(job_id)
    job_root = Path(row["input_path"]).parent
    try:
        if cancellation_requested(job_id):
            finish_cancellation(job_id, job_root)
            return
        update_job(job_id, status="extracting", stage="extracting", error_message=None)
        source = Path(row["input_path"])
        sections = extract_epub(source) if source.suffix.lower() == ".epub" else extract_fb2(source)
        sections = split_sections(sections)
        if not sections:
            raise RuntimeError("В книге не найден текст для озвучивания")
        total = sum(len(text) for _, text in sections)
        update_job(job_id, status="synthesizing", stage="synthesizing", total_characters=total,
                   processed_characters=0, duration_ms=0, chapters_json="[]",
                   current_chapter="Проверка сохранённых фрагментов и озвучивание недостающих…")
        completed = {}
        verified = set()

        def on_event(event):
            if event.get("type") == "chapter":
                index = int(event["index"])
                completed[index] = event
                duration = 0
                markers = []
                for _, chapter in sorted(completed.items()):
                    markers.append({"title": chapter.get("title", "Раздел"), "startMs": duration})
                    duration += round(float(chapter.get("duration_s", 0)) * 1000)
                processed = sum(len(sections[i][1]) for i in completed if 0 <= i < len(sections))
                label = event.get("title")
                if event.get("stage") != "gpu_render":
                    verified.add(index)
                    label = f"Проверка готовых фрагментов: {len(verified)} из {len(sections)}"
                update_job(job_id, processed_characters=processed,
                           current_chapter=label, duration_ms=duration,
                           chapters_json=json.dumps(markers, ensure_ascii=False))
            elif event.get("type") == "started":
                update_job(job_id, current_chapter="Проверка готовых фрагментов перед сборкой M4A…")
            elif event.get("type") == "assembling":
                update_job(job_id, status="encoding", stage="encoding", current_chapter=None)

        result = render_audiobook(sections, row["title"], row["author"], Path(row["output_path"]),
                                 on_event, lambda: cancellation_requested(job_id))
        if cancellation_requested(job_id):
            finish_cancellation(job_id, job_root)
            return
        update_job(job_id, status="ready", stage="ready", processed_characters=total,
                   duration_ms=round(float(result["duration_s"]) * 1000), current_chapter=None)
    except Exception as error:
        if cancellation_requested(job_id):
            finish_cancellation(job_id, job_root)
        else:
            update_job(job_id, status="failed", stage="failed", error_message=f"{type(error).__name__}: {error}")

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:
            env = dict(os.environ, AUDIOBOOK_WORKER_PROCESS="1")
            process = subprocess.Popen([sys.executable, str(Path(__file__).with_name("worker_runner.py")), job_id], env=env)
            while process.poll() is None:
                try:
                    worker_state["modelLoaded"] = get_job(job_id)["status"] == "synthesizing"
                except HTTPException:
                    pass
                time.sleep(0.5)
            if process.returncode:
                try:
                    update_job(job_id, status="failed", stage="failed", error_message="Audiobook worker stopped unexpectedly")
                except HTTPException:
                    pass
        finally:
            worker_state["activeJobId"] = None
            worker_state["modelLoaded"] = False


if os.environ.get("AUDIOBOOK_WORKER_PROCESS") != "1":
    initialize_database()
    if os.environ.get("AUDIOBOOK_DISABLE_WORKER") != "1":
        threading.Thread(target=worker_loop, name="audiobook-worker", daemon=True).start()