#!/usr/bin/env python3
"""Align ASR sidecars to exact EPUB text and create a Piper WAV/metadata dataset."""

from __future__ import annotations

import argparse
import csv
import difflib
import json
import re
import subprocess
import sys
import wave
from dataclasses import dataclass
from pathlib import Path


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


def normalize_word(value: str) -> str:
    return value.lower().replace("ё", "е")


@dataclass(frozen=True)
class BookWord:
    value: str
    start: int
    end: int


@dataclass(frozen=True)
class AsrToken:
    value: str
    track: str
    segment_id: int
    start: float
    end: float


@dataclass
class AlignedSegment:
    track: str
    source: Path
    start: float
    end: float
    book_start: int
    book_end: int
    matched_tokens: int
    total_tokens: int
    avg_logprob: float
    no_speech_prob: float

    @property
    def duration(self) -> float:
        return self.end - self.start

    @property
    def coverage(self) -> float:
        return self.matched_tokens / max(1, self.total_tokens)


def load_book(path: Path) -> tuple[str, list[BookWord]]:
    text = path.read_text(encoding="utf-8")
    words = [
        BookWord(normalize_word(match.group()), match.start(), match.end())
        for match in WORD_RE.finditer(text)
    ]
    return text, words


def load_asr(sidecars: list[Path]) -> tuple[list[AsrToken], list[dict[str, object]], dict[str, Path]]:
    tokens: list[AsrToken] = []
    segments: list[dict[str, object]] = []
    sources: dict[str, Path] = {}
    for sidecar in sidecars:
        payload = json.loads(sidecar.read_text(encoding="utf-8"))
        track = sidecar.stem
        source = Path(payload["source"])
        sources[track] = source
        for segment in payload["segments"]:
            token_start = len(tokens)
            for word in segment.get("words", []):
                normalized = [normalize_word(match.group()) for match in WORD_RE.finditer(word["word"])]
                for value in normalized:
                    tokens.append(
                        AsrToken(
                            value=value,
                            track=track,
                            segment_id=int(segment["id"]),
                            start=float(word["start"]),
                            end=float(word["end"]),
                        )
                    )
            segments.append(
                {
                    "track": track,
                    "source": source,
                    "segment_id": int(segment["id"]),
                    "start": float(segment["start"]),
                    "end": float(segment["end"]),
                    "avg_logprob": float(segment["avg_logprob"]),
                    "no_speech_prob": float(segment["no_speech_prob"]),
                    "token_start": token_start,
                    "token_end": len(tokens),
                }
            )
    return tokens, segments, sources


def token_mapping(asr_tokens: list[AsrToken], book_words: list[BookWord]) -> dict[int, int]:
    matcher = difflib.SequenceMatcher(
        None,
        [token.value for token in asr_tokens],
        [word.value for word in book_words],
        autojunk=True,
    )
    mapping: dict[int, int] = {}
    for block in matcher.get_matching_blocks():
        for offset in range(block.size):
            mapping[block.a + offset] = block.b + offset
    return mapping


def align_segments(
    segments: list[dict[str, object]],
    mapping: dict[int, int],
    book_words: list[BookWord],
) -> list[AlignedSegment]:
    aligned: list[AlignedSegment] = []
    previous_book_end = -1
    for segment in segments:
        start_index = int(segment["token_start"])
        end_index = int(segment["token_end"])
        mapped = [mapping[index] for index in range(start_index, end_index) if index in mapping]
        if not mapped:
            continue
        first, last = min(mapped), max(mapped)
        if first < previous_book_end:
            continue
        asr_count = max(1, end_index - start_index)
        book_count = last - first + 1
        if book_count > asr_count * 1.8 + 8:
            continue
        item = AlignedSegment(
            track=str(segment["track"]),
            source=Path(segment["source"]),
            start=float(segment["start"]),
            end=float(segment["end"]),
            book_start=book_words[first].start,
            book_end=book_words[last].end,
            matched_tokens=len(mapped),
            total_tokens=asr_count,
            avg_logprob=float(segment["avg_logprob"]),
            no_speech_prob=float(segment["no_speech_prob"]),
        )
        aligned.append(item)
        previous_book_end = last
    return aligned


def merge_segments(items: list[AlignedSegment], book_text: str, max_duration: float) -> list[AlignedSegment]:
    merged: list[AlignedSegment] = []
    current: AlignedSegment | None = None
    for item in items:
        eligible = item.coverage >= 0.55 and item.avg_logprob >= -1.2 and item.no_speech_prob <= 0.5
        if not eligible or item.duration <= 0:
            if current is not None:
                merged.append(current)
                current = None
            continue
        if current is None:
            current = item
            continue
        gap = item.start - current.end
        combined_duration = item.end - current.start
        same_track = item.track == current.track and item.source == current.source
        near_in_book = 0 <= item.book_start - current.book_end <= 120
        if same_track and gap <= 0.8 and near_in_book and combined_duration <= max_duration:
            total = current.total_tokens + item.total_tokens
            current.end = item.end
            current.book_end = item.book_end
            current.matched_tokens += item.matched_tokens
            current.avg_logprob = (
                current.avg_logprob * current.total_tokens + item.avg_logprob * item.total_tokens
            ) / total
            current.no_speech_prob = max(current.no_speech_prob, item.no_speech_prob)
            current.total_tokens = total
            text = book_text[current.book_start : current.book_end].rstrip()
            if current.duration >= 3.0 and text.endswith((".", "!", "?", "…", ":", ";")):
                merged.append(current)
                current = None
        else:
            merged.append(current)
            current = item
    if current is not None:
        merged.append(current)
    return merged


def clean_label(value: str) -> str:
    return re.sub(r"\s+", " ", value).strip(" —–-\t\r\n")


def is_valid_wav(path: Path, sample_rate: int) -> bool:
    if not path.is_file():
        return False
    try:
        with wave.open(str(path), "rb") as audio:
            return (
                audio.getnchannels() == 1
                and audio.getsampwidth() == 2
                and audio.getframerate() == sample_rate
                and audio.getnframes() > 0
            )
    except (EOFError, wave.Error):
        return False


def create_dataset(
    items: list[AlignedSegment],
    book_text: str,
    output_dir: Path,
    ffmpeg: str,
    sample_rate: int,
    min_duration: float,
    max_duration: float,
    min_coverage: float,
) -> dict[str, object]:
    wav_dir = output_dir / "wav"
    wav_dir.mkdir(parents=True, exist_ok=True)
    metadata_path = output_dir / "metadata.csv"
    rows: list[tuple[str, str]] = []
    accepted_seconds = 0.0
    rejected = 0
    details: list[dict[str, object]] = []
    for index, item in enumerate(items):
        text = clean_label(book_text[item.book_start : item.book_end])
        duration = item.duration
        if not (min_duration <= duration <= max_duration) or item.coverage < min_coverage:
            rejected += 1
            continue
        if len(text) < 8 or len(text) > 320 or "|" in text:
            rejected += 1
            continue
        clip_name = f"aletheia_ru_{len(rows):06d}.wav"
        clip_path = wav_dir / clip_name
        start = max(0.0, item.start - 0.06)
        end = item.end + 0.08
        if not is_valid_wav(clip_path, sample_rate):
            subprocess.run(
                [
                    ffmpeg,
                    "-hide_banner",
                    "-loglevel",
                    "error",
                    "-y",
                    "-ss",
                    f"{start:.3f}",
                    "-i",
                    str(item.source),
                    "-t",
                    f"{end - start:.3f}",
                    "-ac",
                    "1",
                    "-ar",
                    str(sample_rate),
                    "-sample_fmt",
                    "s16",
                    str(clip_path),
                ],
                check=True,
            )
        rows.append((clip_name, text))
        accepted_seconds += duration
        details.append(
            {
                "clip": clip_name,
                "source": str(item.source),
                "start": round(item.start, 3),
                "end": round(item.end, 3),
                "duration": round(duration, 3),
                "coverage": round(item.coverage, 4),
                "text": text,
            }
        )
    with metadata_path.open("w", encoding="utf-8", newline="") as output:
        writer = csv.writer(output, delimiter="|", lineterminator="\n")
        writer.writerows(rows)
    report = {
        "schema": 1,
        "clips": len(rows),
        "accepted_seconds": round(accepted_seconds, 3),
        "accepted_hours": round(accepted_seconds / 3600, 4),
        "rejected_candidates": rejected,
        "sample_rate": sample_rate,
        "min_duration": min_duration,
        "max_duration": max_duration,
        "min_coverage": min_coverage,
        "items": details,
    }
    (output_dir / "dataset_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    return report


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--book-text", type=Path, required=True)
    parser.add_argument("--transcripts-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--sample-rate", type=int, default=22050)
    parser.add_argument("--min-duration", type=float, default=2.0)
    parser.add_argument("--max-duration", type=float, default=12.0)
    parser.add_argument("--min-coverage", type=float, default=0.62)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    try:
        import imageio_ffmpeg
    except ImportError as error:
        raise SystemExit("imageio-ffmpeg is missing from PYTHONPATH") from error

    sidecars = sorted(args.transcripts_dir.glob("*.json"))
    if not sidecars:
        parser.error("No transcript sidecars found")
    book_text, book_words = load_book(args.book_text)
    asr_tokens, segments, _ = load_asr(sidecars)
    mapping = token_mapping(asr_tokens, book_words)
    aligned = align_segments(segments, mapping, book_words)
    merged = merge_segments(aligned, book_text, args.max_duration)
    summary = {
        "transcript_files": len(sidecars),
        "book_words": len(book_words),
        "asr_tokens": len(asr_tokens),
        "exact_token_matches": len(mapping),
        "exact_token_match_ratio": round(len(mapping) / max(1, len(asr_tokens)), 4),
        "aligned_segments": len(aligned),
        "merged_candidates": len(merged),
    }
    if args.dry_run:
        print(json.dumps(summary, ensure_ascii=False, indent=2))
        return 0
    report = create_dataset(
        merged,
        book_text,
        args.output_dir,
        imageio_ffmpeg.get_ffmpeg_exe(),
        args.sample_rate,
        args.min_duration,
        args.max_duration,
        args.min_coverage,
    )
    report.update(summary)
    (args.output_dir / "dataset_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(json.dumps({key: value for key, value in report.items() if key != "items"}, ensure_ascii=False, indent=2))
    return 0


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