UmbraAudio/umbraaudio/database.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import aiosqlite

SCHEMA = """
CREATE TABLE IF NOT EXISTS project (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    source_path TEXT NOT NULL,
    source_format TEXT NOT NULL,
    created_at TEXT NOT NULL DEFAULT (datetime('now')),
    updated_at TEXT NOT NULL DEFAULT (datetime('now')),
    status TEXT NOT NULL DEFAULT 'initialized'
);

CREATE TABLE IF NOT EXISTS chapters (
    id INTEGER PRIMARY KEY,
    project_id INTEGER NOT NULL REFERENCES project(id),
    chapter_index INTEGER NOT NULL,
    title TEXT,
    source_text TEXT NOT NULL,
    word_count INTEGER NOT NULL DEFAULT 0,
    status TEXT NOT NULL DEFAULT 'parsed',
    audio_path TEXT,
    created_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS characters (
    id INTEGER PRIMARY KEY,
    project_id INTEGER NOT NULL REFERENCES project(id),
    name TEXT NOT NULL,
    aliases TEXT NOT NULL DEFAULT '[]',
    gender TEXT,
    age_range TEXT,
    is_narrator INTEGER NOT NULL DEFAULT 0,
    registry_id INTEGER,
    profile_json TEXT,
    voice_description TEXT,
    voice_ref_path TEXT,
    speaker_embedding BLOB,
    pitch_register TEXT DEFAULT 'mid',
    pace_baseline TEXT DEFAULT 'moderate',
    energy TEXT DEFAULT 'measured',
    timbre TEXT DEFAULT 'clear',
    texture TEXT DEFAULT 'smooth',
    formality TEXT DEFAULT 'conversational',
    emotional_baseline TEXT DEFAULT 'neutral_calm',
    distinguishing_trait TEXT,
    accent_marker TEXT DEFAULT 'none',
    status TEXT NOT NULL DEFAULT 'identified',
    created_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS voice_review_history (
    id INTEGER PRIMARY KEY,
    character_name TEXT NOT NULL,
    action TEXT NOT NULL,
    prompt_text TEXT,
    resulting_exaggeration REAL,
    resulting_cfg_weight REAL,
    taxonomy_snapshot_json TEXT,
    reviewed_by TEXT,
    reviewed_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS pronunciation_guide (
    id INTEGER PRIMARY KEY,
    project_id INTEGER NOT NULL REFERENCES project(id),
    term TEXT NOT NULL,
    phonetic TEXT NOT NULL,
    notes TEXT,
    UNIQUE(project_id, term)
);

CREATE TABLE IF NOT EXISTS utterances (
    id INTEGER PRIMARY KEY,
    chapter_id INTEGER NOT NULL REFERENCES chapters(id),
    character_id INTEGER REFERENCES characters(id),
    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,
    attribution_reasoning TEXT,
    attribution_evidence_kind TEXT,
    attribution_tag_name TEXT,
    attribution_uncertain INTEGER NOT NULL DEFAULT 0,
    needs_review INTEGER NOT NULL DEFAULT 0,
    review_reason TEXT,
    emotion_confidence REAL DEFAULT 1.0,
    user_edited INTEGER NOT NULL DEFAULT 0,
    status TEXT NOT NULL DEFAULT 'pending',
    source_unit_key TEXT
);

CREATE TABLE IF NOT EXISTS generated_audio (
    id INTEGER PRIMARY KEY,
    utterance_id INTEGER NOT NULL REFERENCES utterances(id),
    character_id INTEGER REFERENCES characters(id),
    audio_path TEXT NOT NULL,
    duration_ms INTEGER,
    mos_score REAL,
    wer_score REAL,
    speaker_similarity REAL,
    artifact_score REAL,
    generation_params TEXT,
    status TEXT NOT NULL DEFAULT 'generated',
    created_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS quality_reports (
    id INTEGER PRIMARY KEY,
    project_id INTEGER NOT NULL REFERENCES project(id),
    chapter_id INTEGER REFERENCES chapters(id),
    report_type TEXT NOT NULL,
    metrics_json TEXT NOT NULL,
    passed INTEGER NOT NULL DEFAULT 0,
    notes TEXT,
    created_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS review_decisions (
    id INTEGER PRIMARY KEY,
    project_id INTEGER NOT NULL REFERENCES project(id),
    entity_type TEXT NOT NULL,
    entity_id INTEGER NOT NULL,
    decision TEXT NOT NULL,
    notes TEXT,
    created_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS name_tokens (
    id INTEGER PRIMARY KEY,
    project_id INTEGER NOT NULL REFERENCES project(id),
    token TEXT NOT NULL,
    kind TEXT NOT NULL DEFAULT 'name_tag',
    status TEXT NOT NULL DEFAULT 'unresolved',
    target_speaker TEXT,
    occurrences INTEGER NOT NULL DEFAULT 0,
    chapters_json TEXT NOT NULL DEFAULT '[]',
    evidence_json TEXT NOT NULL DEFAULT '[]',
    suggestion TEXT,
    suggestion_score REAL DEFAULT 0.0,
    updated_at TEXT NOT NULL DEFAULT (datetime('now')),
    UNIQUE(project_id, token)
);
CREATE INDEX IF NOT EXISTS idx_name_tokens_status ON name_tokens(status);
CREATE INDEX IF NOT EXISTS idx_name_tokens_project ON name_tokens(project_id);

CREATE INDEX IF NOT EXISTS idx_utterances_chapter ON utterances(chapter_id);
CREATE INDEX IF NOT EXISTS idx_utterances_source_unit ON utterances(source_unit_key);
CREATE INDEX IF NOT EXISTS idx_utterances_character ON utterances(character_id);
CREATE INDEX IF NOT EXISTS idx_generated_audio_utterance ON generated_audio(utterance_id);
CREATE INDEX IF NOT EXISTS idx_chapters_project ON chapters(project_id);
CREATE INDEX IF NOT EXISTS idx_characters_project ON characters(project_id);
"""


class Database:
    def __init__(self, db_path: Path) -> None:
        self.db_path = db_path
        self._conn: aiosqlite.Connection | None = None

    async def connect(self) -> None:
        self._conn = await aiosqlite.connect(str(self.db_path))
        self._conn.row_factory = aiosqlite.Row
        # WAL prevents SQLITE_BUSY when the web UI and a CLI/background batch
        # touch the same project DB concurrently. Connections to the same DB
        # share a single WAL; each new writer connection still enables it here.
        try:
            await self._conn.execute("PRAGMA journal_mode=WAL")
        except Exception:
            pass
        await self._conn.executescript(SCHEMA)
        await self._migrate()
        await self._conn.commit()

    async def _migrate(self) -> None:
        cursor = await self._conn.execute("PRAGMA table_info(characters)")
        existing_cols = {row[1] for row in await cursor.fetchall()}

        new_cols = {
            "pitch_register": "TEXT DEFAULT 'mid'",
            "pace_baseline": "TEXT DEFAULT 'moderate'",
            "energy": "TEXT DEFAULT 'measured'",
            "timbre": "TEXT DEFAULT 'clear'",
            "texture": "TEXT DEFAULT 'smooth'",
            "formality": "TEXT DEFAULT 'conversational'",
            "emotional_baseline": "TEXT DEFAULT 'neutral_calm'",
            "distinguishing_trait": "TEXT",
            "accent_marker": "TEXT DEFAULT 'none'",
            "registry_id": "INTEGER",
        }
        for col, typedef in new_cols.items():
            if col not in existing_cols:
                await self._conn.execute(f"ALTER TABLE characters ADD COLUMN {col} {typedef}")

        cursor = await self._conn.execute("PRAGMA table_info(utterances)")
        utt_cols = {row[1] for row in await cursor.fetchall()}
        if "user_edited" not in utt_cols:
            await self._conn.execute(
                "ALTER TABLE utterances ADD COLUMN user_edited INTEGER NOT NULL DEFAULT 0"
            )
        for col, typedef in (
            ("attribution_evidence_kind", "TEXT"),
            ("attribution_tag_name", "TEXT"),
            ("attribution_uncertain", "INTEGER NOT NULL DEFAULT 0"),
            ("needs_review", "INTEGER NOT NULL DEFAULT 0"),
            ("review_reason", "TEXT"),
            ("source_unit_key", "TEXT"),
        ):
            if col not in utt_cols:
                await self._conn.execute(
                    f"ALTER TABLE utterances ADD COLUMN {col} {typedef}"
                )
        await self._conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_utterances_source_unit ON utterances(source_unit_key)"
        )

        cursor = await self._conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name='voice_review_history'"
        )
        if not await cursor.fetchone():
            await self._conn.execute("""
                CREATE TABLE IF NOT EXISTS voice_review_history (
                    id INTEGER PRIMARY KEY,
                    character_name TEXT NOT NULL,
                    action TEXT NOT NULL,
                    prompt_text TEXT,
                    resulting_exaggeration REAL,
                    resulting_cfg_weight REAL,
                    taxonomy_snapshot_json TEXT,
                    reviewed_by TEXT,
                    reviewed_at TEXT NOT NULL DEFAULT (datetime('now'))
                )
            """)

    async def close(self) -> None:
        if self._conn:
            await self._conn.close()

    @property
    def conn(self) -> aiosqlite.Connection:
        if not self._conn:
            raise RuntimeError("Database not connected. Call connect() first.")
        return self._conn

    async def create_project(self, name: str, source_path: str, source_format: str) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO project (name, source_path, source_format) VALUES (?, ?, ?)",
            (name, source_path, source_format),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def update_project_status(self, project_id: int, status: str) -> None:
        await self.conn.execute(
            "UPDATE project SET status = ?, updated_at = datetime('now') WHERE id = ?",
            (status, project_id),
        )
        await self.conn.commit()

    async def insert_chapter(
        self,
        project_id: int,
        index: int,
        title: str | None,
        source_text: str,
        word_count: int,
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO chapters (project_id, chapter_index, title, source_text, word_count) "
            "VALUES (?, ?, ?, ?, ?)",
            (project_id, index, title, source_text, word_count),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def insert_character(
        self,
        project_id: int,
        name: str,
        aliases: list[str] | None = None,
        gender: str | None = None,
        age_range: str | None = None,
        is_narrator: bool = False,
        registry_id: int | None = None,
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO characters "
            "(project_id, name, aliases, gender, age_range, is_narrator, registry_id) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            (
                project_id,
                name,
                json.dumps(aliases or []),
                gender,
                age_range,
                int(is_narrator),
                registry_id,
            ),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def update_character_profile(
        self,
        character_id: int,
        profile_json: dict[str, Any],
        voice_description: str | None = None,
    ) -> None:
        taxonomy_fields = {}
        for key in (
            "pitch_register",
            "pace_baseline",
            "energy",
            "timbre",
            "texture",
            "formality",
            "emotional_baseline",
            "distinguishing_trait",
            "accent_marker",
        ):
            if key in profile_json:
                taxonomy_fields[key] = profile_json[key]

        set_clauses = [
            "profile_json = ?",
            "voice_description = COALESCE(?, voice_description)",
            "status = 'profiled'",
        ]
        params: list[Any] = [json.dumps(profile_json), voice_description]

        for field, value in taxonomy_fields.items():
            set_clauses.append(f"{field} = ?")
            params.append(value)

        params.append(character_id)
        await self.conn.execute(
            f"UPDATE characters SET {', '.join(set_clauses)} WHERE id = ?",
            params,
        )
        await self.conn.commit()

    async def update_character_voice(
        self, character_id: int, voice_ref_path: str, speaker_embedding: bytes | None = None
    ) -> None:
        await self.conn.execute(
            "UPDATE characters SET voice_ref_path = ?, speaker_embedding = ?, status = 'voiced' "
            "WHERE id = ?",
            (voice_ref_path, speaker_embedding, character_id),
        )
        await self.conn.commit()

    async def insert_pronunciation(
        self, project_id: int, term: str, phonetic: str, notes: str | None = None
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT OR IGNORE INTO pronunciation_guide (project_id, term, phonetic, notes) "
            "VALUES (?, ?, ?, ?)",
            (project_id, term, phonetic, notes),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def insert_utterance(
        self,
        chapter_id: int,
        character_id: int | None,
        index: int,
        text: str,
        type: str = "narration",
        context_before: str | None = None,
        context_after: str | None = None,
        emotion: str | None = None,
        pace: str | None = None,
        instruct: str | None = None,
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO utterances "
            "(chapter_id, character_id, utterance_index, text, type, context_before, context_after, "
            "emotion, pace, instruct) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                chapter_id,
                character_id,
                index,
                text,
                type,
                context_before,
                context_after,
                emotion,
                pace,
                instruct,
            ),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def insert_generated_audio(
        self,
        utterance_id: int,
        character_id: int | None,
        audio_path: str,
        duration_ms: int | None = None,
        generation_params: dict[str, Any] | None = None,
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO generated_audio "
            "(utterance_id, character_id, audio_path, duration_ms, generation_params) "
            "VALUES (?, ?, ?, ?, ?)",
            (
                utterance_id,
                character_id,
                audio_path,
                duration_ms,
                json.dumps(generation_params) if generation_params else None,
            ),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def update_audio_quality(
        self,
        audio_id: int,
        mos_score: float | None = None,
        wer_score: float | None = None,
        speaker_similarity: float | None = None,
        artifact_score: float | None = None,
    ) -> None:
        await self.conn.execute(
            "UPDATE generated_audio SET mos_score = ?, wer_score = ?, "
            "speaker_similarity = ?, artifact_score = ? WHERE id = ?",
            (mos_score, wer_score, speaker_similarity, artifact_score, audio_id),
        )
        await self.conn.commit()

    async def insert_quality_report(
        self,
        project_id: int,
        chapter_id: int | None,
        report_type: str,
        metrics: dict[str, Any],
        passed: bool,
        notes: str | None = None,
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO quality_reports "
            "(project_id, chapter_id, report_type, metrics_json, passed, notes) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (project_id, chapter_id, report_type, json.dumps(metrics), int(passed), notes),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def insert_review_decision(
        self,
        project_id: int,
        entity_type: str,
        entity_id: int,
        decision: str,
        notes: str | None = None,
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO review_decisions (project_id, entity_type, entity_id, decision, notes) "
            "VALUES (?, ?, ?, ?, ?)",
            (project_id, entity_type, entity_id, decision, notes),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def get_project(self, project_id: int) -> dict[str, Any] | None:
        cursor = await self.conn.execute("SELECT * FROM project WHERE id = ?", (project_id,))
        row = await cursor.fetchone()
        return dict(row) if row else None

    async def active_project_id(self) -> int:
        """Return the id of the project row in this project's DB.

        One project per DB file is the isolation guarantee, so the active id is
        simply the first (lowest) id in the ``project`` table. Raises if the DB
        has no project row yet (e.g. before ``init`` has populated it).
        """
        cursor = await self.conn.execute("SELECT id FROM project ORDER BY id LIMIT 1")
        row = await cursor.fetchone()
        if row is None:
            raise RuntimeError(
                f"No project row in {self.db_path}; run 'umbraaudio init' first."
            )
        return int(row[0])

    async def get_chapters(self, project_id: int) -> list[dict[str, Any]]:
        cursor = await self.conn.execute(
            "SELECT * FROM chapters WHERE project_id = ? ORDER BY chapter_index", (project_id,)
        )
        return [dict(r) for r in await cursor.fetchall()]

    async def get_characters(self, project_id: int) -> list[dict[str, Any]]:
        cursor = await self.conn.execute(
            "SELECT * FROM characters WHERE project_id = ? ORDER BY is_narrator DESC, name",
            (project_id,),
        )
        return [dict(r) for r in await cursor.fetchall()]

    async def get_utterances(self, chapter_id: int) -> list[dict[str, Any]]:
        cursor = await self.conn.execute(
            "SELECT u.*, c.name as character_name FROM utterances u "
            "LEFT JOIN characters c ON u.character_id = c.id "
            "WHERE u.chapter_id = ? ORDER BY u.utterance_index",
            (chapter_id,),
        )
        return [dict(r) for r in await cursor.fetchall()]

    async def get_pronunciation_guide(self, project_id: int) -> dict[str, str]:
        cursor = await self.conn.execute(
            "SELECT term, phonetic FROM pronunciation_guide WHERE project_id = ?",
            (project_id,),
        )
        return {row["term"]: row["phonetic"] for row in await cursor.fetchall()}

    async def get_project_by_path(self, source_path: str) -> dict[str, Any] | None:
        cursor = await self.conn.execute(
            "SELECT * FROM project WHERE source_path = ?", (source_path,)
        )
        row = await cursor.fetchone()
        return dict(row) if row else None

    async def log_defect(
        self,
        project_id: int,
        stage: str,
        severity: str,
        message: str,
        context: str | None = None,
        chapter_id: int | None = None,
    ) -> int:
        metrics = {
            "stage": stage,
            "severity": severity,
            "message": message,
            "context": context,
        }
        return await self.insert_quality_report(
            project_id=project_id,
            chapter_id=chapter_id,
            report_type=f"defect_{severity.lower()}",
            metrics=metrics,
            passed=False,
            notes=message,
        )

    async def log_voice_review(
        self,
        character_name: str,
        action: str,
        prompt_text: str | None = None,
        resulting_exaggeration: float | None = None,
        resulting_cfg_weight: float | None = None,
        taxonomy_snapshot: dict[str, Any] | None = None,
        reviewed_by: str | None = None,
    ) -> int:
        cursor = await self.conn.execute(
            "INSERT INTO voice_review_history "
            "(character_name, action, prompt_text, resulting_exaggeration, "
            "resulting_cfg_weight, taxonomy_snapshot_json, reviewed_by) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            (
                character_name,
                action,
                prompt_text,
                resulting_exaggeration,
                resulting_cfg_weight,
                json.dumps(taxonomy_snapshot) if taxonomy_snapshot else None,
                reviewed_by,
            ),
        )
        await self.conn.commit()
        return cursor.lastrowid  # type: ignore[return-value]

    async def get_voice_review_history(self, character_name: str) -> list[dict[str, Any]]:
        cursor = await self.conn.execute(
            "SELECT * FROM voice_review_history WHERE character_name = ? ORDER BY reviewed_at",
            (character_name,),
        )
        return [dict(r) for r in await cursor.fetchall()]

    async def update_character_status(self, character_id: int, status: str) -> None:
        await self.conn.execute(
            "UPDATE characters SET status = ? WHERE id = ?",
            (status, character_id),
        )
        await self.conn.commit()

    async def update_utterance_speaker(
        self,
        utterance_id: int,
        speaker_name: str,
        attribution_confidence: float | None = None,
        attribution_reasoning: str | None = None,
        emotion: str | None = None,
        emotion_confidence: float | None = None,
        evidence_kind: str | None = None,
        tag_name: str | None = None,
        uncertain: int | None = None,
        needs_review: int | None = None,
        review_reason: str | None = None,
    ) -> None:
        """Update utterance with LLM attribution results."""
        fields = ["speaker_name = ?"]
        params = [speaker_name]

        if attribution_confidence is not None:
            fields.append("attribution_confidence = ?")
            params.append(attribution_confidence)

        if attribution_reasoning is not None:
            fields.append("attribution_reasoning = ?")
            params.append(attribution_reasoning)

        if emotion is not None:
            fields.append("emotion = ?")
            params.append(emotion)

        if emotion_confidence is not None:
            fields.append("emotion_confidence = ?")
            params.append(emotion_confidence)

        if evidence_kind is not None:
            fields.append("attribution_evidence_kind = ?")
            params.append(evidence_kind)

        if tag_name is not None:
            fields.append("attribution_tag_name = ?")
            params.append(tag_name)

        if uncertain is not None:
            fields.append("attribution_uncertain = ?")
            params.append(uncertain)

        if needs_review is not None:
            fields.append("needs_review = ?")
            params.append(needs_review)

        if review_reason is not None:
            fields.append("review_reason = ?")
            params.append(review_reason)

        params.append(utterance_id)
        await self.conn.execute(
            f"UPDATE utterances SET {', '.join(fields)} WHERE id = ?",
            params,
        )
        await self.conn.commit()