"""数据库模型和管理器"""
import json
from datetime import datetime
from typing import Dict

from peewee import Model, CharField, TextField, IntegerField, DateTimeField, BooleanField, ForeignKeyField
from playhouse.sqlite_ext import SqliteExtDatabase


class BaseModel(Model):
    """基础模型类"""
    pass


class Session(BaseModel):
    """聊天会话模型"""
    name = CharField(max_length=255, unique=True)
    model = CharField(max_length=50, default='deepseek-chat')
    system_prompt = TextField(default='你是一个乐于助人的助手。')
    created_at = DateTimeField(default=datetime.now)
    updated_at = DateTimeField(default=datetime.now)
    is_active = BooleanField(default=True)
    metadata_json = TextField(default='{}')

    @property
    def metadata(self) -> Dict:
        """获取元数据"""
        return json.loads(self.metadata_json) if self.metadata_json else {}

    @metadata.setter
    def metadata(self, value: Dict):
        """设置元数据"""
        self.metadata_json = json.dumps(value, ensure_ascii=False)

    def update_timestamp(self):
        """更新时间戳"""
        self.updated_at = datetime.now()
        self.save()

    class Meta:
        table_name = 'session'


class Message(BaseModel):
    """聊天消息模型"""
    session = ForeignKeyField(Session, backref='messages', on_delete='CASCADE')
    role = CharField(max_length=20, index=True)
    content = TextField()
    tokens = IntegerField(default=0)
    timestamp = DateTimeField(default=datetime.now, index=True)
    metadata_json = TextField(default='{}')

    @property
    def metadata(self) -> Dict:
        """获取元数据"""
        return json.loads(self.metadata_json) if self.metadata_json else {}

    @metadata.setter
    def metadata(self, value: Dict):
        """设置元数据"""
        self.metadata_json = json.dumps(value, ensure_ascii=False)

    class Meta:
        table_name = 'message'
        indexes = (
            (('session', 'timestamp'), False),
        )


class DatabaseManager:
    """数据库管理器"""

    def __init__(self, db_path: str = 'chat_sessions.db'):
        """
        初始化数据库管理器

        Args:
            db_path: 数据库文件路径
        """
        self.db = SqliteExtDatabase(
            db_path,
            pragmas={
                'journal_mode': 'wal',
                'cache_size': -1024 * 64,
                'foreign_keys': 1,
                'ignore_check_constraints': 0,
                'synchronous': 1
            }
        )
        self._init_database()

    def _init_database(self):
        """初始化数据库连接和表"""
        # 设置数据库
        BaseModel._meta.database = self.db
        Session._meta.database = self.db
        Message._meta.database = self.db

        # 连接数据库
        self.db.connect()

        # 创建表
        self.db.create_tables([Session, Message], safe=True)

    def close(self):
        """关闭数据库连接"""
        if not self.db.is_closed():
            self.db.close()