#!/usr/bin/env python
# -*- coding: utf-8 -*-

import requests
import json
from datetime import datetime, timedelta
import logging

logger = logging.getLogger(__name__)

class HSGTMonitor:
    """沪深港通资金流向监控类"""
    
    def __init__(self):
        # 初始化请求头
        self.headers = {
            "Accept": "application/json, text/plain, */*",
            "Accept-Encoding": "gzip, deflate, br, zstd",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
            "Connection": "keep-alive",
            "Content-Type": "application/json;charset=UTF-8",
            "Host": "dataq.10jqka.com.cn",
            "Origin": "https://data.10jqka.com.cn",
            "Platform": "web",
            "Referer": "https://data.10jqka.com.cn/",
            "Source-Id": "b2cweb-hsgtconnect",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36",
            "sec-ch-ua": "\"Chromium\";v=\"134\", \"Not:A-Brand\";v=\"24\", \"Google Chrome\";v=\"134\"",
            "sec-ch-ua-mobile": "?0",
            "sec-ch-ua-platform": "\"Windows\""
        }
        
        # 请求URL
        self.url = "https://dataq.10jqka.com.cn/fetch-data-server/fetch/v1/interval_data"
        
        # 北向资金代码映射 (从香港流入A股的资金)
        self.northbound_codes = {
            "total": "48:883957",     # 陆股通整体(北向资金总和)
            "sh": "16:1A0001",        # 沪股通
            "sz": "32:399001"         # 深股通
        }
        
        # 南向资金代码映射 (从内地流入港股的资金)
        self.southbound_codes = {
            "total": "48:883957",    # 港股通整体(南向资金总和)
            "hk_sh": "16:1A0001",    # 沪市港股通
            "hk_sz": "32:399001"     # 深市港股通
        }
        
        # 北向和南向资金的API参数
        self.index_ids = {
            "northbound": "hsgt_main_money",  # 北向资金
            "southbound": "ggt_net_buy"       # 南向资金
        }

    def get_trading_hours_timestamp(self, date=None):
        """获取指定日期的交易时段时间戳
        
        交易时段: 9:30-11:30, 13:00-15:00
        为了完整捕捉开盘前的资金流向，开始时间设为9:20
        """
        if date is None:
            date = datetime.now().date()
        elif isinstance(date, str):
            date = datetime.strptime(date, "%Y-%m-%d").date()
            
        # 创建当天交易开始和结束时间
        morning_start = datetime.combine(date, datetime.strptime("09:20", "%H:%M").time())
        trading_end = datetime.combine(date, datetime.strptime("15:00", "%H:%M").time())
        
        # 检查当前是否已经过了交易结束时间
        current_time = datetime.now()
        if current_time > trading_end:
            end_time = trading_end
        else:
            end_time = current_time
            
        # 转换为时间戳（秒）
        start_timestamp = int(morning_start.timestamp())
        end_timestamp = int(end_time.timestamp())
        
        return start_timestamp, end_timestamp
    
    def fetch_northbound_data(self, start_timestamp=None, end_timestamp=None):
        """获取北向资金流向数据 (从香港流入A股的资金)"""
        return self._fetch_hsgt_data(
            "northbound", 
            self.northbound_codes, 
            self.index_ids["northbound"], 
            start_timestamp, 
            end_timestamp
        )
    
    def fetch_southbound_data(self, start_timestamp=None, end_timestamp=None):
        """获取南向资金流向数据 (从内地流入港股的资金)"""
        return self._fetch_hsgt_data(
            "southbound", 
            self.southbound_codes, 
            self.index_ids["southbound"], 
            start_timestamp, 
            end_timestamp
        )
    
    def _fetch_hsgt_data(self, flow_type, code_map, index_id, start_timestamp=None, end_timestamp=None):
        """通用的沪深港通资金流向数据获取方法
        
        参数:
            flow_type: 资金流向类型，"northbound" 或 "southbound"
            code_map: 对应的代码映射
            index_id: API的index_id参数
            start_timestamp: 开始时间戳
            end_timestamp: 结束时间戳
        """
        # 如果没有提供时间戳，使用当天的交易时段
        if start_timestamp is None or end_timestamp is None:
            start_timestamp, end_timestamp = self.get_trading_hours_timestamp()
        
        # 获取代码列表
        if flow_type == "northbound":
            codes = [
                code_map["total"],
                code_map["sh"],
                code_map["sz"]
            ]
        else:  # southbound
            codes = [
                code_map["total"],
                code_map["hk_sh"],
                code_map["hk_sz"]
            ]
            
        # 构建请求体
        payload = {
            "indexes": [{
                "codes": codes,
                "index_info": [{
                    "index_id": index_id
                }]
            }],
            "time_range": {
                "time_type": "TREND",
                "start": str(start_timestamp),
                "end": str(end_timestamp)
            }
        }
        
        # logger.info(f"请求{flow_type}资金数据: start={start_timestamp}, end={end_timestamp}, index_id={index_id}")
        
        try:
            # 发送请求
            response = requests.post(self.url, headers=self.headers, json=payload)
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            if result["status_code"] == 0 and "data" in result:
                # 提取时间和各通道的数据
                time_range = [int(ts) for ts in result["data"]["time_range"]]
                time_formatted = [datetime.fromtimestamp(ts).strftime('%H:%M') for ts in time_range]
                
                # 提取各通道的数据
                data_dict = {}
                
                # 检查是否有数据 
                if "data" in result["data"] and result["data"]["data"]:
                    for item in result["data"]["data"]:
                        code = item["code"]
                        # 增加安全检查，确保values列表存在且不为空
                        if "values" in item and len(item["values"]) > 0 and "values" in item["values"][0]:
                            # 北向资金
                            if flow_type == "northbound":
                                if code == code_map["total"]:
                                    data_dict["total"] = item["values"][0]["values"]
                                elif code == code_map["sh"]:
                                    data_dict["sh"] = item["values"][0]["values"]
                                elif code == code_map["sz"]:
                                    data_dict["sz"] = item["values"][0]["values"]
                            # 南向资金
                            else:
                                if code == code_map["total"]:
                                    data_dict["total"] = item["values"][0]["values"]
                                elif code == code_map["hk_sh"]:
                                    data_dict["hk_sh"] = item["values"][0]["values"]
                                elif code == code_map["hk_sz"]:
                                    data_dict["hk_sz"] = item["values"][0]["values"]
                
                # 检查是否获取到了任何数据
                if not data_dict:
                    logger.warning(f"{flow_type}资金数据响应中没有有效数据")
                    return {
                        "success": False,
                        "message": f"{flow_type}资金数据响应中没有有效数据"
                    }
                
                # 计算当前值（增加更多的安全检查）
                current_values = {}
                
                # 安全地获取当前值
                for key in data_dict:
                    if data_dict[key] and len(data_dict[key]) > 0:  # 确保数据列表不为空
                        value = data_dict[key][-1]
                        # 确保值不是None
                        if value is not None:
                            current_values[key] = value / 100000000
                        else:
                            current_values[key] = 0
                            logger.warning(f"{flow_type}资金{key}最后一个值为None")
                    else:
                        current_values[key] = 0
                        logger.warning(f"{flow_type}资金{key}数据为空")
                
                # 转换为亿元单位，处理None值
                for key in data_dict:
                    if data_dict[key]:  # 确保数据列表不为空
                        # 安全转换，处理None值
                        data_dict[key] = [
                            (value / 100000000 if value is not None else 0) 
                            for value in data_dict[key]
                        ]
                
                return {
                    "success": True,
                    "timestamps": time_range,
                    "times": time_formatted,
                    "data": data_dict,
                    "current": current_values,
                    "flow_type": flow_type,
                    "update_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                }
            else:
                error_msg = result.get("status_msg", "Unknown error")
                logger.error(f"请求{flow_type}资金数据失败: {error_msg}")
                return {
                    "success": False,
                    "message": f"返回错误: {error_msg}"
                }
                
        except requests.exceptions.RequestException as e:
            logger.error(f"请求{flow_type}资金数据异常: {str(e)}")
            return {
                "success": False,
                "message": f"请求异常: {str(e)}"
            }
        except Exception as e:
            logger.error(f"处理{flow_type}资金数据异常: {str(e)}")
            return {
                "success": False,
                "message": f"处理异常: {str(e)}"
            }

# 测试代码
if __name__ == "__main__":
    monitor = HSGTMonitor()
    
    # 获取北向资金数据
    try:
        north_data = monitor.fetch_northbound_data()
        print("北向资金数据:")
        print(json.dumps(north_data, ensure_ascii=False, indent=2))
    except Exception as e:
        print(f"测试北向资金数据获取失败: {str(e)}")
    
    # 获取南向资金数据
    try:
        south_data = monitor.fetch_southbound_data()
        print("\n南向资金数据:")
        print(json.dumps(south_data, ensure_ascii=False, indent=2))
    except Exception as e:
        print(f"测试南向资金数据获取失败: {str(e)}") 