UmbraAudio/scripts/qa_attribution.py

#!/usr/bin/env python3
"""Mechanical QA cross-check of speaker attribution against speech tags.

Usage:
    python scripts/qa_attribution.py                 # check all chapters
    python scripts/qa_attribution.py --chapters 6    # check chapter id 6
    python scripts/qa_attribution.py --fix           # auto-fix exact name-tag mismatches
    python scripts/qa_attribution.py --type role_tag_guessed   # only one issue type
"""

from __future__ import annotations

import argparse
import json
import sqlite3
import sys
from pathlib import Path

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

from umbraaudio.analysis.attribution_qa import (
    QAIssue,
    check_attribution,
    fix_high_confidence,
    verify_conversation_flow,
)
from umbraaudio.config import load_config


def _build_alias_map(speakers_path: Path) -> dict[str, str]:
    alias_map: dict[str, str] = {}
    if not speakers_path.exists():
        return alias_map
    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
        alias_map[primary.lower()] = primary
        for alias in s.get("aliases", []):
            if alias:
                alias_map[alias.lower()] = primary
        first = primary.split()[0].lower() if primary.split() else ""
        if first and first not in alias_map:
            alias_map[first] = primary
    return alias_map


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    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")
    parser.add_argument("--fix", action="store_true", help="Auto-fix exact speech-tag mismatches")
    parser.add_argument(
        "--type",
        dest="issue_type",
        type=str,
        default=None,
        help=(
            "Only report one issue type "
            "(speech_tag_mismatch, narrator_dialogue, role_tag_guessed, "
            "address_mismatch, narration_speaker_mismatch, long_monologue)"
        ),
    )
    args = parser.parse_args()

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

    alias_map = _build_alias_map(cfg.data_dir / "speakers.json")
    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 FROM chapters "
            f"WHERE id IN ({placeholders}) ORDER BY chapter_index",
            chapter_ids,
        ).fetchall()
    else:
        chapters = db.execute(
            "SELECT id, title, chapter_index FROM chapters ORDER BY chapter_index"
        ).fetchall()

    totals = {
        "speech_tag_mismatch": 0,
        "narrator_dialogue": 0,
        "role_tag_guessed": 0,
        "address_mismatch": 0,
        "narration_speaker_mismatch": 0,
        "long_monologue": 0,
    }
    all_issues: list[QAIssue] = []

    for ch in chapters:
        ch_id = ch["id"]
        rows = db.execute(
            "SELECT id, utterance_index, text, type, speaker_name "
            "FROM utterances WHERE chapter_id=? ORDER BY utterance_index",
            (ch_id,),
        ).fetchall()
        utts = [dict(r) for r in rows]

        result = check_attribution(utts, ch_id, alias_map)

        flow_result = verify_conversation_flow(utts, ch_id, alias_map)
        result.issues.extend(flow_result.issues)

        if args.fix:
            fixed = fix_high_confidence(utts, alias_map)
            for orig, new in zip(utts, fixed):
                if orig.get("speaker_name") != new.get("speaker_name"):
                    db.execute(
                        "UPDATE utterances SET speaker_name=? WHERE id=?",
                        (new["speaker_name"], orig["id"]),
                    )
                    print(
                        f"  FIX ch{ch_id} utt{new.get('utterance_index')}: "
                        f"{orig.get('speaker_name')} -> {new['speaker_name']}"
                    )
            db.commit()

        for issue in result.issues:
            if args.issue_type and issue.issue_type != args.issue_type:
                continue
            totals[issue.issue_type] += 1
            all_issues.append(issue)
            print(
                f"[{issue.issue_type}] ch{issue.chapter_id} utt{issue.utterance_index} "
                f"speaker={issue.speaker!r} expected={issue.expected!r} conf={issue.confidence}"
            )
            if issue.evidence:
                print(f"    evidence: {issue.evidence}")
            print(f"    text: {issue.text}")

    db.close()

    print(f"\nTotals: {totals}")
    if all_issues:
        print(f"Total issues: {len(all_issues)} (review or fix before accepting voices)")


if __name__ == "__main__":
    main()