#!/usr/bin/env python3
"""
通过数据采集模块API导入测试数据到正式数据库
自动化脚本 - 一键导入所有测试数据
"""

import requests
import time
import json
from typing import Dict, Any

API_BASE = "http://localhost:8000/api/v1/data"

class DataImporter:
    def __init__(self):
        self.api_base = API_BASE
        self.uploaded_files = {}

    def upload_file(self, file_path: str, import_type: str) -> Dict[str, Any]:
        """上传文件并返回upload信息"""
        print(f"\n📤 正在上传: {file_path}")
        print(f"   类型: {import_type}")

        with open(file_path, 'rb') as f:
            files = {'file': (file_path.split('/')[-1], f)}
            data = {'import_type': import_type}

            response = requests.post(f"{self.api_base}/upload", files=files, data=data)

        if response.status_code == 200:
            result = response.json()
            upload_id = result['upload_id']
            total_rows = result['total_rows']
            field_mapping = result.get('field_mapping', {})
            print(f"   ✅ 上传成功 (upload_id: {upload_id}, {total_rows}条记录)")
            return {
                'upload_id': upload_id,
                'field_mapping': field_mapping,
                'total_rows': total_rows
            }
        else:
            print(f"   ❌ 上传失败: {response.text}")
            raise Exception(f"上传失败: {response.status_code}")

    def confirm_import(self, upload_id: str, import_type: str, field_mapping: Dict[str, str] = None) -> Dict[str, Any]:
        """确认导入并返回结果"""
        print(f"\n🔄 正在导入数据...")

        # 如果没有提供 field_mapping，创建一个空的映射
        if field_mapping is None:
            field_mapping = {}

        response = requests.post(
            f"{self.api_base}/confirm",
            json={
                'upload_id': upload_id,
                'import_type': import_type,
                'field_mapping': field_mapping
            }
        )

        if response.status_code == 200:
            result = response.json()
            success_rows = result['success_rows']
            failed_rows = result['failed_rows']
            total_rows = result['total_rows']
            task_id = result['task_id']

            print(f"   ✅ 导入完成!")
            print(f"   📊 总计: {total_rows}, 成功: {success_rows}, 失败: {failed_rows}")
            print(f"   🆔 任务ID: {task_id}")

            if failed_rows > 0:
                print(f"   ⚠️  错误详情:")
                for error in result['error_log'][:3]:  # 只显示前3个错误
                    print(f"      行{error['row']}: {error['errorMsg']}")

            return result
        else:
            print(f"   ❌ 导入失败: {response.text}")
            raise Exception(f"导入失败: {response.status_code}")

    def import_data_type(self, file_path: str, import_type: str, name: str):
        """导入单个数据类型"""
        print("\n" + "="*80)
        print(f"📦 正在导入: {name}")
        print("="*80)

        try:
            # 上传文件
            upload_info = self.upload_file(file_path, import_type)

            # 确认导入
            result = self.confirm_import(
                upload_info['upload_id'],
                import_type,
                upload_info['field_mapping']
            )

            print(f"\n✅ {name} 导入成功!")
            return True

        except Exception as e:
            print(f"\n❌ {name} 导入失败: {str(e)}")
            return False

    def check_import_history(self):
        """查看导入历史"""
        print("\n" + "="*80)
        print("📋 导入历史记录")
        print("="*80)

        response = requests.get(f"{self.api_base}/history?page=1&size=10")

        if response.status_code == 200:
            result = response.json()
            records = result.get('records', [])

            for record in records[-5:]:  # 显示最近5条
                status_icon = "✅" if record['status'] == 'success' else "⚠️"
                print(f"{status_icon} {record['import_type_name']}")
                print(f"   文件: {record['file_name']}")
                print(f"   状态: {record['status']} | 成功: {record['success_rows']} | 失败: {record['failed_rows']}")
                print(f"   时间: {record['created_at']}")
                print()
        else:
            print(f"❌ 查询失败: {response.text}")

    def run(self):
        """执行完整导入流程"""
        print("="*80)
        print("🚀 数据采集模块 - 测试数据导入工具")
        print("="*80)

        # 定义要导入的文件
        files_to_import = [
            {
                'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/主播信息导入模板.xlsx',
                'type': 'streamer',
                'name': '主播信息'
            },
            {
                'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/分成协议导入模板.xlsx',
                'type': 'contract',
                'name': '分成协议'
            },
            {
                'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/充值记录导入模板.xlsx',
                'type': 'recharge',
                'name': '充值记录'
            },
            {
                'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/税务申报导入模板.xlsx',
                'type': 'tax',
                'name': '税务申报'
            },
        ]

        success_count = 0
        failed_count = 0

        # 逐个导入
        for item in files_to_import:
            if self.import_data_type(item['file'], item['type'], item['name']):
                success_count += 1
            else:
                failed_count += 1

            time.sleep(1)  # 等待1秒避免请求过快

        # 查看历史
        self.check_import_history()

        # 总结
        print("\n" + "="*80)
        print("📊 导入完成统计")
        print("="*80)
        print(f"✅ 成功导入: {success_count} 个文件")
        print(f"❌ 导入失败: {failed_count} 个文件")
        print(f"📈 总计: {len(files_to_import)} 个文件")
        print("\n🎉 数据导入完成!")
        print("="*80)

        if failed_count == 0:
            print("\n💡 提示:")
            print("   现在您可以运行算法测试:")
            print("   python /Users/liulujian/Documents/code/deeprisk-claude-1/backend/scripts/test_revenue_algorithm.py")
        else:
            print("\n⚠️  部分导入失败，请检查错误信息并重试")

if __name__ == "__main__":
    try:
        importer = DataImporter()
        importer.run()
    except KeyboardInterrupt:
        print("\n\n⏹️  用户取消操作")
    except Exception as e:
        print(f"\n\n❌ 程序执行失败: {str(e)}")
        import traceback
        traceback.print_exc()