Orbitin/src/feishu_v2.py

#!/usr/bin/env python3
"""
飞书表格 API 客户端模块 v2
支持数据库存储和2026年全年排班表
"""
import requests
import json
import os
import time
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
import logging
import hashlib

from src.schedule_database import ScheduleDatabase

logger = logging.getLogger(__name__)


class FeishuSheetsClient:
    """飞书表格 API 客户端"""
    
    def __init__(self, base_url: str, token: str, spreadsheet_token: str):
        """
        初始化客户端
        
        参数:
            base_url: 飞书 API 基础URL
            token: Bearer 认证令牌
            spreadsheet_token: 表格 token
        """
        self.base_url = base_url.rstrip('/')
        self.spreadsheet_token = spreadsheet_token
        self.headers = {
            'Authorization': f'Bearer {token}',
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        }
    
    def get_sheets_info(self) -> List[Dict]:
        """
        获取所有表格信息（sheet_id 和 title）
        
        返回:
            表格信息列表 [{'sheet_id': 'xxx', 'title': 'xxx'}, ...]
        """
        url = f'{self.base_url}/spreadsheets/{self.spreadsheet_token}/sheets/query'
        
        try:
            response = requests.get(url, headers=self.headers, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            if data.get('code') != 0:
                logger.error(f"飞书API错误: {data.get('msg')}")
                return []
            
            sheets = data.get('data', {}).get('sheets', [])
            result = []
            for sheet in sheets:
                result.append({
                    'sheet_id': sheet.get('sheet_id'),
                    'title': sheet.get('title')
                })
            
            logger.info(f"获取到 {len(result)} 个表格")
            return result
            
        except requests.exceptions.RequestException as e:
            logger.error(f"获取表格信息失败: {e}")
            return []
        except Exception as e:
            logger.error(f"解析表格信息失败: {e}")
            return []
    
    def get_sheet_data(self, sheet_id: str, range_: str = 'A:AF') -> Dict:
        """
        获取指定表格的数据
        
        参数:
            sheet_id: 表格ID
            range_: 数据范围，默认 A:AF (31列)
        
        返回:
            飞书API返回的原始数据，包含revision版本号
        """
        # 注意：获取表格数据使用 v2 API，而不是 v3
        url = f'{self.base_url.replace("/v3", "/v2")}/spreadsheets/{self.spreadsheet_token}/values/{sheet_id}!{range_}'
        params = {
            'valueRenderOption': 'ToString',
            'dateTimeRenderOption': 'FormattedString'
        }
        
        try:
            response = requests.get(url, headers=self.headers, params=params, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            if data.get('code') != 0:
                logger.error(f"飞书API错误: {data.get('msg')}")
                return {}
            
            return data.get('data', {})
            
        except requests.exceptions.RequestException as e:
            logger.error(f"获取表格数据失败: {e}")
            return {}
        except Exception as e:
            logger.error(f"解析表格数据失败: {e}")
            return {}


class ScheduleDataParser:
    """排班数据解析器（支持2026年全年排班表）"""
    
    @staticmethod
    def _parse_chinese_date(date_str: str) -> Optional[str]:
        """
        解析中文日期格式
        
        参数:
            date_str: 中文日期，如 "12月30日" 或 "12/30" 或 "12月1日" 或 "1月1日"
        
        返回:
            标准化日期字符串 "M月D日" (不补零)
        """
        if not date_str:
            return None
        
        # 如果是 "12/30" 格式
        if '/' in date_str:
            try:
                month, day = date_str.split('/')
                # 移除可能的空格和前导零
                month = month.strip().lstrip('0')
                day = day.strip().lstrip('0')
                return f"{int(month)}月{int(day)}日"
            except:
                return None
        
        # 如果是 "12月30日" 或 "1月1日" 格式
        if '月' in date_str and '日' in date_str:
            # 移除前导零，如 "01月01日" -> "1月1日"
            parts = date_str.split('月')
            if len(parts) == 2:
                month_part = parts[0].lstrip('0')
                day_part = parts[1].rstrip('日').lstrip('0')
                return f"{month_part}月{day_part}日"
            return date_str
        
        # 如果是 "12月1日" 格式（已经包含"日"字）
        if '月' in date_str:
            # 检查是否已经有"日"字
            if '日' not in date_str:
                return f"{date_str}日"
            return date_str
        
        return None
    
    @staticmethod
    def _find_date_column_index(headers: List[str], target_date: str) -> Optional[int]:
        """
        在表头中查找目标日期对应的列索引
        
        参数:
            headers: 表头行 ["姓名", "12月1日", "12月2日", ...]
            target_date: 目标日期 "12月30日"
        
        返回:
            列索引（从0开始），未找到返回None
        """
        if not headers or not target_date:
            return None
        
        # 标准化目标日期
        target_std = ScheduleDataParser._parse_chinese_date(target_date)
        if not target_std:
            return None
        
        # 遍历表头查找匹配的日期
        for i, header in enumerate(headers):
            header_std = ScheduleDataParser._parse_chinese_date(header)
            if header_std == target_std:
                return i
        
        return None
    
    def parse_monthly_sheet(self, values: List[List[str]], target_date: str) -> Dict:
        """
        解析月度表格数据（如12月表格）
        
        参数:
            values: 飞书表格返回的二维数组
            target_date: 目标日期（格式: "12月30日" 或 "12/30"）
        
        返回:
            排班信息字典
        """
        if not values or len(values) < 2:
            return {
                'day_shift': '',
                'night_shift': '',
                'day_shift_list': [],
                'night_shift_list': []
            }
        
        # 第一行是表头
        headers = values[0]
        date_column_index = self._find_date_column_index(headers, target_date)
        
        if date_column_index is None:
            logger.warning(f"未找到日期列: {target_date}")
            return {
                'day_shift': '',
                'night_shift': '',
                'day_shift_list': [],
                'night_shift_list': []
            }
        
        # 收集白班和夜班人员
        day_shift_names = []
        night_shift_names = []
        
        # 从第二行开始是人员数据
        for row in values[1:]:
            if len(row) <= date_column_index:
                continue
            
            name = row[0] if row else ''
            shift = row[date_column_index] if date_column_index < len(row) else ''
            
            if not name or not shift:
                continue
            
            if shift == '白':
                day_shift_names.append(name)
            elif shift == '夜':
                night_shift_names.append(name)
        
        # 格式化输出
        day_shift_str = '、'.join(day_shift_names) if day_shift_names else ''
        night_shift_str = '、'.join(night_shift_names) if night_shift_names else ''
        
        return {
            'day_shift': day_shift_str,
            'night_shift': night_shift_str,
            'day_shift_list': day_shift_names,
            'night_shift_list': night_shift_names
        }
    
    def parse_yearly_sheet(self, values: List[List[str]], target_date: str) -> Dict:
        """
        解析年度表格数据（如2026年排班表）
        
        参数:
            values: 飞书表格返回的二维数组
            target_date: 目标日期（格式: "12月30日" 或 "12/30"）
        
        返回:
            排班信息字典
        """
        if not values:
            return {
                'day_shift': '',
                'night_shift': '',
                'day_shift_list': [],
                'night_shift_list': []
            }
        
        # 查找目标月份的数据块
        target_month = target_date.split('月')[0] if '月' in target_date else ''
        if not target_month:
            logger.warning(f"无法从 {target_date} 提取月份")
            return {
                'day_shift': '',
                'night_shift': '',
                'day_shift_list': [],
                'night_shift_list': []
            }
        
        # 在年度表格中查找对应的月份块
        current_block_start = -1
        current_month = ''
        
        for i, row in enumerate(values):
            if not row:
                continue
            
            first_cell = str(row[0]) if row else ''
            
            # 检查是否是月份标题行，如 "福州港1月排班表"
            if '排班表' in first_cell and '月' in first_cell:
                # 提取月份数字
                for char in first_cell:
                    if char.isdigit():
                        month_str = ''
                        j = first_cell.index(char)
                        while j < len(first_cell) and first_cell[j].isdigit():
                            month_str += first_cell[j]
                            j += 1
                        if month_str:
                            current_month = month_str
                            current_block_start = i
                            break
            
            # 如果找到目标月份，检查下一行是否是表头行
            if current_month == target_month and i == current_block_start + 1:
                # 当前行是表头行
                headers = row
                date_column_index = self._find_date_column_index(headers, target_date)
                
                if date_column_index is None:
                    logger.warning(f"在年度表格中未找到日期列: {target_date}")
                    return {
                        'day_shift': '',
                        'night_shift': '',
                        'day_shift_list': [],
                        'night_shift_list': []
                    }
                
                # 收集人员数据（从表头行的下一行开始）
                day_shift_names = []
                night_shift_names = []
                
                for j in range(i + 1, len(values)):
                    person_row = values[j]
                    if not person_row:
                        # 遇到空行，继续检查下一行
                        continue
                    
                    # 检查是否是下一个月份块的开始
                    if person_row[0] and isinstance(person_row[0], str) and '排班表' in person_row[0] and '月' in person_row[0]:
                        break
                    
                    # 跳过星期行（第一列为空的行）
                    if not person_row[0]:
                        continue
                    
                    if len(person_row) <= date_column_index:
                        continue
                    
                    name = person_row[0] if person_row else ''
                    shift = person_row[date_column_index] if date_column_index < len(person_row) else ''
                    
                    if not name or not shift:
                        continue
                    
                    if shift == '白':
                        day_shift_names.append(name)
                    elif shift == '夜':
                        night_shift_names.append(name)
                
                # 格式化输出
                day_shift_str = '、'.join(day_shift_names) if day_shift_names else ''
                night_shift_str = '、'.join(night_shift_names) if night_shift_names else ''
                
                return {
                    'day_shift': day_shift_str,
                    'night_shift': night_shift_str,
                    'day_shift_list': day_shift_names,
                    'night_shift_list': night_shift_names
                }
        
        logger.warning(f"在年度表格中未找到 {target_month}月 的数据块")
        return {
            'day_shift': '',
            'night_shift': '',
            'day_shift_list': [],
            'night_shift_list': []
        }
    
    def parse(self, values: List[List[str]], target_date: str, sheet_title: str = '') -> Dict:
        """
        解析排班数据，自动判断表格类型
        
        参数:
            values: 飞书表格返回的二维数组
            target_date: 目标日期（格式: "12月30日" 或 "12/30"）
            sheet_title: 表格标题，用于判断表格类型
        
        返回:
            排班信息字典
        """
        # 根据表格标题判断表格类型
        if '年' in sheet_title and '排班表' in sheet_title:
            # 年度表格
            logger.info(f"使用年度表格解析器: {sheet_title}")
            return self.parse_yearly_sheet(values, target_date)
        else:
            # 月度表格
            logger.info(f"使用月度表格解析器: {sheet_title}")
            return self.parse_monthly_sheet(values, target_date)


class FeishuScheduleManagerV2:
    """飞书排班管理器 v2（使用数据库存储）"""
    
    def __init__(self, base_url: str = None, token: str = None, 
                 spreadsheet_token: str = None):
        """
        初始化管理器
        
        参数:
            base_url: 飞书API基础URL，从环境变量读取
            token: 飞书API令牌，从环境变量读取
            spreadsheet_token: 表格token，从环境变量读取
        """
        # 从环境变量读取配置
        self.base_url = base_url or os.getenv('FEISHU_BASE_URL', 'https://open.feishu.cn/open-apis/sheets/v3')
        self.token = token or os.getenv('FEISHU_TOKEN', '')
        self.spreadsheet_token = spreadsheet_token or os.getenv('FEISHU_SPREADSHEET_TOKEN', '')
        
        if not self.token or not self.spreadsheet_token:
            logger.warning("飞书配置不完整，请检查环境变量")
        
        self.client = FeishuSheetsClient(self.base_url, self.token, self.spreadsheet_token)
        self.parser = ScheduleDataParser()
        self.db = ScheduleDatabase()
    
    def _select_sheet_for_date(self, sheets: List[Dict], target_year_month: str) -> Optional[Dict]:
        """
        为指定日期选择最合适的表格
        
        参数:
            sheets: 表格列表
            target_year_month: 目标年月，格式 "2025-12"
        
        返回:
            选中的表格信息，未找到返回None
        """
        if not sheets:
            return None
        
        # 提取年份和月份
        year = target_year_month[:4]
        month = target_year_month[5:7]
        
        # 对于2026年，优先使用年度表格
        if year == '2026':
            # 查找年度表格，如 "2026年排班表"
            year_name = f"{year}年"
            for sheet in sheets:
                title = sheet.get('title', '')
                if year_name in title and '排班表' in title:
                    logger.info(f"找到2026年年度表格: {title}")
                    return sheet
        
        # 优先查找月份表格，如 "12月"
        month_name = f"{int(month)}月"
        for sheet in sheets:
            title = sheet.get('title', '')
            if month_name in title:
                logger.info(f"找到月份表格: {title}")
                return sheet
        
        # 查找年度表格，如 "2026年排班表"
        year_name = f"{year}年"
        for sheet in sheets:
            title = sheet.get('title', '')
            if year_name in title and '排班表' in title:
                logger.info(f"找到年度表格: {title}")
                return sheet
        
        # 如果没有找到匹配的表格，使用第一个表格
        logger.warning(f"未找到 {target_year_month} 的匹配表格，使用第一个表格: {sheets[0]['title']}")
        return sheets[0]
    
    def get_schedule_for_date(self, date_str: str) -> Dict:
        """
        获取指定日期的排班信息
        
        参数:
            date_str: 日期字符串，格式 "2025-12-30"
        
        返回:
            排班信息字典
        """
        try:
            # 解析日期
            dt = datetime.strptime(date_str, '%Y-%m-%d')
            # 生成两种格式的日期字符串，用于匹配不同表格
            target_date_mm_dd = dt.strftime('%m/%d')  # "01/01" 用于月度表格
            target_date_chinese = f"{dt.month}月{dt.day}日"  # "1月1日" 用于年度表格
            target_year_month = dt.strftime('%Y-%m')  # "2025-12"
            
            logger.info(f"获取 {date_str} 的排班信息 (格式: {target_date_mm_dd}/{target_date_chinese})")
            
            # 1. 首先尝试从数据库获取
            cached_schedule = self.db.get_schedule(date_str)
            if cached_schedule:
                logger.info(f"从数据库获取 {date_str} 的排班信息")
                return {
                    'day_shift': cached_schedule['day_shift'],
                    'night_shift': cached_schedule['night_shift'],
                    'day_shift_list': cached_schedule['day_shift_list'],
                    'night_shift_list': cached_schedule['night_shift_list']
                }
            
            # 2. 数据库中没有，需要从飞书获取
            logger.info(f"数据库中没有 {date_str} 的排班信息，从飞书获取")
            
            # 获取表格信息
            sheets = self.client.get_sheets_info()
            if not sheets:
                logger.error("未获取到表格信息")
                return {
                    'day_shift': '',
                    'night_shift': '',
                    'day_shift_list': [],
                    'night_shift_list': []
                }
            
            # 选择最合适的表格
            selected_sheet = self._select_sheet_for_date(sheets, target_year_month)
            if not selected_sheet:
                logger.error("未找到合适的表格")
                return {
                    'day_shift': '',
                    'night_shift': '',
                    'day_shift_list': [],
                    'night_shift_list': []
                }
            
            sheet_id = selected_sheet['sheet_id']
            sheet_title = selected_sheet['title']
            
            # 3. 获取表格数据
            sheet_data = self.client.get_sheet_data(sheet_id)
            if not sheet_data:
                logger.error("未获取到表格数据")
                return {
                    'day_shift': '',
                    'night_shift': '',
                    'day_shift_list': [],
                    'night_shift_list': []
                }
            
            values = sheet_data.get('valueRange', {}).get('values', [])
            revision = sheet_data.get('revision', 0)
            
            if not values:
                logger.error("表格数据为空")
                return {
                    'day_shift': '',
                    'night_shift': '',
                    'day_shift_list': [],
                    'night_shift_list': []
                }
            
            # 4. 检查表格是否有更新
            need_update = self.db.check_sheet_update(
                sheet_id, sheet_title, revision, {'values': values}
            )
            
            if not need_update and cached_schedule:
                # 表格无更新，且数据库中有缓存，直接返回
                logger.info(f"表格无更新，使用数据库缓存")
                return {
                    'day_shift': cached_schedule['day_shift'],
                    'night_shift': cached_schedule['night_shift'],
                    'day_shift_list': cached_schedule['day_shift_list'],
                    'night_shift_list': cached_schedule['night_shift_list']
                }
            
            # 5. 解析数据 - 根据表格类型选择合适的日期格式
            # 如果是年度表格，使用中文日期格式；否则使用mm/dd格式
            if '年' in sheet_title and '排班表' in sheet_title:
                target_date = target_date_chinese  # "1月1日"
            else:
                target_date = target_date_mm_dd  # "01/01"
            
            logger.info(f"使用日期格式: {target_date} 解析表格: {sheet_title}")
            result = self.parser.parse(values, target_date, sheet_title)
            
            # 6. 保存到数据库
            if result['day_shift'] or result['night_shift']:
                self.db.save_schedule(date_str, result, sheet_id, sheet_title)
                logger.info(f"已保存 {date_str} 的排班信息到数据库")
            
            return result
            
        except Exception as e:
            logger.error(f"获取排班信息失败: {e}")
            # 降级处理：返回空值
            return {
                'day_shift': '',
                'night_shift': '',
                'day_shift_list': [],
                'night_shift_list': []
            }
    
    def get_schedule_for_today(self) -> Dict:
        """获取今天的排班信息"""
        today = datetime.now().strftime('%Y-%m-%d')
        return self.get_schedule_for_date(today)
    
    def get_schedule_for_tomorrow(self) -> Dict:
        """获取明天的排班信息"""
        tomorrow = (datetime.now() + timedelta(days=1)).strftime('%Y-%m-%d')
        return self.get_schedule_for_date(tomorrow)
    
    def refresh_all_schedules(self, days: int = 30):
        """
        刷新未来指定天数的排班信息
        
        参数:
            days: 刷新未来多少天的排班信息
        """
        logger.info(f"开始刷新未来 {days} 天的排班信息")
        
        today = datetime.now()
        for i in range(days):
            date = (today + timedelta(days=i)).strftime('%Y-%m-%d')
            logger.info(f"刷新 {date} 的排班信息...")
            self.get_schedule_for_date(date)
        
        logger.info(f"排班信息刷新完成")


if __name__ == '__main__':
    # 测试代码
    import sys
    
    # 设置日志
    logging.basicConfig(level=logging.INFO)
    
    # 从环境变量读取配置
    manager = FeishuScheduleManagerV2()
    
    if len(sys.argv) > 1:
        date_str = sys.argv[1]
    else:
        date_str = datetime.now().strftime('%Y-%m-%d')
    
    print(f"获取 {date_str} 的排班信息...")
    schedule = manager.get_schedule_for_date(date_str)
    
    print(f"白班人员: {schedule['day_shift']}")
    print(f"夜班人员: {schedule['night_shift']}")
    print(f"白班列表: {schedule['day_shift_list']}")
    print(f"夜班列表: {schedule['night_shift_list']}")