#!/usr/bin/env python3
"""Audit an audiobook/EPUB pair before building a single-speaker TTS corpus."""

from __future__ import annotations

import argparse
import collections
import difflib
import json
import re
import subprocess
import sys
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path, PurePosixPath


WORD_RE = re.compile(r"[а-яёa-z0-9]+", re.IGNORECASE)
WHITESPACE_RE = re.compile(r"\s+")


def normalized_words(text: str) -> list[str]:
    return [word.replace("ё", "е") for word in WORD_RE.findall(text.lower())]


def element_text(root: ET.Element) -> str:
    ignored = {"script", "style", "svg", "math"}
    parts: list[str] = []

    def visit(node: ET.Element) -> None:
        tag = node.tag.rsplit("}", 1)[-1].lower()
        if tag in ignored:
            return
        if node.text:
            parts.append(node.text)
        for child in node:
            visit(child)
            if child.tail:
                parts.append(child.tail)
        if tag in {"p", "div", "section", "h1", "h2", "h3", "h4", "h5", "h6", "li", "br"}:
            parts.append("\n")

    visit(root)
    lines = [WHITESPACE_RE.sub(" ", line).strip() for line in "".join(parts).splitlines()]
    return "\n".join(line for line in lines if line)


def extract_epub(epub_path: Path, output_dir: Path) -> tuple[str, list[dict[str, object]]]:
    with zipfile.ZipFile(epub_path) as archive:
        container = ET.fromstring(archive.read("META-INF/container.xml"))
        rootfile = next(
            node.attrib["full-path"]
            for node in container.iter()
            if node.tag.rsplit("}", 1)[-1] == "rootfile"
        )
        opf = ET.fromstring(archive.read(rootfile))
        opf_dir = PurePosixPath(rootfile).parent
        manifest = {
            node.attrib["id"]: node.attrib["href"]
            for node in opf.iter()
            if node.tag.rsplit("}", 1)[-1] == "item" and "id" in node.attrib and "href" in node.attrib
        }
        spine_ids = [
            node.attrib["idref"]
            for node in opf.iter()
            if node.tag.rsplit("}", 1)[-1] == "itemref" and "idref" in node.attrib
        ]
        sections: list[dict[str, object]] = []
        texts: list[str] = []
        for index, item_id in enumerate(spine_ids):
            href = manifest.get(item_id)
            if not href:
                continue
            member = str(opf_dir / PurePosixPath(href))
            try:
                root = ET.fromstring(archive.read(member))
            except (KeyError, ET.ParseError):
                continue
            text = element_text(root)
            if not text:
                continue
            headings = [
                WHITESPACE_RE.sub(" ", "".join(node.itertext())).strip()
                for node in root.iter()
                if node.tag.rsplit("}", 1)[-1].lower() in {"h1", "h2", "h3"}
            ]
            sections.append(
                {
                    "index": index,
                    "href": member,
                    "characters": len(text),
                    "words": len(normalized_words(text)),
                    "headings": [heading for heading in headings if heading],
                }
            )
            texts.append(text)

    book_text = "\n\n".join(texts)
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "book.txt").write_text(book_text, encoding="utf-8")
    (output_dir / "spine.json").write_text(
        json.dumps(sections, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    return book_text, sections


def best_text_match(transcript: str, book_words: list[str]) -> dict[str, object]:
    spoken = normalized_words(transcript)
    if not spoken or not book_words:
        return {"ratio": 0.0, "book_excerpt": "", "word_offset": None}
    positions: dict[str, list[int]] = collections.defaultdict(list)
    for index, word in enumerate(book_words):
        if len(word) >= 4:
            positions[word].append(index)
    votes: collections.Counter[int] = collections.Counter()
    for spoken_index, word in enumerate(spoken):
        candidates = positions.get(word, ())
        if len(candidates) <= 200:
            votes.update(book_index - spoken_index for book_index in candidates)
    offsets = [offset for offset, _ in votes.most_common(30)] or [0]
    best_ratio = 0.0
    best_offset = 0
    best_window: list[str] = []
    for offset in offsets:
        start = max(0, offset - 8)
        window = book_words[start : start + len(spoken) + 16]
        ratio = difflib.SequenceMatcher(None, spoken, window, autojunk=False).ratio()
        if ratio > best_ratio:
            best_ratio, best_offset, best_window = ratio, start, window
    return {
        "ratio": round(best_ratio, 4),
        "book_excerpt": " ".join(best_window),
        "word_offset": best_offset,
    }


def transcribe_samples(
    audio_files: list[Path],
    book_text: str,
    output_dir: Path,
    model_name: str,
    sample_seconds: int,
) -> list[dict[str, object]]:
    try:
        import av
        import imageio_ffmpeg
        from faster_whisper import WhisperModel
    except ImportError as error:
        raise SystemExit(
            "Install audit dependencies into PYTHONPATH: faster-whisper imageio-ffmpeg"
        ) from error

    chosen = [audio_files[0], audio_files[len(audio_files) // 2], audio_files[-1]]
    samples_dir = output_dir / "samples"
    samples_dir.mkdir(parents=True, exist_ok=True)
    ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
    model = WhisperModel(
        model_name,
        device="cpu",
        compute_type="int8",
        download_root=str(output_dir / "models"),
    )
    book_words = normalized_words(book_text)
    results: list[dict[str, object]] = []
    for source in chosen:
        with av.open(str(source)) as media:
            stream = media.streams.audio[0]
            duration = float(stream.duration * stream.time_base) if stream.duration is not None else 0.0
        start = min(max(45.0, duration * 0.35), max(0.0, duration - sample_seconds - 5.0))
        wav_path = samples_dir / f"{source.stem[:4]}-{int(start):05d}.wav"
        subprocess.run(
            [
                ffmpeg,
                "-hide_banner",
                "-loglevel",
                "error",
                "-y",
                "-ss",
                f"{start:.3f}",
                "-i",
                str(source),
                "-t",
                str(sample_seconds),
                "-ac",
                "1",
                "-ar",
                "16000",
                str(wav_path),
            ],
            check=True,
        )
        segments, info = model.transcribe(
            str(wav_path),
            language="ru",
            beam_size=5,
            vad_filter=True,
            condition_on_previous_text=True,
        )
        transcript = " ".join(segment.text.strip() for segment in segments).strip()
        result = {
            "source": source.name,
            "source_duration_seconds": round(duration, 3),
            "sample_start_seconds": round(start, 3),
            "sample_duration_seconds": sample_seconds,
            "detected_language": info.language,
            "language_probability": round(info.language_probability, 4),
            "transcript": transcript,
        }
        result.update(best_text_match(transcript, book_words))
        results.append(result)
    (output_dir / "alignment_samples.json").write_text(
        json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    return results


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-dir", type=Path, required=True)
    parser.add_argument("--epub", type=Path)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--model", default="base")
    parser.add_argument("--sample-seconds", type=int, default=75)
    parser.add_argument("--skip-asr", action="store_true")
    args = parser.parse_args()

    audio_files = sorted(args.source_dir.glob("*.mp3"))
    epub_path = args.epub or next(args.source_dir.glob("*.epub"), None)
    if not audio_files:
        parser.error("No MP3 files found")
    if epub_path is None or not epub_path.is_file():
        parser.error("EPUB file not found")

    book_text, sections = extract_epub(epub_path, args.output_dir)
    report: dict[str, object] = {
        "source_dir": str(args.source_dir.resolve()),
        "epub": str(epub_path.resolve()),
        "audio_files": [str(path.resolve()) for path in audio_files],
        "audio_file_count": len(audio_files),
        "epub_section_count": len(sections),
        "book_characters": len(book_text),
        "book_words": len(normalized_words(book_text)),
    }
    if not args.skip_asr:
        report["alignment_samples"] = transcribe_samples(
            audio_files, book_text, args.output_dir, args.model, args.sample_seconds
        )
    (args.output_dir / "audit.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())
