UmbraAudio/tests/test_voices/test_registry.py

from __future__ import annotations

from pathlib import Path

import pytest

from umbraaudio.analysis import CharacterProfile
from umbraaudio.voices.registry import VoiceRegistry


@pytest.mark.asyncio
async def test_connect_creates_db(registry_db_path: Path):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    assert registry_db_path.exists()
    await registry.close()


@pytest.mark.asyncio
async def test_find_or_create_new(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    reg_id, created = await registry.find_or_create(character_profile, source_project="test")
    assert created is True
    assert reg_id > 0
    await registry.close()


@pytest.mark.asyncio
async def test_find_or_create_duplicate(
    registry_db_path: Path, character_profile: CharacterProfile
):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    reg_id1, created1 = await registry.find_or_create(character_profile, source_project="test")
    reg_id2, created2 = await registry.find_or_create(character_profile, source_project="test2")
    assert created1 is True
    assert created2 is False
    assert reg_id1 == reg_id2
    await registry.close()


@pytest.mark.asyncio
async def test_find_by_name(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    reg_id, _ = await registry.find_or_create(character_profile)
    found = await registry.find_by_name("John Smith")
    assert found is not None
    assert found["id"] == reg_id
    await registry.close()


@pytest.mark.asyncio
async def test_find_by_alias(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    await registry.find_or_create(character_profile)
    found = await registry.find_by_name("John")
    assert found is not None
    assert found["character_name"] == "John Smith"
    await registry.close()


@pytest.mark.asyncio
async def test_find_by_unknown_name(registry_db_path: Path):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    found = await registry.find_by_name("Nobody")
    assert found is None
    await registry.close()


@pytest.mark.asyncio
async def test_update_voice_ref(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    reg_id, _ = await registry.find_or_create(character_profile)
    await registry.update_voice_ref(reg_id, "/path/to/ref.wav")
    found = await registry.find_by_name("John Smith")
    assert found is not None
    assert found["voice_ref_path"] == "/path/to/ref.wav"
    assert found["status"] == "voiced"
    await registry.close()


@pytest.mark.asyncio
async def test_update_profile(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    reg_id, _ = await registry.find_or_create(character_profile)
    updated = CharacterProfile(
        name="John Smith",
        aliases=["John", "Johnny"],
        gender="male",
        age_range="adult",
        energy="high",
        timbre="bright",
    )
    await registry.update_profile(reg_id, updated)
    found = await registry.find_by_name("John Smith")
    assert found is not None
    assert found["energy"] == "high"
    assert found["timbre"] == "bright"
    # new alias should be findable
    found_by_new = await registry.find_by_name("Johnny")
    assert found_by_new is not None
    await registry.close()


@pytest.mark.asyncio
async def test_get_all_and_voiced(
    registry_db_path: Path, character_profile: CharacterProfile, narrator_profile: CharacterProfile
):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    await registry.find_or_create(character_profile)
    await registry.find_or_create(narrator_profile)
    all_chars = await registry.get_all()
    assert len(all_chars) >= 2
    voiced = await registry.get_all_voiced()
    assert len(voiced) == 0  # none voiced yet
    await registry.close()


@pytest.mark.asyncio
async def test_search(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    await registry.find_or_create(character_profile)
    results = await registry.search("Smith")
    assert len(results) >= 1
    assert results[0]["character_name"] == "John Smith"
    results_none = await registry.search("xyzzy")
    assert len(results_none) == 0
    await registry.close()


@pytest.mark.asyncio
async def test_link_to_project(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    reg_id, _ = await registry.find_or_create(character_profile)
    await registry.link_to_project(reg_id, 1, role_in_book="Protagonist")
    await registry.link_to_project(reg_id, 2)
    await registry.close()


@pytest.mark.asyncio
async def test_set_vmatrix(registry_db_path: Path, character_profile: CharacterProfile):
    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    reg_id, _ = await registry.find_or_create(character_profile)
    await registry.set_vmatrix(reg_id, "/umbra/voice_prompts/nanny.pt")
    found = await registry.find_by_name("John Smith")
    assert found is not None
    assert found["vmatrix_path"] == "/umbra/voice_prompts/nanny.pt"
    await registry.close()


@pytest.mark.asyncio
async def test_migrate_adds_vmatrix_column(
    registry_db_path: Path, character_profile: CharacterProfile
):
    """A DB created by the previous schema (no vmatrix_path) gains the column."""

    legacy_schema = """
    CREATE TABLE 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,
        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'))
    );
    """
    import aiosqlite

    conn = await aiosqlite.connect(str(registry_db_path))
    await conn.executescript(legacy_schema)
    await conn.commit()
    await conn.close()

    registry = VoiceRegistry(registry_db_path)
    await registry.connect()
    cursor = await registry.conn.execute("PRAGMA table_info(voice_registry)")
    cols = {row["name"] for row in await cursor.fetchall()}
    assert "vmatrix_path" in cols
    reg_id, _ = await registry.find_or_create(character_profile)
    await registry.set_vmatrix(reg_id, "/umbra/voice_prompts/nanny.pt")
    found = await registry.find_by_name("John Smith")
    assert found["vmatrix_path"] == "/umbra/voice_prompts/nanny.pt"
    await registry.close()