"""
修复document_service.py中的tenant_id查询问题

问题：get_file_config_by_id方法没有检查tenant_id，导致查询可能失败
解决方案：在查询中添加tenant_id检查
"""
import re
from pathlib import Path

def fix_document_service():
    """修复document_service.py中的查询逻辑"""
    file_path = Path("services/document_service.py")
    
    if not file_path.exists():
        print(f"[错误] 文件不存在: {file_path}")
        return False
    
    # 读取文件
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 查找get_file_config_by_id方法
    pattern = r'(def get_file_config_by_id\(self, file_id: int\) -> Optional\[Dict\]:.*?)(\s+sql = """.*?WHERE id = %s\s+AND state = 1\s+""".*?cursor\.execute\(sql, \(file_id,\)\))'
    
    match = re.search(pattern, content, re.DOTALL)
    
    if not match:
        print("[错误] 未找到get_file_config_by_id方法或查询语句")
        return False
    
    old_code = match.group(0)
    
    # 检查是否已经包含tenant_id
    if 'tenant_id' in old_code:
        print("[信息] 查询已经包含tenant_id检查，无需修复")
        return True
    
    # 生成新的代码
    new_sql = '''            sql = """
                SELECT id, name, file_path
                FROM f_polic_file_config
                WHERE id = %s
                AND tenant_id = %s
                AND state = 1
            """
            # 获取tenant_id（从环境变量或请求中获取）
            tenant_id = self.tenant_id if self.tenant_id else os.getenv('TENANT_ID', '1')
            try:
                tenant_id = int(tenant_id)
            except (ValueError, TypeError):
                tenant_id = 1  # 默认值
            
            cursor.execute(sql, (file_id, tenant_id))'''
    
    # 替换
    new_code = re.sub(
        r'sql = """.*?WHERE id = %s\s+AND state = 1\s+""".*?cursor\.execute\(sql, \(file_id,\)\)',
        new_sql,
        old_code,
        flags=re.DOTALL
    )
    
    new_content = content.replace(old_code, new_code)
    
    # 检查是否需要导入os
    if 'import os' not in new_content and 'os.getenv' in new_content:
        # 在文件开头添加import os（如果还没有）
        if 'from dotenv import load_dotenv' in new_content:
            new_content = new_content.replace('from dotenv import load_dotenv', 'from dotenv import load_dotenv\nimport os')
        elif 'import pymysql' in new_content:
            new_content = new_content.replace('import pymysql', 'import pymysql\nimport os')
        else:
            # 在文件开头添加
            lines = new_content.split('\n')
            import_line = 0
            for i, line in enumerate(lines):
                if line.startswith('import ') or line.startswith('from '):
                    import_line = i + 1
            lines.insert(import_line, 'import os')
            new_content = '\n'.join(lines)
    
    # 写回文件
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(new_content)
    
    print("[成功] 已修复get_file_config_by_id方法，添加了tenant_id检查")
    return True


if __name__ == "__main__":
    print("="*70)
    print("修复document_service.py中的tenant_id查询问题")
    print("="*70)
    
    if fix_document_service():
        print("\n修复完成！")
        print("\n注意：")
        print("1. 请确保.env文件中配置了TENANT_ID")
        print("2. 或者确保应用程序在调用时正确传递tenant_id")
        print("3. 建议在app.py中从请求中获取tenant_id并传递给document_service")
    else:
        print("\n修复失败，请手动检查代码")