meeting_recorder/database.py

"""
会议记录数据库 - SQLite 存储
"""

import os
import json
import sqlite3
from datetime import datetime
from typing import List, Optional, Dict
from contextlib import contextmanager

from sqlalchemy import create_engine, Column, String, Float, Integer, DateTime, JSON, Text
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy.pool import StaticPool

Base = declarative_base()


class Meeting(Base):
    """会议表"""
    __tablename__ = "meetings"

    meeting_id = Column(String, primary_key=True)
    title = Column(String, nullable=False)
    date = Column(DateTime, default=datetime.now)
    duration = Column(Float, default=0)  # 秒
    audio_path = Column(String, nullable=True)  # 音频文件路径
    transcript_path = Column(String, nullable=True)  # 转录文件路径
    speaker_count = Column(Integer, default=0)
    status = Column(String, default="pending")  # pending, processing, completed, failed
    segments_json = Column(Text, nullable=True)  # 存储片段 JSON
    brief_summary = Column(Text, nullable=True)  # 简要摘要
    detailed_summary = Column(Text, nullable=True)  # 详细摘要
    created_at = Column(DateTime, default=datetime.now)
    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)


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

    def __init__(self, db_path: str = None):
        if db_path is None:
            # 默认放在项目根目录
            project_root = os.path.dirname(os.path.abspath(__file__))
            db_path = os.path.join(project_root, "data", "meetings.db")
        
        # 确保目录存在
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        
        self.db_path = db_path
        self.engine = create_engine(
            f"sqlite:///{db_path}",
            connect_args={"check_same_thread": False},
            poolclass=StaticPool
        )
        Base.metadata.create_all(self.engine)
        self.Session = sessionmaker(bind=self.engine)

    @contextmanager
    def get_session(self):
        """获取数据库会话"""
        session = self.Session()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def create_meeting(
        self,
        meeting_id: str,
        title: str,
        duration: float = 0,
        audio_path: str = None,
        status: str = "pending"
    ) -> Meeting:
        """创建会议记录"""
        meeting = Meeting(
            meeting_id=meeting_id,
            title=title,
            duration=duration,
            audio_path=audio_path,
            status=status
        )
        with self.get_session() as session:
            session.add(meeting)
        return meeting

    def update_meeting(
        self,
        meeting_id: str,
        segments: List[Dict] = None,
        status: str = None,
        transcript_path: str = None,
        speaker_count: int = None,
        duration: float = None,
        brief_summary: str = None,
        detailed_summary: str = None,
        title: str = None
    ):
        """更新会议记录"""
        with self.get_session() as session:
            meeting = session.query(Meeting).filter(Meeting.meeting_id == meeting_id).first()
            if meeting:
                if title is not None:
                    meeting.title = title
                if segments is not None:
                    meeting.segments_json = json.dumps(segments, ensure_ascii=False)
                if status is not None:
                    meeting.status = status
                if transcript_path is not None:
                    meeting.transcript_path = transcript_path
                if speaker_count is not None:
                    meeting.speaker_count = speaker_count
                if duration is not None:
                    meeting.duration = duration
                if brief_summary is not None:
                    meeting.brief_summary = brief_summary
                if detailed_summary is not None:
                    meeting.detailed_summary = detailed_summary

    def get_meeting(self, meeting_id: str) -> Optional[Dict]:
        """获取单个会议记录"""
        with self.get_session() as session:
            meeting = session.query(Meeting).filter(Meeting.meeting_id == meeting_id).first()
            if meeting:
                return self._meeting_to_dict(meeting)
            return None

    def get_all_meetings(self, limit: int = 100) -> List[Dict]:
        """获取所有会议记录"""
        with self.get_session() as session:
            meetings = session.query(Meeting).order_by(Meeting.created_at.desc()).limit(limit).all()
            return [self._meeting_to_dict(m) for m in meetings]

    def search_meetings(self, keyword: str) -> List[Dict]:
        """搜索会议记录"""
        with self.get_session() as session:
            meetings = session.query(Meeting).filter(
                Meeting.title.like(f"%{keyword}%")
            ).order_by(Meeting.created_at.desc()).all()
            return [self._meeting_to_dict(m) for m in meetings]

    def delete_meeting(self, meeting_id: str):
        """删除会议记录"""
        with self.get_session() as session:
            meeting = session.query(Meeting).filter(Meeting.meeting_id == meeting_id).first()
            if meeting:
                session.delete(meeting)

    def _meeting_to_dict(self, meeting: Meeting) -> Dict:
        """转换为字典"""
        segments = []
        if meeting.segments_json:
            try:
                segments = json.loads(meeting.segments_json)
            except json.JSONDecodeError:
                segments = []

        return {
            "meeting_id": meeting.meeting_id,
            "title": meeting.title,
            "date": meeting.date.isoformat() if meeting.date else None,
            "duration": meeting.duration,
            "audio_path": meeting.audio_path,
            "transcript_path": meeting.transcript_path,
            "speaker_count": meeting.speaker_count,
            "status": meeting.status,
            "segments": segments,
            "brief_summary": meeting.brief_summary,
            "detailed_summary": meeting.detailed_summary,
            "created_at": meeting.created_at.isoformat() if meeting.created_at else None,
            "updated_at": meeting.updated_at.isoformat() if meeting.updated_at else None
        }

    def get_meetings_by_date(self, date: str) -> List[Dict]:
        """按日期获取会议"""
        with self.get_session() as session:
            meetings = session.query(Meeting).filter(
                Meeting.date >= datetime.fromisoformat(date),
                Meeting.date < datetime.fromisoformat(date + "T23:59:59")
            ).order_by(Meeting.date.desc()).all()
            return [self._meeting_to_dict(m) for m in meetings]


"""数据库单例实例"""
_db_instance = None

def init_db(db_path: str = None) -> Database:
    """初始化数据库（单例模式）"""
    global _db_instance
    if _db_instance is None:
        _db_instance = Database(db_path)
    return _db_instance


if __name__ == "__main__":
    # 测试
    db = init_db(":memory:")
    
    # 创建测试数据
    meeting_id = "test_001"
    db.create_meeting(
        meeting_id=meeting_id,
        title="测试会议",
        duration=300
    )
    
    db.update_meeting(
        meeting_id=meeting_id,
        segments=[
            {"start": 0, "end": 10, "speaker": 0, "text": "测试文本1"},
            {"start": 10, "end": 20, "speaker": 1, "text": "测试文本2"}
        ],
        status="completed",
        speaker_count=2
    )
    
    # 查询
    meeting = db.get_meeting(meeting_id)
    print(f"会议: {meeting['title']}")
    print(f"状态: {meeting['status']}")
    print(f"片段数: {len(meeting['segments'])}")