utils/database.py

"""
Database Module for User Authentication
========================================
SQLite-based user storage for Pharma K platform.
"""

import hashlib
import logging
import os
import secrets
import sqlite3
from datetime import datetime
from typing import Optional, List, Dict, Any
from pathlib import Path

logger = logging.getLogger(__name__)

#: 模块级标志：管理员未初始化的告警只打印一次，避免每次建连/重渲染刷屏。
_ADMIN_INIT_WARNED = False


def _salted_hash(password: str) -> str:
    """生成加盐密码哈希 ``"<salt_hex>$<hash_hex>"``（需求 5.4）。

    优先复用底座 :mod:`services.auth_service` 的实现，确保全平台哈希格式一致；
    在 ``platform/`` 目录未加入导入路径的运行环境下，退回到与之等价的最小
    标准库实现（PBKDF2-HMAC-SHA256），保证产出的哈希仍是加盐的、且与
    AuthService 的校验逻辑兼容（相同算法 / 迭代轮数 / 盐长度）。
    """
    try:
        from services.auth_service import hash_password  # type: ignore

        return hash_password(password)
    except Exception:  # pragma: no cover - 仅在 platform 目录未入路径时触发
        salt = secrets.token_bytes(16)
        dk = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, 200_000)
        return f"{salt.hex()}${dk.hex()}"

# Database file location.
# 数据目录可经环境变量 ``PHARMAK_DATA_DIR`` 覆盖（HF Spaces 等部署环境用），
# 未设置时回退到仓库内 ``data/`` 目录（本地开发）。
_DEFAULT_DB_DIR = Path(__file__).parent.parent / "data"
DB_DIR = Path(os.environ.get("PHARMAK_DATA_DIR") or _DEFAULT_DB_DIR)
DB_PATH = DB_DIR / "pharma_k.db"

#: 进程级标志：Dataset 拉取每个进程仅执行一次（启动时），避免每次建连重复下载。
_DB_PULLED = False


def _maybe_pull_once() -> None:
    """首次建连时从私有 HF Dataset 拉取数据库（若启用同步）。

    HF 免费档文件系统临时、重启即清空，故进程启动后先尝试从 Dataset 恢复最新 DB；
    未启用同步（本地开发 / 未配置环境变量）时静默 no-op。
    """
    global _DB_PULLED
    if _DB_PULLED:
        return
    _DB_PULLED = True  # 无论成败都只尝试一次，避免反复触网
    try:
        from utils import hf_storage

        if hf_storage.is_enabled():
            hf_storage.pull_db(DB_PATH)
    except Exception as exc:  # noqa: BLE001 - 同步不可用绝不影响本地存储
        logger.warning("启动期数据库同步拉取不可用：%s", exc)


def _push_db() -> None:
    """把本地数据库推回私有 HF Dataset（若启用同步）。写操作后调用。"""
    try:
        from utils import hf_storage

        if hf_storage.is_enabled():
            hf_storage.push_db(DB_PATH)
    except Exception as exc:  # noqa: BLE001
        logger.warning("数据库同步推送不可用：%s", exc)


def get_db_connection() -> sqlite3.Connection:
    """Get database connection, creating tables if needed."""
    # Ensure data directory exists
    DB_DIR.mkdir(parents=True, exist_ok=True)

    # 启动期：先从 Dataset 拉取最新 DB（仅一次），再建连建表。
    _maybe_pull_once()

    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    
    # Create tables if not exist
    _init_tables(conn)
    
    return conn


def _init_tables(conn: sqlite3.Connection):
    """Initialize database tables."""
    cursor = conn.cursor()
    
    # Users table
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            email TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            role TEXT DEFAULT 'user',
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            last_login DATETIME
        )
    ''')
    
    # Admin config table
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS admin_config (
            key TEXT PRIMARY KEY,
            value TEXT NOT NULL,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')

    # Per-user LLM config table（用户自配模型，密钥加密存储）。
    # api_key_enc 存 Fernet 密文（见 utils.crypto）；解密失败 / 未加密时按无配置处理。
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS user_llm_config (
            email TEXT PRIMARY KEY,
            provider TEXT NOT NULL DEFAULT '',
            model TEXT NOT NULL DEFAULT '',
            api_key_enc TEXT NOT NULL DEFAULT '',
            base_url TEXT NOT NULL DEFAULT '',
            enabled INTEGER NOT NULL DEFAULT 1,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')

    conn.commit()
    
    # Initialize default admin if not exists
    _init_default_admin(conn)
    
    # Initialize default LLM config
    _init_default_llm_config(conn)


def _init_default_admin(conn: sqlite3.Connection):
    """初始化管理员账户（需求 5.1）。

    管理员初始密码从环境变量 ``ADMIN_INIT_PASSWORD`` 读取，并以**加盐哈希**存储；
    未设置该环境变量时**不创建**任何默认账户，仅记录「管理员未初始化」日志
    （不含任何明文密码或哈希）。代码与日志中均不出现明文密码。
    """
    cursor = conn.cursor()
    cursor.execute("SELECT id FROM users WHERE email = 'admin'")

    if cursor.fetchone() is not None:
        return  # 管理员已存在，跳过

    password = os.environ.get("ADMIN_INIT_PASSWORD")
    if not password:
        global _ADMIN_INIT_WARNED
        if not _ADMIN_INIT_WARNED:
            logger.warning(
                "管理员未初始化：未设置 ADMIN_INIT_PASSWORD 环境变量，未创建默认管理员账户。"
                "如需管理员登录，请设置该环境变量后重启。"
            )
            _ADMIN_INIT_WARNED = True
        return

    password_hash = _salted_hash(password)
    cursor.execute('''
        INSERT INTO users (email, password_hash, role)
        VALUES (?, ?, 'admin')
    ''', ('admin', password_hash))
    conn.commit()
    logger.info("管理员账户已初始化（初始密码取自环境变量，未记录明文）。")


def _init_default_llm_config(conn: sqlite3.Connection):
    """Initialize default LLM configuration (admin-managed model config)."""
    cursor = conn.cursor()

    # 管理员级模型配置默认占位（管理员在后台设置后生效）。
    defaults = {
        "admin_llm_enabled": "0",          # 是否启用管理员配置的模型（"1"/"0"）
        "admin_llm_provider": "",          # 提供商 id（如 deepseek / custom）
        "admin_llm_model": "",             # 模型名（留空则用该提供商默认模型）
        "admin_llm_api_key": "",           # 管理员配置的密钥（存于本地 DB，不入公开库）
        "admin_llm_base_url": "",          # 仅 custom 提供商需要
    }
    for key, value in defaults.items():
        cursor.execute(
            "INSERT OR IGNORE INTO admin_config (key, value) VALUES (?, ?)",
            (key, value),
        )
    conn.commit()


# =============================================================================
# User Operations
# =============================================================================

def create_user(email: str, password_hash: str) -> bool:
    """Create a new user account."""
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute('''
            INSERT INTO users (email, password_hash, role)
            VALUES (?, ?, 'user')
        ''', (email, password_hash))
        conn.commit()
        conn.close()
        _push_db()  # 注册成功后同步回 Dataset（启用同步时），确保账号持久
        return True
    except sqlite3.IntegrityError:
        return False  # Email already exists


def get_user_by_email(email: str) -> Optional[Dict[str, Any]]:
    """Get user by email."""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM users WHERE email = ?", (email,))
    row = cursor.fetchone()
    conn.close()
    
    if row:
        return dict(row)
    return None


def update_last_login(email: str):
    """Update user's last login time."""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
        UPDATE users SET last_login = ?
        WHERE email = ?
    ''', (datetime.now().isoformat(), email))
    conn.commit()
    conn.close()


def get_all_users() -> List[Dict[str, Any]]:
    """Get all users (for admin)."""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT id, email, role, created_at, last_login FROM users ORDER BY created_at DESC")
    rows = cursor.fetchall()
    conn.close()
    return [dict(row) for row in rows]


# =============================================================================
# Admin Config Operations
# =============================================================================

def get_admin_config(key: str) -> Optional[str]:
    """Get admin configuration value."""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT value FROM admin_config WHERE key = ?", (key,))
    row = cursor.fetchone()
    conn.close()
    
    if row:
        return row['value']
    return None


def set_admin_config(key: str, value: str):
    """Set admin configuration value."""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
        INSERT OR REPLACE INTO admin_config (key, value, updated_at)
        VALUES (?, ?, ?)
    ''', (key, value, datetime.now().isoformat()))
    conn.commit()
    conn.close()
    _push_db()  # 管理员配置变更后同步回 Dataset（启用同步时）


def set_admin_configs(items: Dict[str, str]):
    """批量写入多个管理员配置项，单次事务 + 单次同步推送。

    供 :func:`set_default_llm_config` 等一次性写多键的场景使用，避免逐键多次触网。
    """
    conn = get_db_connection()
    cursor = conn.cursor()
    now = datetime.now().isoformat()
    for key, value in items.items():
        cursor.execute('''
            INSERT OR REPLACE INTO admin_config (key, value, updated_at)
            VALUES (?, ?, ?)
        ''', (key, value, now))
    conn.commit()
    conn.close()
    _push_db()  # 批量配置变更后单次同步回 Dataset（启用同步时）


def get_default_llm_config() -> Dict[str, str]:
    """返回管理员配置的 LLM 模型设置（admin-managed）。

    结构：``{enabled, provider, model, api_key, base_url}``。``enabled`` 为布尔，
    其余为字符串。未配置时各项为空 / False。该配置在用户未自行配置时作为兜底
    （优先级：用户 > 管理员 > 环境变量）。
    """
    return {
        "enabled": (get_admin_config("admin_llm_enabled") or "0") == "1",
        "provider": get_admin_config("admin_llm_provider") or "",
        "model": get_admin_config("admin_llm_model") or "",
        "api_key": get_admin_config("admin_llm_api_key") or "",
        "base_url": get_admin_config("admin_llm_base_url") or "",
    }


def set_default_llm_config(
    provider: str,
    api_key: str,
    *,
    model: str = "",
    base_url: str = "",
    enabled: bool = True,
):
    """保存管理员级 LLM 模型设置（写入本地 DB，不入公开仓库）。"""
    set_admin_configs({
        "admin_llm_enabled": "1" if enabled else "0",
        "admin_llm_provider": provider or "",
        "admin_llm_model": model or "",
        "admin_llm_api_key": api_key or "",
        "admin_llm_base_url": base_url or "",
    })


def clear_default_llm_config():
    """清除 / 停用管理员级模型配置。"""
    set_admin_config("admin_llm_enabled", "0")


# =============================================================================
# Per-User LLM Config Operations（用户自配模型，密钥加密持久化）
# =============================================================================

def get_user_llm_config(email: str) -> Optional[Dict[str, Any]]:
    """读取某用户持久化的自配模型设置（密钥已解密）。

    返回 ``{provider, model, api_key, base_url, enabled}``；无记录、加密不可用
    或密钥解密失败（secret 轮换 / 密文损坏）时返回 ``None``，调用方按"无持久配置"
    处理（退回仅会话存储）。
    """
    if not email:
        return None
    from utils import crypto

    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM user_llm_config WHERE email = ?", (email,))
    row = cursor.fetchone()
    conn.close()
    if not row:
        return None

    enc = row["api_key_enc"] or ""
    api_key = crypto.decrypt(enc) if enc else ""
    # 有密文但解密失败 → 视为失效配置，按无配置处理。
    if enc and not api_key:
        return None

    return {
        "provider": row["provider"] or "",
        "model": row["model"] or "",
        "api_key": api_key or "",
        "base_url": row["base_url"] or "",
        "enabled": bool(row["enabled"]),
    }


def set_user_llm_config(
    email: str,
    provider: str,
    api_key: str,
    *,
    model: str = "",
    base_url: str = "",
    enabled: bool = True,
) -> bool:
    """持久化某用户的自配模型设置（密钥**加密后**落库并同步）。

    仅在加密可用（已配置 ``PHARMAK_SECRET_KEY`` 且安装 ``cryptography``）时持久化，
    避免明文落库。加密不可用时返回 ``False``，调用方应退回仅会话存储。
    """
    if not email:
        return False
    from utils import crypto

    if not crypto.is_available():
        return False
    enc = crypto.encrypt(api_key) or "" if api_key else ""

    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute('''
        INSERT OR REPLACE INTO user_llm_config
            (email, provider, model, api_key_enc, base_url, enabled, updated_at)
        VALUES (?, ?, ?, ?, ?, ?, ?)
    ''', (
        email, provider or "", model or "", enc, base_url or "",
        1 if enabled else 0, datetime.now().isoformat(),
    ))
    conn.commit()
    conn.close()
    _push_db()  # 用户配置变更后同步回 Dataset（启用同步时）
    return True


def clear_user_llm_config(email: str):
    """删除某用户的持久化自配模型设置。"""
    if not email:
        return
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("DELETE FROM user_llm_config WHERE email = ?", (email,))
    conn.commit()
    conn.close()
    _push_db()