#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复年度数据中年份字段的数据类型问题
将所有年份从数字类型转换为字符串类型
"""

import sys
import os
import json

# 添加父目录到路径以便导入模块
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from sqlalchemy.orm import sessionmaker
from database import SessionLocal
from models import Lab

def fix_year_data_types():
    """修复所有工程研究中心年度数据中的年份类型问题"""
    try:
        print("🔄 正在连接数据库...")
        db = SessionLocal()
        
        print("🔄 正在查询所有工程研究中心...")
        labs = db.query(Lab).all()
        
        print(f"✅ 找到 {len(labs)} 个工程研究中心")
        
        fixed_count = 0
        error_count = 0
        
        for lab in labs:
            try:
                if not lab.annual_data:
                    continue
                
                if isinstance(lab.annual_data, str):
                    # 解析JSON数据
                    data = json.loads(lab.annual_data)
                    
                    # 检查是否需要修复
                    needs_fix = False
                    for year_data in data:
                        year_field = year_data.get("归属年份")
                        if isinstance(year_field, int):
                            needs_fix = True
                            break
                    
                    if needs_fix:
                        print(f"🔧 修复工程研究中心: {lab.name}")
                        
                        # 修复所有年份字段
                        for year_data in data:
                            year_field = year_data.get("归属年份")
                            if isinstance(year_field, int):
                                year_data["归属年份"] = str(year_field)
                                print(f"   - 将年份 {year_field} 转换为字符串")
                        
                        # 保存修复后的数据
                        lab.annual_data = json.dumps(data, ensure_ascii=False)
                        fixed_count += 1
                
            except Exception as e:
                print(f"❌ 处理工程研究中心 {lab.name} 时出错: {str(e)}")
                error_count += 1
                continue
        
        # 提交更改
        if fixed_count > 0:
            print(f"\n💾 正在保存修复结果...")
            db.commit()
            print(f"✅ 修复完成!")
        else:
            print("ℹ️  没有需要修复的数据")
        
        print(f"📊 统计信息:")
        print(f"   - 修复的工程研究中心: {fixed_count}")
        print(f"   - 错误数量: {error_count}")
        print(f"   - 总计处理: {len(labs)}")
        
        db.close()
        return True
        
    except Exception as e:
        print(f"❌ 修复失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def verify_fix():
    """验证修复结果"""
    try:
        print("\n🔍 验证修复结果...")
        db = SessionLocal()
        
        # 检查西部优势矿产资源高效利用工程研究中心
        lab = db.query(Lab).filter(Lab.name == "西部优势矿产资源高效利用").first()
        
        if lab and lab.annual_data:
            data = json.loads(lab.annual_data)
            for i, year_data in enumerate(data):
                year_field = year_data.get("归属年份")
                print(f"📅 年度 {i+1}: {year_field} (类型: {type(year_field)})")
        
        db.close()
        return True
        
    except Exception as e:
        print(f"❌ 验证失败: {str(e)}")
        return False

def main():
    """主函数"""
    print("🚀 开始修复年度数据类型...")
    print("=" * 50)
    
    success = fix_year_data_types()
    
    if success:
        verify_fix()
    
    print("=" * 50)
    if success:
        print("🎉 修复完成！")
    else:
        print("💥 修复失败！")

if __name__ == "__main__":
    main() 