UmbraAudio/umbraaudio/assembly/__init__.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import logging
import sqlite3
from pathlib import Path

import numpy as np
import soundfile as sf

logger = logging.getLogger(__name__)


class ChapterQualityIssue:
    def __init__(self, severity: str, message: str, utterance_id: int | None = None) -> None:
        self.severity = severity
        self.message = message
        self.utterance_id = utterance_id

    def to_dict(self) -> dict:
        d = {"severity": self.severity, "message": self.message}
        if self.utterance_id is not None:
            d["utterance_id"] = self.utterance_id
        return d


class ChapterQualityChecker:
    def __init__(self, data_dir: Path, known_speakers: list[str] | None = None) -> None:
        self._data_dir = data_dir
        self._known_speakers = set(s.lower() for s in known_speakers) if known_speakers else None

    def check_chapter(self, chapter_id: int) -> list[ChapterQualityIssue]:
        issues: list[ChapterQualityIssue] = []
        db_path = self._data_dir / "project.db"
        db = sqlite3.connect(str(db_path))
        db.row_factory = sqlite3.Row

        utterances = db.execute(
            "SELECT id, utterance_index, text, type, speaker_name FROM utterances "
            "WHERE chapter_id = ? ORDER BY utterance_index",
            (chapter_id,),
        ).fetchall()
        db.close()

        if not utterances:
            issues.append(ChapterQualityIssue("error", "No utterances found for chapter"))
            return issues

        audio_dir = self._data_dir / "chapter_audio" / f"ch{chapter_id}"
        seen_indices: set[int] = set()

        for utt in utterances:
            idx = utt["utterance_index"]
            utt_id = utt["id"]
            text = utt["text"] or ""
            speaker = utt["speaker_name"] or ""
            utt_type = utt["type"] or "narration"

            if idx in seen_indices:
                issues.append(
                    ChapterQualityIssue("error", f"Duplicate utterance index {idx}", utt_id)
                )
            seen_indices.add(idx)

            if not text.strip():
                issues.append(
                    ChapterQualityIssue("warning", f"Empty text at index {idx}", utt_id)
                )

            if utt_type == "dialogue" and not speaker:
                issues.append(
                    ChapterQualityIssue("warning", f"Dialogue at index {idx} has no speaker", utt_id)
                )

            if self._known_speakers and speaker and speaker.lower() not in self._known_speakers:
                issues.append(
                    ChapterQualityIssue(
                        "info",
                        f"Speaker '{speaker}' at index {idx} not in known speakers list",
                        utt_id,
                    )
                )

            audio_file = audio_dir / f"utt_{idx:04d}.wav"
            if not audio_file.exists():
                issues.append(
                    ChapterQualityIssue("error", f"Missing audio for utterance {idx}", utt_id)
                )
            elif audio_file.stat().st_size < 100:
                issues.append(
                    ChapterQualityIssue("warning", f"Audio file too small for utterance {idx}", utt_id)
                )

        errors = sum(1 for i in issues if i.severity == "error")
        warnings = sum(1 for i in issues if i.severity == "warning")

        if errors == 0 and warnings == 0:
            logger.info("Chapter %d: all checks passed (%d utterances)", chapter_id, len(utterances))
        else:
            logger.warning(
                "Chapter %d: %d errors, %d warnings out of %d utterances",
                chapter_id, errors, warnings, len(utterances),
            )

        return issues

    def check_all_chapters(self) -> dict[int, list[ChapterQualityIssue]]:
        db_path = self._data_dir / "project.db"
        db = sqlite3.connect(str(db_path))
        db.row_factory = sqlite3.Row
        chapters = db.execute("SELECT id FROM chapters ORDER BY chapter_index").fetchall()
        db.close()

        results: dict[int, list[ChapterQualityIssue]] = {}
        for ch in chapters:
            issues = self.check_chapter(ch["id"])
            if issues:
                results[ch["id"]] = issues
        return results


class AudioStitcher:
    def __init__(
        self,
        silence_between_utterances_ms: int = 500,
        silence_between_paragraphs_ms: int = 2000,
        silence_between_chapters_ms: int = 4000,
        sample_rate: int = 24000,
    ) -> None:
        self.sr = sample_rate
        self.silence_utt = silence_between_utterances_ms
        self.silence_para = silence_between_paragraphs_ms
        self.silence_chap = silence_between_chapters_ms

    def _make_silence(self, duration_ms: int) -> np.ndarray:
        samples = int(self.sr * duration_ms / 1000)
        return np.zeros(samples, dtype=np.float32)

    def stitch_utterances(self, audio_files: list[Path], output_path: Path) -> Path:
        segments = []
        for f in audio_files:
            audio, _ = sf.read(str(f), dtype="float32")
            if audio.ndim > 1:
                audio = audio.mean(axis=1)
            segments.append(audio)
            segments.append(self._make_silence(self.silence_utt))

        if segments:
            segments.pop()

        combined = np.concatenate(segments) if segments else np.array([], dtype=np.float32)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        sf.write(str(output_path), combined, self.sr)
        return output_path

    def stitch_chapters(self, chapter_files: list[Path], output_path: Path) -> Path:
        segments = []
        for f in chapter_files:
            audio, _ = sf.read(str(f), dtype="float32")
            if audio.ndim > 1:
                audio = audio.mean(axis=1)
            segments.append(audio)
            segments.append(self._make_silence(self.silence_chap))

        if segments:
            segments.pop()

        combined = np.concatenate(segments) if segments else np.array([], dtype=np.float32)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        sf.write(str(output_path), combined, self.sr)
        return output_path