Aletheia/server/audiobook/audio_cleanup.py

"""Durable post-delivery cleanup. A ready-job DELETE is the phone's acknowledgement.

Runs separately so it can be deployed without interrupting an active renderer.
Never deletes voice profiles, models, or another book's finished output.
"""
import hashlib
import json
import logging
import os
from pathlib import Path
import re
import shutil
import sqlite3
import time

ROOT = Path(os.environ.get('AUDIOBOOK_SERVICE_ROOT', '/home/sevenhill/apps/aletheia-audiobook'))
OUTPUTS = Path(os.environ.get('VOICESTUDIO_OUTPUTS', '/home/sevenhill/voicestudio-data/outputs'))
DATABASE = ROOT / 'audiobook-jobs/jobs.db3'
VOICE_DATABASE = OUTPUTS.parent / 'omnivoice.db'
LOG = logging.getLogger(__name__)


def connect():
    connection = sqlite3.connect(DATABASE, timeout=30)
    connection.row_factory = sqlite3.Row
    return connection


def initialize():
    with connect() as db:
        db.executescript('''
            CREATE TABLE IF NOT EXISTS delivered_audio_cleanup (
                id TEXT PRIMARY KEY, output_path TEXT NOT NULL,
                acknowledged INTEGER NOT NULL DEFAULT 0, cutoff REAL,
                artifact TEXT, completed INTEGER NOT NULL DEFAULT 0
            );
            CREATE TRIGGER IF NOT EXISTS acknowledge_delivered_audio
            BEFORE DELETE ON jobs WHEN OLD.status = 'ready'
            BEGIN
                INSERT INTO delivered_audio_cleanup(id, output_path, acknowledged, cutoff)
                VALUES(OLD.id, OLD.output_path, 1, unixepoch())
                ON CONFLICT(id) DO UPDATE SET acknowledged=1, cutoff=unixepoch();
            END;
        ''')


def remember_output(destination, source):
    """Persist outside the job directory, before the result becomes downloadable."""
    source = Path(source).resolve()
    if source.parent != OUTPUTS.resolve() or not re.fullmatch(r'audiobook_[a-zA-Z0-9_-]+\.m4b', source.name):
        raise ValueError('Unexpected VoiceStudio output')
    with connect() as db:
        db.execute('''INSERT INTO delivered_audio_cleanup(id, output_path, artifact)
            VALUES(?, ?, ?) ON CONFLICT(id) DO UPDATE SET artifact=excluded.artifact''',
            (Path(destination).parent.name, str(destination), source.name))


def digest(path):
    with path.open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').digest()


def capture_legacy_outputs():
    """Also covers a worker that was already running when cleanup was installed."""
    with connect() as db:
        rows = db.execute('''SELECT j.id, j.output_path FROM jobs j
            LEFT JOIN delivered_audio_cleanup c ON c.id=j.id
            WHERE j.status IN ('encoding','ready') AND c.artifact IS NULL''').fetchall()
    for row in rows:
        local = Path(row['output_path'])
        if not local.is_file() or not local.stat().st_size:
            continue
        candidates = [p for p in OUTPUTS.glob('audiobook_*.m4b')
                      if p.stat().st_size == local.stat().st_size]
        if not candidates:
            continue
        expected = digest(local)
        matches = [p for p in candidates if digest(p) == expected]
        if len(matches) == 1:
            remember_output(local, matches[0])


def renderers_idle():
    with connect() as db:
        if db.execute("SELECT 1 FROM jobs WHERE status IN ('queued','extracting','synthesizing','encoding') LIMIT 1").fetchone():
            return False
    # VoiceStudio's own jobs may belong to a different client.
    if not VOICE_DATABASE.is_file():
        return False
    with sqlite3.connect(f'file:{VOICE_DATABASE}?mode=ro', uri=True) as db:
        return not db.execute("SELECT 1 FROM jobs WHERE status IN ('queued','running','pending') LIMIT 1").fetchone()


def remove_artifact(name):
    if not re.fullmatch(r'audiobook_[a-zA-Z0-9_-]+\.m4b', name):
        raise ValueError('Unsafe artifact name')
    root = OUTPUTS.resolve()
    output = (root / name).resolve()
    work = (root / Path(name).stem).resolve()
    if output.parent != root or work.parent != root:
        raise ValueError('Artifact escapes output directory')
    cache = (root / 'longform_cache').resolve()
    concat = work / 'concat.txt'
    removed = 0
    if concat.is_file():
        for line in concat.read_text().splitlines():
            match = re.fullmatch(r"file '([^']+)'", line)
            if not match:
                raise ValueError('Unexpected concat entry')
            wav = Path(match[1]).resolve()
            if wav.parent != cache or not re.fullmatch(r'[a-f0-9]{20}\.wav', wav.name):
                raise ValueError('Unsafe cache path')
            if wav.is_file():
                removed += wav.stat().st_size
                wav.unlink()
    if output.is_file():
        removed += output.stat().st_size
        output.unlink()
    if work.exists():
        shutil.rmtree(work)
    return removed


def sweep():
    capture_legacy_outputs()
    with connect() as db:
        # Serialize against receiver job creation/state changes. A new render
        # cannot begin between the idle check and the last unlink.
        db.execute('BEGIN IMMEDIATE')
        pending = db.execute('''SELECT * FROM delivered_audio_cleanup
            WHERE acknowledged=1 AND completed=0 AND artifact IS NOT NULL''').fetchall()
        if not pending or not VOICE_DATABASE.is_file():
            return
        with sqlite3.connect(VOICE_DATABASE, timeout=30) as voice_db:
            voice_db.execute('BEGIN IMMEDIATE')
            if not renderers_idle():
                return
            for row in pending:
                removed = remove_artifact(row['artifact'])
                # Segment WAVs are a disposable shared cache. Only prune old
                # entries while idle; newly created files are never candidates.
                segment_root = OUTPUTS / 'longform_cache/segments'
                for path in segment_root.glob('*.wav'):
                    if path.is_symlink() or not re.fullmatch(r'[a-f0-9]{20}\.wav', path.name):
                        continue
                    if path.stat().st_mtime < row['cutoff']:
                        removed += path.stat().st_size
                        path.unlink()
                db.execute('UPDATE delivered_audio_cleanup SET completed=1 WHERE id=?', (row['id'],))
                LOG.info('Delivered audiobook %s cleaned: %d bytes freed', row['id'], removed)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    initialize()
    while True:
        try:
            sweep()
        except Exception:
            LOG.exception('Cleanup deferred; durable acknowledgement retained')
        time.sleep(5)