# -*- coding: utf-8 -*-
import os
import json
import random
import time
import csv
import httpx
from httpx import BasicAuth


class DataSetDownloader:
    def __init__(self):
        self.base_api_url = 'https://api.worldquantbrain.com'
        self.client = self.login()

    def login(self):
        """登录并返回客户端实例"""
        username, password = "jack0210_@hotmail.com", "!QAZ2wsx+0913"
        client = httpx.Client(auth=BasicAuth(username, password))

        try:
            response = client.post(f'{self.base_api_url}/authentication')
            print(f"登录状态: {response.status_code}")

            if response.status_code in [200, 201]:
                print("登录成功!")
                return client
            else:
                print(f"登录失败: {response.json()}")
                return None
        except Exception as e:
            print(f"登录过程中出现错误: {e}")
            return None

    def _debug_response(self, endpoint, data_set_id, offset=0, limit=20):
        """调试请求响应"""
        print(f"\n=== 调试请求: {endpoint} ===")
        url = f"{self.base_api_url}/{data_set_id}"
        params = self._build_params(data_set_id, offset, limit)

        response = self.client.get(url, params=params)

        if response.status_code == 200:
            data = response.json()
            print(f"count: {data.get('count')}")
            print(f"results 长度: {len(data.get('results', []))}")
            print(f"响应键: {list(data.keys())}")

    def _build_params(self, data_set_id, region, universe, offset=0, limit=50):
        """构建请求参数"""
        return {
            'dataset.id': data_set_id,
            'delay': 1,
            'instrumentType': 'EQUITY',
            'limit': limit,
            'offset': offset,
            'region': region,
            'universe': universe
        }

    def _process_item(self, item):
        """处理单个数据项"""
        return {
            'id': item.get('id', ''),
            'description': item.get('description', ''),
            'dataset_id': item.get('dataset', {}).get('id', ''),
            'dataset_name': item.get('dataset', {}).get('name', ''),
            'category_id': item.get('category', {}).get('id', ''),
            'category_name': item.get('category', {}).get('name', ''),
            'region': item.get('region', ''),
            'delay': item.get('delay', ''),
            'universe': item.get('universe', ''),
            'type': item.get('type', '')
        }

    def _process_data(self, raw_data):
        """批量处理数据"""
        return [self._process_item(item) for item in raw_data]

    def download_data_set(self, data_set_id, region, universe):
        endpoint = 'data-fields'
        """下载数据集"""
        # 检查登录状态
        if not self.client:
            print("❌ 客户端未初始化，无法下载数据")
            return

        # 调试请求
        self._debug_response(endpoint, data_set_id, offset=0, limit=20)

        # 获取数据总数
        url = f"{self.base_api_url}/{endpoint}"
        params = self._build_params(data_set_id, region, universe, limit=1)

        response = self.client.get(url, params=params)
        data = response.json()
        total_count = data.get('count', 0)

        print(f"📊 数据集总数: {total_count}")

        if total_count == 0:
            print("❌ 没有找到数据")
            return

        # 下载所有数据
        limit = 50
        all_data = []
        print("🚀 开始下载数据...")

        for offset in range(0, total_count, limit):
            time.sleep(random.uniform(1.0, 1.5))

            params = self._build_params(data_set_id, region, universe, offset, limit)  # 修正参数
            print(f"📥 下载进度: {offset}/{total_count} ({offset / total_count * 100:.1f}%)")

            results = []

            retry = 3
            while retry > 0:
                try:
                    response = self.client.get(url, params=params)

                    if response.status_code == 200:
                        data = response.json()
                        results = data.get('results', [])

                        print(f"✅ 本页获取到 {len(results)} 条记录")
                        all_data.extend(results)

                        # 成功时退出重试循环
                        break

                    else:
                        print(f"❌ 请求失败: {response.status_code}")
                        retry -= 1
                        if retry > 0:
                            print(f"🔄 重试中... ({retry}次剩余)")
                            time.sleep(random.uniform(2, 3))

                except Exception as e:
                    print(f"❌ 下载过程中出错: {e}")
                    retry -= 1
                    if retry > 0:
                        print(f"🔄 重试中... ({retry}次剩余)")
                        time.sleep(random.uniform(2, 3))

            # 如果重试用完仍失败，跳过当前offset继续下一个
            if retry == 0:
                print(f"⚠️  跳过 offset {offset}")
                continue

            if len(results) < limit:
                print("🎯 到达数据末尾")
                break

        time.sleep(random.uniform(4, 6))

        # 处理数据
        print("🔄 处理数据中...")
        processed_data = self._process_data(all_data)

        # 确保输出目录存在
        output_dir = 'reference_fields'
        os.makedirs(output_dir, exist_ok=True)

        # 保存为CSV
        output_file = os.path.join(output_dir, f"{data_set_id}_{region.lower()}_{universe.lower()}.csv")

        if processed_data:
            fieldnames = list(processed_data[0].keys())
            with open(output_file, 'w', encoding='utf-8', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(processed_data)

        print(f"💾 处理后的数据已保存到: {output_file}")
        print(f"🎉 总共处理了 {len(processed_data)} 条记录")


if __name__ == "__main__":
    downloader = DataSetDownloader()

    if downloader.client:
        plan_to_download = [
            {
                'id': 'analyst10',
                'region': 'USA',
                'universe': 'TOP3000'
            },
            {
                'id': 'analyst10',
                'region': 'GLB',
                'universe': 'TOP3000'
            }
        ]

        for item in plan_to_download:
            downloader.download_data_set(item['id'], item['region'], item['universe'])
            time.sleep(random.uniform(20, 30))
    else:
        print("❌ 登录失败，无法下载数据")