UmbraAudio/umbraaudio/voices/registry.py

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

import aiosqlite

from umbraaudio.analysis import CharacterProfile

logger = logging.getLogger(__name__)

REGISTRY_SCHEMA = """
CREATE TABLE IF NOT EXISTS voice_registry (
    id INTEGER PRIMARY KEY,
    character_name TEXT NOT NULL UNIQUE,
    aliases TEXT NOT NULL DEFAULT '[]',
    gender TEXT,
    age_range TEXT,
    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',
    voice_description TEXT,
    voice_ref_path TEXT,
    vmatrix_path TEXT,
    sample_audio_path TEXT,
    speaker_embedding BLOB,
    source_project TEXT,
    is_narrator INTEGER NOT NULL DEFAULT 0,
    is_general INTEGER NOT NULL DEFAULT 0,
    status TEXT NOT NULL DEFAULT 'pending',
    created_at TEXT NOT NULL DEFAULT (datetime('now')),
    updated_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS registry_aliases (
    id INTEGER PRIMARY KEY,
    registry_id INTEGER NOT NULL REFERENCES voice_registry(id),
    alias_name TEXT NOT NULL,
    UNIQUE(registry_id, alias_name)
);

CREATE INDEX IF NOT EXISTS idx_registry_alias ON registry_aliases(alias_name);
"""


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

    async def connect(self) -> None:
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = await aiosqlite.connect(str(self.db_path))
        self._conn.row_factory = aiosqlite.Row
        await self._conn.executescript(REGISTRY_SCHEMA)
        await self._migrate()
        await self._conn.commit()

    async def _migrate(self) -> None:
        """Add columns introduced after the table was first created."""
        cursor = await self._conn.execute("PRAGMA table_info(voice_registry)")
        cols = {row["name"] for row in await cursor.fetchall()}
        if "vmatrix_path" not in cols:
            # vmatrix_path: path to the captured vmatrix (.pt) for the approved voice
            await self._conn.execute(
                "ALTER TABLE voice_registry ADD COLUMN vmatrix_path TEXT"
            )

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

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

    async def find_by_name(self, name: str) -> dict[str, Any] | None:
        cursor = await self.conn.execute(
            "SELECT * FROM voice_registry WHERE LOWER(character_name) = LOWER(?)",
            (name,),
        )
        row = await cursor.fetchone()
        if row:
            return dict(row)
        cursor2 = await self.conn.execute(
            """SELECT vr.* FROM voice_registry vr
               JOIN registry_aliases ra ON ra.registry_id = vr.id
               WHERE LOWER(ra.alias_name) = LOWER(?)""",
            (name,),
        )
        row2 = await cursor2.fetchone()
        return dict(row2) if row2 else None

    async def find_or_create(
        self, profile: CharacterProfile, source_project: str = ""
    ) -> tuple[int, bool]:
        existing = await self.find_by_name(profile.name)
        if existing:
            return existing["id"], False

        cursor = await self.conn.execute(
            """INSERT INTO voice_registry
               (character_name, aliases, gender, age_range, pitch_register,
                pace_baseline, energy, timbre, texture, formality,
                emotional_baseline, distinguishing_trait, accent_marker,
                voice_description, source_project, is_narrator, is_general, status)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                profile.name,
                json.dumps(profile.aliases),
                profile.gender,
                profile.age_range,
                profile.pitch_register,
                profile.pace_baseline,
                profile.energy,
                profile.timbre,
                profile.texture,
                profile.formality,
                profile.emotional_baseline,
                profile.distinguishing_trait,
                profile.accent_marker,
                profile.voice_description,
                source_project,
                int(profile.is_narrator),
                0,
                "identified",
            ),
        )
        registry_id = cursor.lastrowid
        await self._insert_aliases(registry_id, profile.aliases)
        await self.conn.commit()
        return registry_id, True

    async def _insert_aliases(self, registry_id: int, aliases: list[str]) -> None:
        for alias in aliases:
            if alias.strip():
                await self.conn.execute(
                    "INSERT OR IGNORE INTO registry_aliases (registry_id, alias_name) VALUES (?, ?)",
                    (registry_id, alias.strip()),
                )

    async def update_voice_ref(
        self, registry_id: int, voice_ref_path: str, embedding: bytes | None = None
    ) -> None:
        await self.conn.execute(
            """UPDATE voice_registry
               SET voice_ref_path = ?, speaker_embedding = ?, status = 'voiced',
                   updated_at = datetime('now')
               WHERE id = ?""",
            (voice_ref_path, embedding, registry_id),
        )
        await self.conn.commit()

    async def set_vmatrix(
        self, registry_id: int, vmatrix_path: str | None
    ) -> None:
        """Record the captured vmatrix (.pt) for the approved voice.

        The vmatrix is the set of numbers inside Qwen3-TTS that IS the voice
        (ICL "8k matrix" ref_code + speaker embedding). An approved speaker is
        therefore: the reference WAV (voice_ref_path, what you listen to) plus
        this cached .pt (ponder the consistency across every utterance).
        """
        await self.conn.execute(
            """UPDATE voice_registry
               SET vmatrix_path = ?, updated_at = datetime('now')
               WHERE id = ?""",
            (vmatrix_path, registry_id),
        )
        await self.conn.commit()

    async def update_profile(self, registry_id: int, profile: CharacterProfile) -> None:
        await self.conn.execute(
            """UPDATE voice_registry
               SET aliases = ?, gender = ?, age_range = ?,
                   pitch_register = ?, pace_baseline = ?, energy = ?,
                   timbre = ?, texture = ?, formality = ?,
                   emotional_baseline = ?, distinguishing_trait = ?,
                   accent_marker = ?, voice_description = ?,
                   updated_at = datetime('now')
               WHERE id = ?""",
            (
                json.dumps(profile.aliases),
                profile.gender,
                profile.age_range,
                profile.pitch_register,
                profile.pace_baseline,
                profile.energy,
                profile.timbre,
                profile.texture,
                profile.formality,
                profile.emotional_baseline,
                profile.distinguishing_trait,
                profile.accent_marker,
                profile.voice_description,
                registry_id,
            ),
        )
        await self._insert_aliases(registry_id, profile.aliases)
        await self.conn.commit()

    async def link_to_project(
        self, registry_id: int, project_id: int, role_in_book: str = ""
    ) -> None:
        await self.conn.execute(
            "CREATE TABLE IF NOT EXISTS voice_projects ("
            "id INTEGER PRIMARY KEY,"
            "registry_id INTEGER NOT NULL REFERENCES voice_registry(id),"
            "project_id INTEGER NOT NULL,"
            "role_in_book TEXT,"
            "UNIQUE(registry_id, project_id)"
            ")"
        )
        await self.conn.execute(
            "INSERT OR IGNORE INTO voice_projects (registry_id, project_id, role_in_book) "
            "VALUES (?, ?, ?)",
            (registry_id, project_id, role_in_book),
        )
        await self.conn.commit()

    async def get_all_voiced(self) -> list[dict[str, Any]]:
        cursor = await self.conn.execute(
            "SELECT * FROM voice_registry WHERE status = 'voiced' ORDER BY character_name"
        )
        return [dict(r) for r in await cursor.fetchall()]

    async def get_all(self) -> list[dict[str, Any]]:
        cursor = await self.conn.execute(
            "SELECT * FROM voice_registry ORDER BY is_narrator DESC, character_name"
        )
        return [dict(r) for r in await cursor.fetchall()]

    async def search(self, query: str) -> list[dict[str, Any]]:
        like = f"%{query}%"
        cursor = await self.conn.execute(
            """SELECT vr.* FROM voice_registry vr
               LEFT JOIN registry_aliases ra ON ra.registry_id = vr.id
               WHERE vr.character_name LIKE ?
                  OR ra.alias_name LIKE ?
               GROUP BY vr.id
               ORDER BY vr.character_name""",
            (like, like),
        )
        return [dict(r) for r in await cursor.fetchall()]