storage/database.py

"""DuckDB persistence layer for all analysis runs."""
import json
import time
from pathlib import Path
from typing import Any, Optional

import duckdb
import pandas as pd


class Database:
    def __init__(self, path: str = "lie_detector.duckdb"):
        self.path = path
        self.con = duckdb.connect(path)
        self._init_schema()

    # ------------------------------------------------------------------
    # Schema
    # ------------------------------------------------------------------
    def _init_schema(self) -> None:
        self.con.execute("""
            CREATE TABLE IF NOT EXISTS runs (
                run_id        VARCHAR PRIMARY KEY,
                created_at    DOUBLE NOT NULL,
                source_path   VARCHAR,
                duration_sec  DOUBLE,
                unified_score DOUBLE,
                facial_score  DOUBLE,
                audio_score   DOUBLE,
                linguistic_score DOUBLE,
                facial_detail    JSON,
                audio_detail     JSON,
                linguistic_detail JSON,
                fusion_weights   JSON,
                metadata         JSON
            )
        """)

        self.con.execute("""
            CREATE TABLE IF NOT EXISTS frame_signals (
                run_id       VARCHAR,
                frame_idx    INTEGER,
                timestamp_s  DOUBLE,
                blink        BOOLEAN,
                asymmetric_smile DOUBLE,
                micro_expression BOOLEAN,
                eye_gaze_shift   DOUBLE,
                lip_compress     DOUBLE,
                head_nod         DOUBLE
            )
        """)

        self.con.execute("""
            CREATE TABLE IF NOT EXISTS sentence_signals (
                run_id          VARCHAR,
                sentence_idx    INTEGER,
                start_s         DOUBLE,
                end_s           DOUBLE,
                text            VARCHAR,
                pronoun_distance DOUBLE,
                neg_emotion      DOUBLE,
                cognitive_complexity DOUBLE,
                hedging          DOUBLE,
                coherence        DOUBLE,
                over_explanation DOUBLE,
                contradiction    DOUBLE,
                llm_score        DOUBLE,
                llm_reasoning    VARCHAR
            )
        """)

        self.con.execute("""
            CREATE TABLE IF NOT EXISTS audio_segments (
                run_id        VARCHAR,
                segment_idx   INTEGER,
                start_s       DOUBLE,
                end_s         DOUBLE,
                pitch_mean    DOUBLE,
                pitch_std     DOUBLE,
                speech_rate   DOUBLE,
                pause_count   INTEGER,
                energy_mean   DOUBLE,
                zcr_mean      DOUBLE,
                filler_count  INTEGER
            )
        """)

    # ------------------------------------------------------------------
    # Write
    # ------------------------------------------------------------------
    def save_run(
        self,
        run_id: str,
        source_path: str,
        duration_sec: float,
        unified_score: float,
        facial_score: float,
        audio_score: float,
        linguistic_score: float,
        facial_detail: dict,
        audio_detail: dict,
        linguistic_detail: dict,
        fusion_weights: dict,
        metadata: Optional[dict] = None,
    ) -> None:
        self.con.execute(
            """
            INSERT OR REPLACE INTO runs VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)
            """,
            [
                run_id,
                time.time(),
                source_path,
                duration_sec,
                unified_score,
                facial_score,
                audio_score,
                linguistic_score,
                json.dumps(facial_detail),
                json.dumps(audio_detail),
                json.dumps(linguistic_detail),
                json.dumps(fusion_weights),
                json.dumps(metadata or {}),
            ],
        )

    def save_frame_signals(self, run_id: str, frames: list[dict]) -> None:
        if not frames:
            return
        df = pd.DataFrame(frames)
        df.insert(0, "run_id", run_id)
        self.con.execute("INSERT INTO frame_signals SELECT * FROM df")

    def save_sentence_signals(self, run_id: str, sentences: list[dict]) -> None:
        if not sentences:
            return
        df = pd.DataFrame(sentences)
        df.insert(0, "run_id", run_id)
        self.con.execute("INSERT INTO sentence_signals SELECT * FROM df")

    def save_audio_segments(self, run_id: str, segments: list[dict]) -> None:
        if not segments:
            return
        df = pd.DataFrame(segments)
        df.insert(0, "run_id", run_id)
        self.con.execute("INSERT INTO audio_segments SELECT * FROM df")

    # ------------------------------------------------------------------
    # Read
    # ------------------------------------------------------------------
    def list_runs(self) -> pd.DataFrame:
        return self.con.execute(
            """
            SELECT run_id, created_at, source_path, duration_sec,
                   unified_score, facial_score, audio_score, linguistic_score
            FROM runs ORDER BY created_at DESC
            """
        ).df()

    def get_run(self, run_id: str) -> Optional[dict]:
        rows = self.con.execute(
            "SELECT * FROM runs WHERE run_id = ?", [run_id]
        ).fetchall()
        if not rows:
            return None
        cols = [d[0] for d in self.con.description]
        row = dict(zip(cols, rows[0]))
        for field in ("facial_detail", "audio_detail", "linguistic_detail",
                      "fusion_weights", "metadata"):
            if row.get(field):
                row[field] = json.loads(row[field])
        return row

    def get_frame_signals(self, run_id: str) -> pd.DataFrame:
        return self.con.execute(
            "SELECT * FROM frame_signals WHERE run_id = ? ORDER BY frame_idx",
            [run_id],
        ).df()

    def get_sentence_signals(self, run_id: str) -> pd.DataFrame:
        return self.con.execute(
            "SELECT * FROM sentence_signals WHERE run_id = ? ORDER BY sentence_idx",
            [run_id],
        ).df()

    def get_audio_segments(self, run_id: str) -> pd.DataFrame:
        return self.con.execute(
            "SELECT * FROM audio_segments WHERE run_id = ? ORDER BY segment_idx",
            [run_id],
        ).df()

    def close(self) -> None:
        self.con.close()