#!/usr/bin/env python3
"""Find chapter boundaries in a long transcript, and title them from what is said.

An hour of talking has structure, and that structure is visible in the transcript
without understanding the content: a speaker who finishes a topic pauses longer
than usual and then starts using different words. So a boundary needs both signals:

  * a pause that is long *for this speaker* — the threshold is a percentile of their
    own gaps, not a fixed number, because a fast speaker's "long pause" is 0.6 s and
    a deliberate one's is 2 s;
  * a lexical shift — the vocabulary before and after the pause has to actually
    differ, measured as cosine distance between word bags. A breath between two
    sentences about the same thing is not a chapter.

Titles come from the first sentence after the boundary, trimmed to something that
fits a YouTube description. They are a starting point for a human to rewrite, and
the report says so rather than pretending they are editorial.
"""
from __future__ import annotations

import argparse
import json
import math
import re
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from skill_config import emit, utf8_stdout  # noqa: E402

# Words that carry no topic information in Russian; they would dominate every bag.
STOPWORDS = {
    "и", "в", "во", "не", "что", "он", "на", "я", "с", "со", "как", "а", "то", "все",
    "она", "так", "его", "но", "да", "ты", "к", "у", "же", "вы", "за", "бы", "по",
    "только", "ее", "мне", "было", "вот", "от", "меня", "еще", "нет", "о", "из",
    "ему", "теперь", "когда", "даже", "ну", "вдруг", "ли", "если", "уже", "или",
    "ни", "быть", "был", "него", "до", "вас", "нибудь", "опять", "уж", "вам", "ведь",
    "там", "потом", "себя", "ничего", "ей", "может", "они", "тут", "где", "есть",
    "надо", "ней", "для", "мы", "тебя", "их", "чем", "была", "сам", "чтоб", "без",
    "будто", "чего", "раз", "тоже", "себе", "под", "будет", "ж", "тогда", "кто",
    "этот", "того", "потому", "этого", "какой", "совсем", "ним", "здесь", "этом",
    "один", "почти", "мой", "тем", "чтобы", "нее", "были", "куда", "зачем", "всех",
    "никогда", "можно", "при", "наконец", "два", "об", "другой", "хоть", "после",
    "над", "больше", "тот", "через", "эти", "нас", "про", "всего", "них", "какая",
    "много", "разве", "три", "эту", "моя", "впрочем", "хорошо", "свою", "этой",
    "перед", "иногда", "лучше", "чуть", "том", "нельзя", "такой", "им", "более",
    "всегда", "конечно", "всю", "между", "это",
}

SENTENCE_END = re.compile(r"[.!?…]+$")
TOKEN = re.compile(r"[^\W\d_]+", re.UNICODE)


def load_words(path: Path) -> list[dict]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    words = payload.get("words") if isinstance(payload, dict) else payload
    cleaned = []
    for index, word in enumerate(words or []):
        token = str(word.get("word", "")).strip()
        if not token:
            continue
        try:
            start = float(word["start"])
            end = float(word["end"])
        except (KeyError, TypeError, ValueError):
            continue
        cleaned.append({"id": word.get("id", index), "word": token,
                        "start": start, "end": max(end, start + 0.01)})
    cleaned.sort(key=lambda item: item["start"])
    return cleaned


def bag(words: list[dict]) -> Counter:
    counter: Counter = Counter()
    for word in words:
        for match in TOKEN.findall(word["word"].casefold()):
            if len(match) > 2 and match not in STOPWORDS:
                counter[match] += 1
    return counter


def cosine_distance(first: Counter, second: Counter) -> float:
    if not first or not second:
        return 1.0
    shared = set(first) | set(second)
    dot = sum(first[key] * second[key] for key in shared)
    norm_first = math.sqrt(sum(value * value for value in first.values()))
    norm_second = math.sqrt(sum(value * value for value in second.values()))
    if norm_first == 0 or norm_second == 0:
        return 1.0
    return 1.0 - dot / (norm_first * norm_second)


def percentile(values: list[float], fraction: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, int(round(fraction * (len(ordered) - 1)))))
    return ordered[index]


def make_title(words: list[dict], limit: int = 60) -> str:
    """First sentence after the boundary, trimmed. Meant to be rewritten."""
    parts: list[str] = []
    for word in words[:24]:
        parts.append(word["word"])
        if SENTENCE_END.search(word["word"]) and len(" ".join(parts)) > 18:
            break
    text = " ".join(parts).strip()
    text = re.sub(r"\s+([,.!?;:])", r"\1", text)
    if len(text) > limit:
        text = text[:limit].rsplit(" ", 1)[0] + "…"
    return text[:1].upper() + text[1:] if text else "Без названия"


def timestamp(seconds: float) -> str:
    seconds = max(0, int(round(seconds)))
    hours, remainder = divmod(seconds, 3600)
    minutes, secs = divmod(remainder, 60)
    return f"{hours}:{minutes:02d}:{secs:02d}" if hours else f"{minutes}:{secs:02d}"


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("words", help="timeline.json or *.words.json")
    parser.add_argument("--out", default="chapters.json")
    parser.add_argument("--target", type=int, default=0,
                        help="Preferred chapter count; 0 lets the content decide")
    parser.add_argument("--min-seconds", type=float, default=45.0,
                        help="Shortest acceptable chapter")
    parser.add_argument("--pause-percentile", type=float, default=0.92,
                        help="Which of the speaker's own gaps counts as long")
    parser.add_argument("--min-shift", type=float, default=0.72,
                        help="Minimum cosine distance between the vocabulary either side")
    parser.add_argument("--window", type=float, default=25.0,
                        help="Seconds of speech compared on each side of a candidate")
    args = parser.parse_args()

    words = load_words(Path(args.words))
    if len(words) < 20:
        raise SystemExit("Too few words for chapter detection")
    total = words[-1]["end"] - words[0]["start"]

    gaps = []
    for index in range(len(words) - 1):
        gaps.append(words[index + 1]["start"] - words[index]["end"])
    threshold = max(0.5, percentile(gaps, args.pause_percentile))

    candidates = []
    for index in range(len(words) - 1):
        gap = gaps[index]
        if gap < threshold:
            continue
        if not SENTENCE_END.search(words[index]["word"]):
            continue
        boundary = words[index + 1]["start"]
        before = [w for w in words if boundary - args.window <= w["end"] <= boundary]
        after = [w for w in words if boundary <= w["start"] <= boundary + args.window]
        if len(before) < 8 or len(after) < 8:
            continue
        shift = cosine_distance(bag(before), bag(after))
        candidates.append({
            "at": round(boundary, 3),
            "word_index": index + 1,
            "pause": round(gap, 3),
            "lexical_shift": round(shift, 3),
            # Pause and vocabulary change weighted together: either alone produces
            # boundaries in the middle of a topic.
            "score": round(shift * 2.0 + min(gap / threshold, 3.0), 3),
        })

    strong = [item for item in candidates if item["lexical_shift"] >= args.min_shift]
    strong.sort(key=lambda item: -item["score"])

    if args.target > 0:
        wanted = args.target - 1
    else:
        # Roughly one chapter every two to three minutes, which is what a viewer
        # can navigate without the list becoming a wall.
        wanted = max(1, min(len(strong), int(total / 150)))

    accepted: list[dict] = []
    for item in strong:
        if len(accepted) >= wanted:
            break
        if all(abs(item["at"] - existing["at"]) >= args.min_seconds for existing in accepted):
            accepted.append(item)
    accepted.sort(key=lambda item: item["at"])

    boundaries = [words[0]["start"]] + [item["at"] for item in accepted]
    chapters = []
    for index, start in enumerate(boundaries):
        end = boundaries[index + 1] if index + 1 < len(boundaries) else words[-1]["end"]
        inside = [w for w in words if start <= w["start"] < end]
        if not inside:
            continue
        source = next((item for item in accepted if abs(item["at"] - start) < 0.01), None)
        chapters.append({
            "index": index,
            "start": round(start, 3),
            "end": round(end, 3),
            "duration": round(end - start, 3),
            "timestamp": timestamp(start),
            "title": make_title(inside),
            "words": len(inside),
            "keywords": [word for word, _count in bag(inside).most_common(6)],
            "detected_by": None if source is None else {
                "pause": source["pause"], "lexical_shift": source["lexical_shift"],
                "score": source["score"],
            },
        })

    short = [item for item in chapters if item["duration"] < args.min_seconds]
    payload = {
        "version": "1.0",
        "source": str(Path(args.words).resolve()),
        "total_seconds": round(total, 3),
        "chapters": chapters,
        "detection": {
            "speaker_pause_threshold": round(threshold, 3),
            "pause_percentile": args.pause_percentile,
            "min_lexical_shift": args.min_shift,
            "comparison_window_seconds": args.window,
            "candidates_found": len(candidates),
            "candidates_passing_shift": len(strong),
            "accepted": len(accepted),
        },
        "youtube_description": "\n".join(
            f"{item['timestamp']} {item['title']}" for item in chapters
        ),
        "ffmpeg_chapters_note": (
            "Для встраивания глав в MP4 используй scripts/render_chunked.py --chapters "
            "или ffmpeg -i in.mp4 -i chapters.txt -map_metadata 1 -codec copy out.mp4"
        ),
        "titles_are_drafts": (
            "Заголовки взяты из первой фразы главы механически. Это отправная точка: "
            "перепиши их под смысл, иначе описание будет выглядеть автоматическим."
        ),
    }
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out).write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
                              encoding="utf-8")

    # ffmetadata alongside, since that is what ffmpeg needs to embed the chapters.
    metadata = [";FFMETADATA1"]
    for item in chapters:
        metadata += [
            "[CHAPTER]", "TIMEBASE=1/1000",
            f"START={int(item['start'] * 1000)}",
            f"END={int(item['end'] * 1000)}",
            f"title={item['title']}",
        ]
    metadata_path = Path(args.out).with_suffix(".ffmetadata.txt")
    metadata_path.write_text("\n".join(metadata) + "\n", encoding="utf-8")

    emit({
        "status": "ok",
        "out": str(Path(args.out).resolve()),
        "ffmetadata": str(metadata_path.resolve()),
        "total_minutes": round(total / 60, 1),
        "chapters": len(chapters),
        "average_chapter_minutes": round(total / 60 / max(1, len(chapters)), 1),
        "shorter_than_minimum": [item["timestamp"] for item in short],
        "speaker_pause_threshold": round(threshold, 3),
        "candidates": len(candidates),
        "listing": [{"timestamp": item["timestamp"], "title": item["title"],
                     "minutes": round(item["duration"] / 60, 1),
                     "keywords": item["keywords"][:4]}
                    for item in chapters],
        "reminder": payload["titles_are_drafts"],
    })


if __name__ == "__main__":
    main()