import json
import pandas as pd
from typing import List, Dict, Any, Optional, Union
from chuangliangTool.db_base import Chuangliang_ad_magic_task_skit, Chuangliang_ad_magic_cut
import requests
from functools import lru_cache


def get_task_param(task_id: int) -> Union[Dict[str, Any], str]:
    """获取任务参数"""
    try:
        # 使用字符串格式化而不是参数化查询
        task_param_sql = f"""
            SELECT status, task_param 
            FROM chuangliang_ad_task_skit.task_log_client_mix_shear 
            WHERE task_id = {task_id}
        """
        task_param_result = Chuangliang_ad_magic_task_skit.query_database_select_to_dataframe(task_param_sql)

        # 检查结果是否为空
        if task_param_result and len(task_param_result) > 0:
            row = task_param_result[0]  # 列表中的第一个元素
            status, task_param = row['status'], row['task_param']

            if status == 'success':
                full_url = f"https://tos.mobgi.com/{task_param}"
                response = requests.get(full_url, timeout=10)
                response.raise_for_status()
                return response.json()
            return task_param
        return "No task parameter found for the given task_id"
    except Exception as e:
        return f"Error fetching task param: {str(e)}"


class RuleTaskQuery:
    def __init__(self, rule_name: str):
        self.rule_name = rule_name
        self._rule_id = None
        self._rule_info = None
        self._number_man_rule = None

    @lru_cache(maxsize=None)
    def get_rule_id(self) -> Optional[int]:
        """获取规则ID，使用缓存避免重复查询"""
        # 使用字符串格式化而不是参数化查询
        rule_id_sql = f"""
            SELECT id 
            FROM chuangliang_ad_magic_cut.magic_fission_rule 
            WHERE name = '{self.rule_name}'
        """
        result = Chuangliang_ad_magic_cut.query_database_select_to_dataframe(rule_id_sql)

        # 检查结果是否为空
        if result and len(result) > 0:
            return result[0]['id']  # 列表中的第一个元素
        print(f"未找到规则 '{self.rule_name}' 对应的ID")
        return None

    def get_task_list(self) -> List[int]:
        """获取任务ID列表"""
        rule_id = self.get_rule_id()
        if not rule_id:
            return []

        # 使用字符串格式化而不是参数化查询
        task_list_sql = f"""
            SELECT task_id 
            FROM chuangliang_ad_magic_cut.magic_rule_result 
            WHERE rule_id = {rule_id}
        """
        task_list = Chuangliang_ad_magic_cut.query_database_select_to_dataframe(task_list_sql)

        # 提取任务ID列表
        if task_list and len(task_list) > 0:
            return [task['task_id'] for task in task_list]
        return []

    def _parse_json_field(self, json_str: str) -> Union[Dict[str, Any], str]:
        """安全解析JSON字符串"""
        try:
            return json.loads(json_str)
        except (json.JSONDecodeError, TypeError):
            return json_str

    def _process_task_result(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """处理单个任务结果"""
        if 'result_data' in result:
            result['result_data'] = self._parse_json_field(result['result_data'])
        return result

    def _process_number_man_result(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """处理数字人合成任务结果"""
        if 'task_param' in result:
            result['task_param'] = self._parse_json_field(result['task_param'])

        if 'result_data' in result:
            result_data = self._parse_json_field(result['result_data'])
            if isinstance(result_data, dict) and 'result' in result_data:
                result_data['result'] = self._parse_json_field(result_data['result'])
            result['result_data'] = result_data

        return result

    def get_task_list_result(self) -> List[Dict[str, Any]]:
        """获取任务结果列表"""
        task_ids = self.get_task_list()
        if not task_ids:
            return []

        # 分批处理避免SQL过长
        batch_size = 1000
        results = []
        for i in range(0, len(task_ids), batch_size):
            batch_ids = task_ids[i:i + batch_size]
            # 使用字符串格式化而不是参数化查询
            task_ids_str = ",".join(map(str, batch_ids))

            task_result_sql = f"""
                SELECT task_id, status, result_data 
                FROM chuangliang_ad_task_skit.task_log_client_mix_shear 
                WHERE task_id IN ({task_ids_str})
            """
            batch_results = Chuangliang_ad_magic_task_skit.query_database_select_to_dataframe(task_result_sql)

            if batch_results and len(batch_results) > 0:
                results.extend(batch_results)

        return [self._process_task_result(r) for r in results] if results else []

    def get_number_man_task_list_result(self) -> List[Dict[str, Any]]:
        """获取数字人合成任务结果列表"""
        task_ids = self.get_task_list()
        if not task_ids:
            return []

        # 分批处理避免SQL过长
        batch_size = 1000
        results = []
        for i in range(0, len(task_ids), batch_size):
            batch_ids = task_ids[i:i + batch_size]
            # 使用字符串格式化而不是参数化查询
            task_ids_str = ",".join(map(str, batch_ids))

            number_man_sql = f"""
                SELECT task_id, status, task_param, result_data 
                FROM chuangliang_ad_task_skit.task_log_client_number_man 
                WHERE task_id IN ({task_ids_str})
            """
            batch_results = Chuangliang_ad_magic_task_skit.query_database_select_to_dataframe(number_man_sql)

            if batch_results and len(batch_results) > 0:
                results.extend(batch_results)

        return [self._process_number_man_result(r) for r in results] if results else []

    def get_rule_info(self) -> Union[Dict[str, Any], str]:
        """获取规则批次提交信息"""
        if self._rule_info is not None:
            return self._rule_info

        rule_id = self.get_rule_id()
        if not rule_id:
            return {}

        # 使用字符串格式化而不是参数化查询
        rule_info_sql = f"""
            SELECT value 
            FROM chuangliang_ad_magic_cut.magic_fission_rule_extend 
            WHERE fission_rule_id = {rule_id} AND field = 'preview'
        """
        result = Chuangliang_ad_magic_cut.query_database_select_to_dataframe(rule_info_sql)

        if result and len(result) > 0:
            value = result[0]['value']  # 列表中的第一个元素
            self._rule_info = self._parse_json_field(value)
            return self._rule_info

        print(f"未找到规则ID {rule_id} 的批次提交信息")
        return {}

    def get_number_man_rule(self) -> Union[Dict[str, Any], str]:
        """获取数字人合成批次提交信息"""
        if self._number_man_rule is not None:
            return self._number_man_rule

        rule_id = self.get_rule_id()
        if not rule_id:
            return {}

        # 使用字符串格式化而不是参数化查询
        number_man_sql = f"""
            SELECT value 
            FROM chuangliang_ad_magic_cut.magic_fission_rule_extend 
            WHERE fission_rule_id = {rule_id} AND field = 'content'
        """
        result = Chuangliang_ad_magic_cut.query_database_select_to_dataframe(number_man_sql)

        if result and len(result) > 0:
            value = result[0]['value']  # 列表中的第一个元素
            self._number_man_rule = self._parse_json_field(value)
            return self._number_man_rule

        print(f"未找到数字人合成ID {rule_id} 的批次提交信息")
        return {}

    def get_comprehensive_data(self) -> Dict[str, Any]:
        """获取完整的规则数据"""
        rule_id = self.get_rule_id()
        if not rule_id:
            return {
                "rule_name": self.rule_name,
                "error": "Rule ID not found"
            }

        rule_info = self.get_rule_info()
        use_fallback = not rule_info

        rule_data = self.get_number_man_rule() if use_fallback else rule_info
        task_results = self.get_number_man_task_list_result() if use_fallback else self.get_task_list_result()

        return {
            "rule_name": self.rule_name,
            "rule_id": rule_id,
            "info_task_count": len(rule_data) if isinstance(rule_data, dict) else 0,
            "rule_info": rule_data,
            "task_count": len(task_results),
            "tasks": task_results
        }


def main():
    """主函数示例"""
    rule_name = 'lyy-20251110-191909'
    rule_query = RuleTaskQuery(rule_name)

    # 获取完整数据
    # result = rule_query.get_comprehensive_data()
    man_result = rule_query.get_number_man_task_list_result()
    result = rule_query.get_task_list_result()
    # result = get_task_param(32743070)
    print(len(result))
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()