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

"""Regression tests for resegment_merge.py."""

from __future__ import annotations

import sqlite3
from typing import TYPE_CHECKING

from scripts.resegment_merge import _merge_chapter, _norm, _TargetPool

if TYPE_CHECKING:
    from pathlib import Path

# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #

def _make_db(tmp_path: Path, chapter_rows: list[dict], title: str, source: str) -> Path:
    """Create a minimal project.db with a single chapter + utterance rows."""
    db_path = tmp_path / "project.db"
    db = sqlite3.connect(str(db_path))
    db.execute("""
        CREATE TABLE chapters (
            id INTEGER PRIMARY KEY,
            title TEXT,
            chapter_index INTEGER NOT NULL,
            source_text TEXT,
            status TEXT
        )
    """)
    db.execute("""
        CREATE TABLE utterances (
            id INTEGER PRIMARY KEY,
            chapter_id INTEGER NOT NULL,
            utterance_index INTEGER NOT NULL,
            text TEXT NOT NULL,
            type TEXT NOT NULL DEFAULT 'narration',
            context_before TEXT,
            context_after TEXT,
            emotion TEXT,
            pace TEXT,
            emphasis_words TEXT,
            pause_before_ms INTEGER,
            pause_after_ms INTEGER,
            instruct TEXT,
            speaker_name TEXT,
            attribution_confidence REAL DEFAULT 1.0,
            attribution_reasoning TEXT,
            emotion_confidence REAL DEFAULT 1.0,
            user_edited INTEGER NOT NULL DEFAULT 0,
            status TEXT NOT NULL DEFAULT 'pending'
        )
    """)
    db.execute("INSERT INTO chapters (id, title, chapter_index, source_text) VALUES (1,?,?,?)",
               (title, 1, source))
    for r in chapter_rows:
        db.execute(
            "INSERT INTO utterances (id, chapter_id, utterance_index, text, type, "
            "speaker_name, user_edited, emotion, pace, instruct, status) "
            "VALUES (?, 1, ?, ?, ?, ?, ?, ?, ?, ?, 'pending')",
            (r["id"], r["index"], r["text"], r["type"],
             r.get("speaker", "narrator"), r.get("user_edited", 0),
             r.get("emotion"), r.get("pace"), r.get("instruct")),
        )
    db.commit()
    db.close()
    return db_path


def _read_rows(db_path: Path, chapter_id: int = 1) -> list[dict]:
    db = sqlite3.connect(str(db_path))
    db.row_factory = sqlite3.Row
    rows = db.execute(
        "SELECT id, utterance_index, text, type, speaker_name, user_edited "
        "FROM utterances WHERE chapter_id=? ORDER BY utterance_index",
        (chapter_id,),
    ).fetchall()
    db.close()
    return [dict(r) for r in rows]


def _make_seg():
    from umbraaudio.mapping.segmenter import Segmenter
    return Segmenter()


# --------------------------------------------------------------------------- #
# tests
# --------------------------------------------------------------------------- #

class TestNorm:
    def test_strips_quotes_and_commas(self):
        assert _norm("\u201cHello,\u201d he said.") == "hello, he said"

    def test_case_insensitive(self):
        assert _norm("Hallr says,") == _norm("hallr says,")


class TestTargetPool:
    def test_exact_pop(self):
        pool = _TargetPool([
            {"text": "Hallr shrugs.", "type": "narration"},
            {"text": "Hallr shrugs.", "type": "narration"},
        ])
        s = pool.exact("hallr shrugs")
        assert s is not None
        assert s["text"] == "Hallr shrugs."
        s2 = pool.exact("hallr shrugs")
        assert s2 is not None
        assert pool.exact("hallr shrugs") is None

    def test_narration_rewrite_containment(self):
        pool = _TargetPool([
            {"text": "Ready to Fight\nHallr was up early.", "type": "narration"},
        ])
        s = pool.narration_rewrite("hallr was up early")
        assert s is not None
        assert "Hallr" in s["text"]


class TestMergeChapter:
    def test_dialogue_preserved(self, tmp_path):
        source = 'Hallr says, \u201cI will fight.\u201d\nFeliu nods.\n\u201cSounds good.\u201d'
        db = _make_db(tmp_path, [
            {"id": 1, "index": 0, "text": "\u201cI will fight.\u201d",
             "type": "dialogue", "speaker": "Hallr"},
            {"id": 2, "index": 1, "text": "Feliu nods.", "type": "narration"},
            {"id": 3, "index": 2, "text": "\u201cSounds good.\u201d",
             "type": "dialogue", "speaker": "Feliu"},
        ], "Test", source)
        seg = _make_seg()
        kept, new, rewritten, dropped = _merge_chapter(
            seg, db, _read_db_chapter(db, 1), ["Hallr", "Feliu"], 0.82, dry_run=True
        )
        rows = _read_rows(db, 1)
        # dialogue ids and speakers unchanged
        assert rows[0]["id"] == 1 and rows[0]["speaker_name"] == "Hallr"
        assert rows[2]["id"] == 3 and rows[2]["speaker_name"] == "Feliu"

    def test_leading_comma_rewrite(self, tmp_path):
        source = "Hallr shrugs.\nHe moves on."
        db = _make_db(tmp_path, [
            {"id": 1, "index": 0, "text": ", Hallr shrugs.", "type": "narration"},
            {"id": 2, "index": 1, "text": "He moves on.", "type": "narration"},
        ], "Test", source)
        seg = _make_seg()
        _merge_chapter(
            seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=False
        )
        rows = _read_rows(db, 1)
        assert rows[0]["text"] == "Hallr shrugs."
        assert rows[0]["speaker_name"] == "narrator"

    def test_user_edited_row_not_dropped(self, tmp_path):
        source = "Hallr was tired."
        db = _make_db(tmp_path, [
            {"id": 1, "index": 0, "text": "Hallr was tired.", "type": "narration"},
            {"id": 2, "index": 1, "text": "I hope you like the story.",
             "type": "narration", "user_edited": 1},
        ], "Test", source)
        seg = _make_seg()
        kept, new, rewritten, dropped = _merge_chapter(
            seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=True
        )
        assert dropped == []
        rows = _read_rows(db, 1)
        assert len(rows) == 2

    def test_scene_break_preserved_and_added(self, tmp_path):
        source = "Scene one.\n<<<>>>\nScene two."
        db = _make_db(tmp_path, [
            {"id": 1, "index": 0, "text": "Scene one.", "type": "narration"},
            {"id": 2, "index": 1, "text": "<<<>>>", "type": "scene_break"},
            {"id": 3, "index": 2, "text": "Scene two.", "type": "narration"},
        ], "Test", source)
        seg = _make_seg()
        _merge_chapter(
            seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=False
        )
        rows = _read_rows(db, 1)
        scene_rows = [r for r in rows if r["type"] == "scene_break"]
        assert len(scene_rows) == 1
        assert scene_rows[0]["speaker_name"] == "scene_break"

    def test_indices_renumbered(self, tmp_path):
        source = "One.\n<<<>>>\nTwo.\nThree."
        db = _make_db(tmp_path, [
            {"id": 1, "index": 0, "text": "One.", "type": "narration"},
            {"id": 3, "index": 2, "text": "Three.", "type": "narration"},
            {"id": 2, "index": 1, "text": "Two.", "type": "narration"},
        ], "Test", source)
        seg = _make_seg()
        _merge_chapter(
            seg, db, _read_db_chapter(db, 1), [], 0.82, dry_run=False
        )
        rows = _read_rows(db, 1)
        indices = [r["utterance_index"] for r in rows]
        assert indices == list(range(len(rows)))


def _read_db_chapter(db_path: Path, chapter_id: int) -> dict:
    db = sqlite3.connect(str(db_path))
    db.row_factory = sqlite3.Row
    r = db.execute("SELECT id, title, chapter_index, source_text FROM chapters WHERE id=?",
                   (chapter_id,)).fetchone()
    db.close()
    return dict(r)