#!/usr/bin/env python3
"""Re-segment all chapters using mechanical splitter + LLM attribution."""

from __future__ import annotations

import json
import sqlite3
import sys
from pathlib import Path

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

from umbraaudio.analysis.llm_attributor import LLMAttributor
from umbraaudio.config import LLMConfig
from umbraaudio.mapping.segmenter import Segmenter


def main():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--project", type=str, default=None, help="Project name or path"
    )
    parser.add_argument(
        "--chapters", type=str, default=None, help="Comma-separated chapter IDs (e.g. '5,6,7')"
    )
    parser.add_argument("--model", type=str, default="gemma-4-12b-it-qat", help="LLM model name")
    parser.add_argument(
        "--mechanical-only", action="store_true", help="Skip LLM attribution, use mechanical only"
    )
    args = parser.parse_args()

    from umbraaudio.config import load_config

    cfg = load_config(args.project)
    config = cfg.llm

    db_path = cfg.db_path
    if not db_path.exists():
        print(f"Database not found: {db_path}")
        return

    db = sqlite3.connect(str(db_path))
    db.row_factory = sqlite3.Row

    if args.chapters:
        chapter_ids = [int(x.strip()) for x in args.chapters.split(",")]
        placeholders = ",".join("?" * len(chapter_ids))
        chapters = db.execute(
            f"SELECT id, title, chapter_index, source_text FROM chapters "
            f"WHERE id IN ({placeholders}) ORDER BY chapter_index",
            chapter_ids,
        ).fetchall()
    else:
        chapters = db.execute(
            "SELECT id, title, chapter_index, source_text FROM chapters ORDER BY chapter_index"
        ).fetchall()

    print(f"Found {len(chapters)} chapters to segment")
    print(f"Model: {args.model}")
    print(f"LLM attribution: {'disabled' if args.mechanical_only else 'enabled'}")

    # Load character names and alias map from speakers.json
    speakers_path = cfg.data_dir / "speakers.json"
    character_names = []
    alias_map = {}
    if speakers_path.exists():
        with open(speakers_path) as f:
            data = json.load(f)
        for s in data.get("speakers", []):
            primary = s.get("primary_name", "")
            if not primary:
                continue
            character_names.append(primary)
            for alias in s.get("aliases", []):
                alias_map[alias] = primary

    print(f"Loaded {len(character_names)} character names, {len(alias_map)} aliases")

    seg = Segmenter(alias_map=alias_map)
    attributor = LLMAttributor(config, model=args.model, data_dir=cfg.data_dir) if not args.mechanical_only else None

    total_utterances = 0

    for ch in chapters:
        ch_id = ch["id"]
        title = ch["title"] or f"Chapter {ch['chapter_index']}"
        source_text = ch["source_text"]

        old_count = db.execute(
            "SELECT COUNT(*) as cnt FROM utterances WHERE chapter_id = ?", (ch_id,)
        ).fetchone()["cnt"]

        # Phase 1: Mechanical segmentation
        mechanical = seg.segment(source_text, ch_id, character_names, chapter_title=title)

        utterances = []
        for i, u in enumerate(mechanical):
            speaker = u.get("character_name", "narrator")
            if speaker is None:
                speaker = "narrator"

            is_scene_break = u["text"].strip() in ("<<<>>>",)

            utterances.append(
                {
                    "utterance_index": i,
                    "text": u["text"],
                    "type": u.get("type", "narration"),
                    "speaker_name": "scene_break" if is_scene_break else speaker,
                }
            )

        # Phase 2: LLM attribution (if enabled)
        if attributor and not args.mechanical_only:
            utterances = attributor.attribute_speakers(utterances, title, ch_id)

        # Write to database
        db.execute("DELETE FROM utterances WHERE chapter_id = ?", (ch_id,))

        for u in utterances:
            db.execute(
                "INSERT INTO utterances "
                "(chapter_id, utterance_index, text, type, speaker_name) "
                "VALUES (?, ?, ?, ?, ?)",
                (ch_id, u["utterance_index"], u["text"], u["type"], u["speaker_name"]),
            )

        db.commit()
        total_utterances += len(utterances)

        print(f"  {title}: {old_count} -> {len(utterances)} utterances")

    db.close()

    print(f"\nTotal: {total_utterances} utterances across {len(chapters)} chapters")
    print("Done!")


if __name__ == "__main__":
    main()