persistence/database.py

"""
数据持久化层 - SQLite数据库管理 (v2.0新增)

提供评估记录的持久化存储、查询和导出功能，
替代原有的session_state易失性存储。

模块组成：
- database.py: 数据库连接管理、建表语句、迁移版本控制
- models.py: AssessmentRecord数据模型定义
- repositories.py: CRUD操作封装

作者: Emotion Fusion System v2.0
日期: 2026-05-26
"""

from __future__ import annotations

import json
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


# ============================================================
# 全局配置
# ============================================================

DB_DIR = None  # 由init_database()设置
DB_FILENAME = "emotion_assessments.db"
_db_lock = threading.Lock()


def init_database(db_dir: str | Path | None = None) -> sqlite3.Connection:
    """
    初始化数据库连接（线程安全，懒加载单例模式）
    
    Args:
        db_dir: 数据库文件存放目录，None则使用默认位置
        
    Returns:
        sqlite3.Connection 对象
    """
    global DB_DIR
    
    if db_dir is not None:
        DB_DIR = Path(db_dir)
    elif DB_DIR is None:
        # 默认位置: 项目根目录下的 data/ 目录
        from config import DATA_DIR
        DB_DIR = DATA_DIR
    
    db_path = DB_DIR / DB_FILENAME
    
    with _db_lock:
        conn = sqlite3.connect(str(db_path), check_same_thread=False)
        conn.row_factory = sqlite3.Row  # 支持字典式访问
        _ensure_tables(conn)
        return conn


def get_connection() -> sqlite3.Connection | None:
    """获取当前数据库连接（如果已初始化）"""
    if DB_DIR is None:
        return None
    return init_database(DB_DIR)


# ============================================================
# 建表和迁移
# ============================================================

def _ensure_tables(conn: sqlite3.Connection) -> None:
    """确保所有必需的表存在，支持增量迁移"""
    
    # 检查并创建assessment_records表
    cursor = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name='assessment_records'"
    )
    
    if cursor.fetchone() is None:
        _create_assessment_table(conn)
        print("[persistence] OK: Created assessment_records table")
    else:
        # 检查是否需要添加新列（增量迁移）
        _migrate_schema(conn)
        print("[persistence] OK: Database schema ready")


def _create_assessment_table(conn: sqlite3.Connection) -> None:
    """创建主表 assessment_records"""
    conn.execute("""
        CREATE TABLE IF NOT EXISTS assessment_records (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            
            -- 时间戳
            timestamp TEXT NOT NULL DEFAULT (datetime('now')),
            created_at TEXT NOT NULL DEFAULT (datetime('now')),
            
            -- 可选患者标识（预留扩展）
            patient_id TEXT DEFAULT '',
            patient_name TEXT DEFAULT '',
            
            -- 四模态原始结果 (JSON序列化)
            text_result TEXT,
            speech_result TEXT,
            face_result TEXT,
            ecg_result TEXT,
            
            -- 融合结果
            final_emotion TEXT NOT NULL DEFAULT 'unknown',
            valence REAL,
            arousal REAL,
            confidence REAL DEFAULT 0.0,
            quality REAL DEFAULT 0.0,
            
            -- 元数据
            modality_count INTEGER DEFAULT 0,
            fusion_mode TEXT DEFAULT '',
            available_modalities TEXT DEFAULT '',
            
            -- 不确定性信息 (v2.0新增)
            uncertainty_level TEXT DEFAULT '',
            uncertainty_score REAL DEFAULT 0.5,
            suggestion TEXT DEFAULT '',
            
            -- 警告信息 (JSON数组)
            warnings TEXT DEFAULT '[]',
            
            -- 会话元数据
            session_info TEXT DEFAULT '{}',
            
            -- 唯一性约束（防重复）
            content_hash TEXT DEFAULT ''
        )
    """)
    
    # 创建索引以加速查询
    conn.execute("CREATE INDEX IF NOT EXISTS idx_assessment_time ON assessment_records(timestamp)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_assessment_patient ON assessment_records(patient_id)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_assessment_emotion ON assessment_records(final_emotion)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_assessment_hash ON assessment_records(content_hash)")
    
    conn.commit()


def _migrate_schema(conn: sqlite3.Connection) -> None:
    """增量迁移：添加新列而不破坏现有数据"""
    migrations = [
        ("uncertainty_level", "TEXT DEFAULT ''"),
        ("uncertainty_score", "REAL DEFAULT 0.5"),
        ("suggestion", "TEXT DEFAULT ''"),
        ("modality_count", "INTEGER DEFAULT 0"),
        ("fusion_mode", "TEXT DEFAULT ''"),
        ("available_modalities", "TEXT DEFAULT ''"),
        ("session_info", "TEXT DEFAULT '{}'"),
        ("content_hash", "TEXT DEFAULT ''"),
    ]
    
    for col_name, col_def in migrations:
        try:
            # 尝试选择该列，如果不存在会抛出异常
            conn.execute(f"SELECT {col_name} FROM assessment_records LIMIT 1")
        except sqlite3.OperationalError:
            try:
                conn.execute(f"ALTER TABLE assessment_records ADD COLUMN {col_name} {col_def}")
                print(f"[persistence] Added column: {col_name}")
            except Exception as e:
                pass  # 列可能已存在
    
    conn.commit()


# ============================================================
# 核心操作函数
# ============================================================

def save_assessment(
    conn: sqlite3.Connection,
    fusion_result: dict[str, Any],
    raw_results: list[dict[str, Any]] | None = None,
) -> int:
    """
    保存一次完整的融合评估结果到数据库。
    
    Args:
        conn: 数据库连接
        fusion_result: fuse_multimodal_va()返回的完整结果字典
        raw_results: 各模态的原始结果列表(可选)
        
    Returns:
        新插入记录的ID
    """
    import hashlib
    
    # 计算内容hash用于去重
    content_str = json.dumps({
        "v": fusion_result.get("valence"),
        "a": fusion_result.get("arousal"),
        "e": fusion_result.get("final_emotion"),
        "t": fusion_result.get("timestamp") or datetime.now().isoformat(),
    }, sort_keys=True)
    content_hash = hashlib.md5(content_str.encode()).hexdigest()[:16]
    
    # 准备各模态原始结果的JSON序列化
    def safe_json_dump(obj):
        return json.dumps(obj, ensure_ascii=False, default=str) if obj else None
    
    text_result = next((safe_json_dump(r) for r in (raw_results or []) if r.get("modality") == "text"), None)
    speech_result = next((safe_json_dump(r) for r in (raw_results or []) if r.get("modality") == "speech"), None)
    face_result = next((safe_json_dump(r) for r in (raw_results or []) if r.get("modality") == "face"), None)
    ecg_result = next((safe_json_dump(r) for r in (raw_results or []) if r.get("modality") == "ecg"), None)
    
    warnings_json = json.dumps(fusion_result.get("warnings", []), ensure_ascii=False)
    session_info = {
        "modalities_used": [r.get("modality") for r in (raw_results or []) if r.get("available") is True],
        "fusion_mode": fusion_result.get("fusion_mode", "adaptive_v2"),
        "saved_at": datetime.now().isoformat(),
    }
    
    cursor = conn.execute("""
        INSERT INTO assessment_records (
            timestamp, patient_id, patient_name,
            text_result, speech_result, face_result, ecg_result,
            final_emotion, valence, arousal, confidence, quality,
            modality_count, fusion_mode, available_modalities,
            uncertainty_level, uncertainty_score, suggestion,
            warnings, session_info, content_hash
        ) VALUES (
            ?, ?, ?,
            ?, ?, ?, ?,
            ?, ?, ?, ?, ?,
            ?, ?, ?,
            ?, ?, ?,
            ?, ?, ?
        )
    """, (
        datetime.now(timezone.utc).isoformat(),
        "",  # patient_id 预留
        "",  # patient_name 预留
        text_result,
        speech_result,
        face_result,
        ecg_result,
        fusion_result.get("final_emotion", "unknown"),
        fusion_result.get("valence"),
        fusion_result.get("arousal"),
        fusion_result.get("confidence", 0),
        fusion_result.get("quality", 0),
        len([r for r in (raw_results or []) if r.get("available") is True]),
        fusion_result.get("fusion_mode", "adaptive_v2"),
        ",".join([str(r.get("modality")) for r in (raw_results or []) if r.get("available")]),
        fusion_result.get("uncertainty_level", ""),
        fusion_result.get("uncertainty_score", 0.5),
        fusion_result.get("suggestion", ""),
        warnings_json,
        json.dumps(session_info, ensure_ascii=False),
        content_hash,
    ))
    
    record_id = int(cursor.lastrowid)
    conn.commit()
    
    return record_id


def query_assessments(
    conn: sqlite3.Connection,
    limit: int = 50,
    offset: int = 0,
    patient_id: str | None = None,
    emotion_filter: str | None = None,
    min_confidence: float = 0.0,
    start_date: str | None = None,
    end_date: str | None = None,
) -> tuple[list[dict[str, Any]], int]:
    """
    查询历史评估记录。
    
    Returns:
        (records_list, total_count) 元组
        records_list中的每条记录为dict，包含所有字段
    """
    conditions = []
    params = []
    
    if patient_id:
        conditions.append("patient_id = ?")
        params.append(patient_id)
    
    if emotion_filter and emotion_filter != "all":
        conditions.append("final_emotion = ?")
        params.append(emotion_filter)
    
    if min_confidence > 0:
        conditions.append("confidence >= ?")
        params.append(min_confidence)
    
    if start_date:
        conditions.append("timestamp >= ?")
        params.append(start_date)
    
    if end_date:
        conditions.append("timestamp <= ?")
        params.append(end_date)
    
    where_clause = " AND ".join(conditions) if conditions else "1=1"
    
    # 先查总数
    count_sql = f"SELECT COUNT(*) FROM assessment_records WHERE {where_clause}"
    total = int(conn.execute(count_sql, params).fetchone()[0])
    
    # 再查分页数据
    sql = f"""
        SELECT * FROM assessment_records 
        WHERE {where_clause}
        ORDER BY id DESC
        LIMIT ? OFFSET ?
    """
    rows = conn.execute(sql, [*params, limit, offset]).fetchall()
    
    records = [dict(row) for row in rows]
    
    return records, total


def export_to_csv(
    conn: sqlite3.Connection,
    output_path: str | Path,
    **query_filters,
) -> bool:
    """
    导出评估记录为CSV格式。
    
    Returns:
        True表示成功
    """
    import csv as csv_module
    
    records, _ = query_assessments(conn, limit=10000, **query_filters)
    
    if not records:
        return False
    
    fieldnames = [
        "id", "timestamp", "final_emotion", "valence", "arousal",
        "confidence", "quality", "modality_count", "uncertainty_level",
        "available_modalities"
    ]
    
    with open(output_path, "w", newline="", encoding="utf-8-sig") as f:
        writer = csv_module.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        
        for rec in records:
            # 只导出指定字段
            row = {k: rec.get(k, "") for k in fieldnames}
            writer.writerow(row)
    
    return True


def export_to_json(
    conn: sqlite3.Connection,
    output_path: str | Path,
    **query_filters,
) -> bool:
    """
    导出评估记录为JSON格式。
    
    Returns:
        True表示成功
    """
    records, _ = query_assessments(conn, limit=10000, **query_filters)
    
    if not records:
        return False
    
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(records, f, ensure_ascii=False, indent=2, default=str)
    
    return True


def get_statistics(conn: sqlite3.Connection) -> dict[str, Any]:
    """
    返回数据库统计概览：
    - 总记录数
    - 最近30天记录数
    - 平均valence/arousal/confidence
    - 情绪分布统计
    """
    stats = {}
    
    # 总记录数
    stats["total_records"] = int(conn.execute("SELECT COUNT(*) FROM assessment_records").fetchone()[0])
    
    # 最近30天
    thirty_days_ago = (datetime.now().replace(tzinfo=None).__sub__(__import__("datetime").timedelta(days=30))).isoformat()
    stats["recent_30d"] = int(conn.execute(
        "SELECT COUNT(*) FROM assessment_records WHERE timestamp >= ?", (thirty_days_ago,)
    ).fetchone()[0])
    
    # 平均值
    avg_row = conn.execute("""
        SELECT AVG(valence), AVG(arousal), AVG(confidence), AVG(quality)
        FROM assessment_records
    """).fetchone()
    
    stats["avg_valence"] = round(float(avg_row[0] or 0), 4) if avg_row[0] else 0
    stats["avg_arousal"] = round(float(avg_row[1] or 0), 4) if avg_row[1] else 0
    stats["avg_confidence"] = round(float(avg_row[2] or 0), 4) if avg_row[2] else 0
    stats["avg_quality"] = round(float(avg_row[3] or 0), 4) if avg_row[3] else 0
    
    # 情绪分布
    emotion_dist = conn.execute("""
        SELECT final_emotion, COUNT(*) as cnt 
        FROM assessment_records 
        GROUP BY final_emotion 
        ORDER BY cnt DESC
    """).fetchall()
    
    stats["emotion_distribution"] = {row["final_emotion"]: row["cnt"] for row in emotion_dist}
    
    # 不确定性分布
    uncertainty_dist = conn.execute("""
        SELECT COALESCE(uncertainty_level, 'unknown') as level, COUNT(*) as cnt
        FROM assessment_records
        GROUP BY level
        ORDER BY level
    """).fetchall()
    
    stats["uncertainty_distribution"] = {row["level"]: row["cnt"] for row in uncertainty_dist}
    
    return stats


def delete_old_records(
    conn: sqlite3.Connection,
    days_threshold: float = 90,
) -> int:
    """删除超过指定天数的旧记录，返回删除数量"""
    cutoff = (datetime.now().replace(tzinfo=None).__sub__(
        __import__("datetime").timedelta(days=days_threshold)
    )).isoformat()
    
    cursor = conn.execute("DELETE FROM assessment_records WHERE timestamp < ?", (cutoff,))
    deleted = cursor.rowcount
    conn.commit()
    
    return deleted


def backup_database(
    conn: sqlite3.Connection,
    output_path: str | Path | None = None,
) -> str | None:
    """
    备份整个数据库文件到指定路径。
    
    Returns:
        备份文件的路径，失败返回None
    """
    if DB_DIR is None:
        return None
    
    src_path = DB_DIR / DB_FILENAME
    
    if output_path is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = DB_DIR / f"backup_{timestamp}.db"
    
    import shutil
    shutil.copy2(str(src_path), str(output_path))
    
    return str(output_path)