#!/usr/bin/env python3
# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Run LLM emotion/tone analysis on existing utterances without re-segmenting."""
from __future__ import annotations

import sqlite3
import sys
from pathlib import Path

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

from umbraaudio.config import load_config
from umbraaudio.mapping.emotion import EmotionMapper
from umbraaudio.voices.model_coordinator import ModelCoordinator, PipelineStage

MODEL = "google/gemma-4-e4b"


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(
        "--dry-run",
        action="store_true",
        help="Print stats without writing to DB",
    )
    args = parser.parse_args()

    cfg = load_config(args.project)
    coordinator = ModelCoordinator(cfg)
    if not coordinator.switch_to(
        PipelineStage.CHAPTER_ANALYSIS, lm_model=MODEL
    ).get("ok"):
        print("Failed to switch models to CHAPTER_ANALYSIS stage")
        sys.exit(1)

    from openai import OpenAI

    client = OpenAI(base_url=cfg.llm.lmstudio_base_url, api_key="lm-studio")
    mapper = EmotionMapper(llm_client=client, model=MODEL)

    db_path = cfg.db_path
    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 analyze")
    print(f"Model: {MODEL}")

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

        utts = db.execute(
            "SELECT id, utterance_index, text, type, speaker_name "
            "FROM utterances WHERE chapter_id=? ORDER BY utterance_index",
            (ch_id,),
        ).fetchall()
        if not utts:
            print(f"  {title}: no utterances, skipping")
            continue

        utt_list = [dict(u) for u in utts]
        context = (ch["source_text"] or "")[:500]
        emotions = mapper.analyze_utterances(utt_list, context)

        if args.dry_run:
            print(f"  {title} (id={ch_id}): {len(emotions)} emotions returned")
            total += len(emotions)
            continue

        for u, emo in zip(utt_list, emotions, strict=False):
            db.execute(
                "UPDATE utterances SET emotion=?, pace=?, instruct=? WHERE id=?",
                (
                    emo.get("emotion", "neutral"),
                    emo.get("pace", "moderate"),
                    emo.get("instruct", ""),
                    u["id"],
                ),
            )
        db.commit()
        total += len(emotions)
        print(f"  {title} (id={ch_id}): analyzed {len(emotions)} utterances")

    db.close()
    coordinator.switch_to(PipelineStage.IDLE)
    print(f"\nTotal: {total} utterances analyzed")
    print("Done!")


if __name__ == "__main__":
    main()