Orbitin/src/schedule_database.py

#!/usr/bin/env python3
"""
排班人员数据库模块
"""
import sqlite3
import os
import json
from datetime import datetime
from typing import List, Dict, Optional, Tuple
import hashlib


class ScheduleDatabase:
    """排班人员数据库"""
    
    def __init__(self, db_path: str = 'data/daily_logs.db'):
        """
        初始化数据库
        
        参数:
            db_path: 数据库文件路径
        """
        self.db_path = db_path
        self._ensure_directory()
        self.conn = self._connect()
        self._init_schema()
    
    def _ensure_directory(self):
        """确保数据目录存在"""
        data_dir = os.path.dirname(self.db_path)
        if data_dir and not os.path.exists(data_dir):
            os.makedirs(data_dir)
    
    def _connect(self) -> sqlite3.Connection:
        """连接数据库"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def _init_schema(self):
        """初始化表结构"""
        cursor = self.conn.cursor()
        
        # 创建排班人员表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS schedule_personnel (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                date TEXT NOT NULL,
                day_shift TEXT,
                night_shift TEXT,
                day_shift_list TEXT,  -- JSON数组
                night_shift_list TEXT, -- JSON数组
                sheet_id TEXT,
                sheet_title TEXT,
                data_hash TEXT,        -- 数据哈希，用于检测更新
                created_at TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_at TEXT DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(date)
            )
        ''')
        
        # 创建表格版本表（用于检测表格是否有更新）
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS sheet_versions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                sheet_id TEXT NOT NULL,
                sheet_title TEXT NOT NULL,
                revision INTEGER NOT NULL,
                data_hash TEXT,
                last_checked_at TEXT DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(sheet_id)
            )
        ''')
        
        # 索引
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_schedule_date ON schedule_personnel(date)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_schedule_sheet ON schedule_personnel(sheet_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_sheet_versions ON sheet_versions(sheet_id)')
        
        self.conn.commit()
    
    def _calculate_hash(self, data: Dict) -> str:
        """计算数据哈希值"""
        data_str = json.dumps(data, sort_keys=True, ensure_ascii=False)
        return hashlib.md5(data_str.encode('utf-8')).hexdigest()
    
    def check_sheet_update(self, sheet_id: str, sheet_title: str, revision: int, data: Dict) -> bool:
        """
        检查表格是否有更新
        
        参数:
            sheet_id: 表格ID
            sheet_title: 表格标题
            revision: 表格版本号
            data: 表格数据
        
        返回:
            True: 有更新，需要重新获取
            False: 无更新，可以使用缓存
        """
        cursor = self.conn.cursor()
        
        # 查询当前版本
        cursor.execute(
            'SELECT revision, data_hash FROM sheet_versions WHERE sheet_id = ?',
            (sheet_id,)
        )
        result = cursor.fetchone()
        
        if not result:
            # 第一次获取，记录版本
            data_hash = self._calculate_hash(data)
            cursor.execute('''
                INSERT INTO sheet_versions (sheet_id, sheet_title, revision, data_hash, last_checked_at)
                VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
            ''', (sheet_id, sheet_title, revision, data_hash))
            self.conn.commit()
            return True
        
        # 检查版本号或数据是否有变化
        old_revision = result['revision']
        old_hash = result['data_hash']
        new_hash = self._calculate_hash(data)
        
        if old_revision != revision or old_hash != new_hash:
            # 有更新，更新版本信息
            cursor.execute('''
                UPDATE sheet_versions 
                SET revision = ?, data_hash = ?, last_checked_at = CURRENT_TIMESTAMP
                WHERE sheet_id = ?
            ''', (revision, new_hash, sheet_id))
            self.conn.commit()
            return True
        
        # 无更新，更新检查时间
        cursor.execute('''
            UPDATE sheet_versions 
            SET last_checked_at = CURRENT_TIMESTAMP
            WHERE sheet_id = ?
        ''', (sheet_id,))
        self.conn.commit()
        return False
    
    def save_schedule(self, date: str, schedule_data: Dict, sheet_id: str = None, sheet_title: str = None) -> bool:
        """
        保存排班信息到数据库
        
        参数:
            date: 日期 (YYYY-MM-DD)
            schedule_data: 排班数据
            sheet_id: 表格ID
            sheet_title: 表格标题
        
        返回:
            是否成功
        """
        try:
            cursor = self.conn.cursor()
            
            # 准备数据
            day_shift = schedule_data.get('day_shift', '')
            night_shift = schedule_data.get('night_shift', '')
            day_shift_list = json.dumps(schedule_data.get('day_shift_list', []), ensure_ascii=False)
            night_shift_list = json.dumps(schedule_data.get('night_shift_list', []), ensure_ascii=False)
            data_hash = self._calculate_hash(schedule_data)
            
            # 使用 INSERT OR REPLACE 来更新已存在的记录
            cursor.execute('''
                INSERT OR REPLACE INTO schedule_personnel 
                (date, day_shift, night_shift, day_shift_list, night_shift_list, 
                 sheet_id, sheet_title, data_hash, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
            ''', (
                date, day_shift, night_shift, day_shift_list, night_shift_list,
                sheet_id, sheet_title, data_hash
            ))
            
            self.conn.commit()
            return True
            
        except sqlite3.Error as e:
            print(f"数据库错误: {e}")
            return False
    
    def get_schedule(self, date: str) -> Optional[Dict]:
        """
        获取指定日期的排班信息
        
        参数:
            date: 日期 (YYYY-MM-DD)
        
        返回:
            排班信息字典，未找到返回None
        """
        cursor = self.conn.cursor()
        cursor.execute(
            'SELECT * FROM schedule_personnel WHERE date = ?',
            (date,)
        )
        result = cursor.fetchone()
        
        if not result:
            return None
        
        # 解析JSON数组
        day_shift_list = json.loads(result['day_shift_list']) if result['day_shift_list'] else []
        night_shift_list = json.loads(result['night_shift_list']) if result['night_shift_list'] else []
        
        return {
            'date': result['date'],
            'day_shift': result['day_shift'],
            'night_shift': result['night_shift'],
            'day_shift_list': day_shift_list,
            'night_shift_list': night_shift_list,
            'sheet_id': result['sheet_id'],
            'sheet_title': result['sheet_title'],
            'updated_at': result['updated_at']
        }
    
    def get_schedule_by_range(self, start_date: str, end_date: str) -> List[Dict]:
        """
        获取日期范围内的排班信息
        
        参数:
            start_date: 开始日期 (YYYY-MM-DD)
            end_date: 结束日期 (YYYY-MM-DD)
        
        返回:
            排班信息列表
        """
        cursor = self.conn.cursor()
        cursor.execute('''
            SELECT * FROM schedule_personnel 
            WHERE date >= ? AND date <= ?
            ORDER BY date
        ''', (start_date, end_date))
        
        results = []
        for row in cursor.fetchall():
            day_shift_list = json.loads(row['day_shift_list']) if row['day_shift_list'] else []
            night_shift_list = json.loads(row['night_shift_list']) if row['night_shift_list'] else []
            
            results.append({
                'date': row['date'],
                'day_shift': row['day_shift'],
                'night_shift': row['night_shift'],
                'day_shift_list': day_shift_list,
                'night_shift_list': night_shift_list,
                'sheet_id': row['sheet_id'],
                'sheet_title': row['sheet_title'],
                'updated_at': row['updated_at']
            })
        
        return results
    
    def delete_old_schedules(self, before_date: str) -> int:
        """
        删除指定日期之前的排班记录
        
        参数:
            before_date: 日期 (YYYY-MM-DD)
        
        返回:
            删除的记录数
        """
        cursor = self.conn.cursor()
        cursor.execute(
            'DELETE FROM schedule_personnel WHERE date < ?',
            (before_date,)
        )
        deleted_count = cursor.rowcount
        self.conn.commit()
        return deleted_count
    
    def get_stats(self) -> Dict:
        """获取统计信息"""
        cursor = self.conn.cursor()
        
        cursor.execute('SELECT COUNT(*) FROM schedule_personnel')
        total = cursor.fetchone()[0]
        
        cursor.execute('SELECT MIN(date), MAX(date) FROM schedule_personnel')
        date_range = cursor.fetchone()
        
        cursor.execute('SELECT COUNT(DISTINCT sheet_id) FROM schedule_personnel')
        sheet_count = cursor.fetchone()[0]
        
        return {
            'total': total,
            'date_range': {'start': date_range[0], 'end': date_range[1]},
            'sheet_count': sheet_count
        }
    
    def close(self):
        """关闭连接"""
        if self.conn:
            self.conn.close()


if __name__ == '__main__':
    # 测试代码
    db = ScheduleDatabase()
    
    # 测试保存
    test_schedule = {
        'day_shift': '张勤、杨俊豪',
        'night_shift': '刘炜彬、梁启迟',
        'day_shift_list': ['张勤', '杨俊豪'],
        'night_shift_list': ['刘炜彬', '梁启迟']
    }
    
    success = db.save_schedule('2025-12-31', test_schedule, 'zcYLIk', '12月')
    print(f"保存成功: {success}")
    
    # 测试获取
    schedule = db.get_schedule('2025-12-31')
    print(f"获取结果: {schedule}")
    
    # 测试统计
    stats = db.get_stats()
    print(f"统计: {stats}")
    
    db.close()