"""
数据库模块

提供 SQLite 数据库连接和会话管理功能。
"""

import os
import sqlite3
from pathlib import Path
from contextlib import contextmanager
from typing import Optional

# 默认数据库路径
DEFAULT_DB_PATH = Path(__file__).parent.parent / "data" / "chat.db"


def init_db(db_path: Optional[str] = None):
    """
    初始化数据库
    创建必要的表结构
    """
    if db_path is None:
        db_path = os.getenv("DB_PATH", str(DEFAULT_DB_PATH))
    
    # 确保数据目录存在
    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 创建会话表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS conversations (
            id TEXT PRIMARY KEY,
            title TEXT,
            model TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # 创建消息表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS messages (
            id TEXT PRIMARY KEY,
            conversation_id TEXT,
            role TEXT,
            content TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
        )
    """)
    
    # 创建文件表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS files (
            id TEXT PRIMARY KEY,
            conversation_id TEXT,
            filename TEXT,
            file_path TEXT,
            file_type TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
        )
    """)
    
    conn.commit()
    conn.close()
    
    print(f"[数据库] 初始化完成: {db_path}")


@contextmanager
def get_db(db_path: Optional[str] = None):
    """
    获取数据库连接的上下文管理器
    
    用法:
        with get_db() as db:
            cursor = db.execute("SELECT * FROM conversations")
            rows = cursor.fetchall()
    """
    if db_path is None:
        db_path = os.getenv("DB_PATH", str(DEFAULT_DB_PATH))
    
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()