import requests
import pandas as pd
from datetime import datetime
import sys
import os
import redis
import json
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import time

# 添加项目根目录到路径，便于导入scripts.config
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

# 读取雪球headers和Redis配置
try:
    from src.scripts.config import XUEQIU_HEADERS
    from src.valuation_analysis.config import REDIS_CONFIG
except ImportError:
    XUEQIU_HEADERS = {
        'User-Agent': 'Mozilla/5.0',
        'Cookie': '',  # 需要填写雪球cookie
    }
    REDIS_CONFIG = {
        'host': 'localhost',
        'port': 6379,
        'db': 0,
        'password': None
    }

REDIS_KEY = 'xq_hk_stock_changes_latest'  # 存放港股行情的主键

# 条件导入代理管理器
proxy_manager = None
try:
    from src.scripts.ProxyIP import EnhancedProxyManager
    proxy_manager = EnhancedProxyManager()
except ImportError:
    print("代理管理器导入失败，将使用直接请求模式")


def get_redis_conn():
    """获取Redis连接"""
    pool = redis.ConnectionPool(
        host=REDIS_CONFIG['host'],
        port=REDIS_CONFIG['port'],
        db=REDIS_CONFIG.get('db', 0),
        password=REDIS_CONFIG.get('password', None),
        decode_responses=True
    )
    return redis.Redis(connection_pool=pool)


def fetch_and_store_hk_stock_data(page_size=90, max_workers=10, use_proxy=False):
    """
    批量采集雪球港股所有股票的最新行情数据，并保存到Redis。
    使用线程池并行请求，提高采集效率。
    
    :param page_size: 每页采集数量
    :param max_workers: 线程池最大工作线程数
    :param use_proxy: 是否使用代理（默认False）
    """
    base_url = 'https://stock.xueqiu.com/v5/stock/screener/quote/list.json'
    headers = XUEQIU_HEADERS

    all_data = []
    data_lock = threading.Lock()  # 线程安全锁

    def fetch_page_data(page):
        """获取单页数据的函数"""
        params = {
            'page': page,
            'size': page_size,
            'order': 'desc',
            'order_by': 'dividend_yield',
            'market': 'HK',  # 港股市场
            'type': 'hk'     # 港股类型
        }
        
        try:
            # 根据配置选择是否使用代理
            if use_proxy and proxy_manager:
                response = proxy_manager.request_with_proxy('get', base_url, headers=headers, params=params)
            else:
                response = requests.get(base_url, headers=headers, params=params, timeout=10)
                
            if response.status_code == 200:
                data = response.json()
                page_data = data['data']['list']
                
                # 线程安全地添加数据
                with data_lock:
                    all_data.extend(page_data)
                
                print(f"成功采集港股第 {page} 页数据，获取 {len(page_data)} 条记录")
                return len(page_data)
            else:
                print(f"请求港股数据第 {page} 页失败，状态码：{response.status_code}")
                return 0
        except Exception as e:
            print(f"请求港股数据第 {page} 页异常：{e}")
            return 0

    # 先获取总页数
    params = {
        'page': 1,
        'size': page_size,
        'order': 'desc',
        'order_by': 'dividend_yield',
        'market': 'HK',
        'type': 'hk'
    }
    
    try:
        # 根据配置选择是否使用代理
        if use_proxy and proxy_manager:
            response = proxy_manager.request_with_proxy('get', base_url, headers=headers, params=params)
        else:
            response = requests.get(base_url, headers=headers, params=params, timeout=10)
            
        if response.status_code != 200:
            print(f"请求港股数据失败，状态码：{response.status_code}")
            return pd.DataFrame()
        
        data = response.json()
        total_count = data['data']['count']
        total_pages = (total_count // page_size) + 1
        
        print(f"开始采集港股数据，共 {total_pages} 页，总计 {total_count} 条记录")
        
        # 使用线程池并行采集数据
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = []
            
            # 提交所有页面的采集任务
            for page in range(1, total_pages + 1):
                future = executor.submit(fetch_page_data, page)
                futures.append(future)
            
            # 等待所有任务完成
            print(f"正在并行采集港股数据，使用 {max_workers} 个线程...")
            start_time = time.time()
            
            completed_count = 0
            for future in as_completed(futures):
                completed_count += 1
                try:
                    result = future.result()
                    if result > 0:
                        print(f"进度: {completed_count}/{len(futures)} 页完成")
                except Exception as e:
                    print(f"采集任务异常：{e}")
            
            end_time = time.time()
            print(f"港股数据采集完成，耗时: {end_time - start_time:.2f} 秒")
            
    except Exception as e:
        print(f"获取港股总页数失败：{e}")
        return pd.DataFrame()

    # 转换为 DataFrame
    df = pd.DataFrame(all_data)

    if not df.empty:
        df['fetch_time'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        # 存入Redis，使用hash结构，key为symbol，value为json字符串
        r = get_redis_conn()
        pipe = r.pipeline()
        
        # 先清空旧数据
        r.delete(REDIS_KEY)
        
        print(f"正在将 {len(df)} 条港股记录写入Redis...")
        
        for _, row in df.iterrows():
            symbol = row.get('symbol')
            if not symbol:
                continue
            # 只保留必要字段，也可直接存row.to_dict()
            value = row.to_dict()
            pipe.hset(REDIS_KEY, symbol, json.dumps(value, ensure_ascii=False))
        
        pipe.execute()
        print(f"成功将港股数据写入Redis哈希 {REDIS_KEY}，共{len(df)}条记录。")
        
        # 返回DataFrame供其他脚本使用
        return df
    else:
        print("未获取到任何港股数据。")
        return pd.DataFrame()


def format_hk_stock_code(stock_code):
    """
    统一港股代码格式，支持0700.HK、HK0700、00700等
    返回雪球格式（如0700.HK）和Redis存储格式（如00700）
    """
    stock_code = stock_code.upper()
    if '.HK' in stock_code:
        code = stock_code.replace('.HK', '')
        # 确保是5位数字格式
        if code.isdigit():
            return f'{code.zfill(5)}.HK', code.zfill(5)
        else:
            return stock_code, code.zfill(5)
    elif stock_code.startswith('HK'):
        code = stock_code[2:]
        if code.isdigit():
            return f'{code.zfill(5)}.HK', code.zfill(5)
        else:
            return stock_code, code.zfill(5)
    else:
        # 假设是纯数字，添加.HK后缀并确保5位格式
        if stock_code.isdigit():
            return f'{stock_code.zfill(5)}.HK', stock_code.zfill(5)
        else:
            return stock_code, stock_code


def get_hk_stock_realtime_info_from_redis(stock_code):
    """
    根据港股代码从Redis查询实时行情，并封装为指定结构。
    :param stock_code: 支持0700.HK、HK0700、00700等格式
    :return: dict or None
    """
    _, redis_code = format_hk_stock_code(stock_code)
    r = get_redis_conn()
    value = r.hget(REDIS_KEY, redis_code)
    if not value:
        return None
    try:
        data = json.loads(value)
    except Exception:
        return None
    
    # 封装为指定结构
    result = {
        "code": None,
        "crawlDate": None,
        "marketValue": None,
        "maxPrice": None,
        "minPrice": None,
        "nowPrice": None,
        "pbRate": None,
        "rangeRiseAndFall": None,
        "shortName": None,
        "todayStartPrice": None,
        "ttm": None,
        "turnoverRate": None,
        "yesterdayEndPrice": None
    }
    
    # 赋值映射
    result["code"] = data.get("symbol")
    result["crawlDate"] = data.get("fetch_time")
    result["marketValue"] = data.get("market_capital")
    result["maxPrice"] = data.get("high") if "high" in data else data.get("high52w")
    result["minPrice"] = data.get("low") if "low" in data else data.get("low52w")
    result["nowPrice"] = data.get("current")
    result["pbRate"] = data.get("pb")
    result["rangeRiseAndFall"] = data.get("percent")
    result["shortName"] = data.get("name")
    result["todayStartPrice"] = data.get("open")
    result["ttm"] = data.get("pe_ttm")
    result["turnoverRate"] = data.get("turnover_rate")
    result["yesterdayEndPrice"] = data.get("last_close") if "last_close" in data else data.get("pre_close")
    
    # 兼容部分字段缺失
    if result["maxPrice"] is None and "high" in data:
        result["maxPrice"] = data["high"]
    if result["minPrice"] is None and "low" in data:
        result["minPrice"] = data["low"]
    return result


def fetch_and_store_hk_stock_data_optimized(page_size=90, max_workers=15, use_proxy=False):
    """
    优化版本的港股批量采集函数，支持更灵活的配置
    
    :param page_size: 每页采集数量
    :param max_workers: 线程池最大工作线程数（建议10-20之间）
    :param use_proxy: 是否使用代理（默认False）
    """
    print(f"开始批量采集港股数据...")
    print(f"配置: 每页 {page_size} 条记录，最大线程数 {max_workers}")
    print(f"代理模式: {'启用' if use_proxy else '禁用'}")
    print(f"预计采集: 港股所有股票数据")
    print("-" * 50)
    
    try:
        result = fetch_and_store_hk_stock_data(page_size, max_workers, use_proxy)
        if not result.empty:
            print(f"港股采集完成！共获取 {len(result)} 只股票的数据")
            print(f"数据已保存到Redis键: {REDIS_KEY}")
        else:
            print("港股采集完成，但未获取到数据")
    except Exception as e:
        print(f"港股采集过程中发生错误: {e}")
        return None
    
    return result


if __name__ == '__main__':
    # 可以根据需要调整参数
    # fetch_and_store_hk_stock_data_optimized(page_size=100, max_workers=15, use_proxy=True)
    fetch_and_store_hk_stock_data_optimized(use_proxy=False)  # 默认不使用代理 