"""
===========================================
路由标签识别模块 - RouteLabel
===========================================
功能：使用LLM识别用户查询中涉及的节点类型和实体
     作为后续Neo4j查询的入口节点
===========================================
"""

import json
import logging
import os
import time
import asyncio
from typing import List, Dict, Any, Optional, Tuple
from pydantic import BaseModel, Field
from graph_search.init_prompts import get_extract_entities_prompt, get_rewrite_query_prompt

logger = logging.getLogger(__name__)
#logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)   # 不依赖继承链


class RouteItem(BaseModel):
    """路由项：节点类型和实体"""
    label: str = Field(..., description="节点类型，比如'Drug'、'Disease'")
    entity: str = Field(..., description="实体文本")


class RouteOutput(BaseModel):
    """路由输出：包含多个路由项"""
    outputs: List[RouteItem]


class RouteLabel:
    """
    路由标签识别类
    
    功能：使用LLM识别用户查询中涉及的节点类型和实体
    """
    
    def __init__(self, llm_api, optional_labels: list[str], schema: str = None):
        """
        初始化路由识别器
        
        Args:
            llm_api: LLM API实例（OpenaiAPI）
            optional_labels: 可选节点标签列表
            schema: Neo4j图谱schema信息（可选，用于帮助模型识别节点类型）
        """
        self.llm_api = llm_api
        self.optional_labels = optional_labels
        self.schema = schema
    
    async def _extract_entities(self, query: str) -> List[Dict[str, str]]:
        """
        任务1：抽取主语实体
        
        Args:
            query: 用户查询
            
        Returns:
            List[Dict]: 路由结果列表，格式如 [{"label": "Drug", "entity": "xxx"}]
        """
        # 使用 init_prompts 中的 prompt 生成函数
        prompt = get_extract_entities_prompt(
            query=query,
            optional_labels=self.optional_labels,
            schema=self.schema
        )
        
        try:
            llm_start = time.time()
            model = os.getenv("OPENAI_MODEL")
            response = await self.llm_api.open_api_chat_async_json(prompt, model, temperature=0)
            llm_elapsed = (time.time() - llm_start) * 1000
            
            # 解析JSON响应
            parse_start = time.time()
            content = response.strip()
            
            # 如果包含markdown代码块，提取JSON
            if "```json" in content:
                json_start = content.find("```json") + 7
                json_end = content.find("```", json_start)
                if json_end > json_start:
                    content = content[json_start:json_end].strip()
            elif "```" in content:
                json_start = content.find("```") + 3
                json_end = content.find("```", json_start)
                if json_end > json_start:
                    content = content[json_start:json_end].strip()
            
            # 尝试找到JSON数组或对象的开始和结束位置
            if not content.startswith('[') and not content.startswith('{'):
                for start_char in ['[', '{']:
                    start_idx = content.find(start_char)
                    if start_idx != -1:
                        content = content[start_idx:]
                        break
            
            # 解析JSON
            try:
                outputs_data = json.loads(content)
            except json.JSONDecodeError as e:
                logger.warning(f"任务1 JSON解析失败，尝试修复: {e}")
                if content.rfind(']') > content.rfind('}'):
                    end_idx = content.rfind(']') + 1
                    content = content[:end_idx]
                elif content.rfind('}') > -1:
                    end_idx = content.rfind('}') + 1
                    content = content[:end_idx]
                try:
                    outputs_data = json.loads(content)
                except json.JSONDecodeError:
                    logger.error(f"任务1 JSON解析最终失败，原始内容: {content[:200]}...")
                    return []
            
            # 提取entities
            if isinstance(outputs_data, dict) and "entities" in outputs_data:
                route_results = [
                    {"label": item.get("label", ""), "entity": item.get("entity", "")}
                    for item in outputs_data["entities"]
                ]
            elif isinstance(outputs_data, list):
                # 兼容旧格式：直接是列表
                route_results = [
                    {"label": item.get("label", ""), "entity": item.get("entity", "")}
                    for item in outputs_data
               ]
            elif isinstance(outputs_data, dict) and "outputs" in outputs_data:
                # 兼容旧格式：包含 outputs 字段
                route_results = [
                    {"label": item.get("label", ""), "entity": item.get("entity", "")}
                    for item in outputs_data["outputs"]
                ]
            else:
                route_results = []
            
            return route_results
            
        except Exception as e:
            logger.error(f"任务1（实体抽取）失败: {e}")
            return []
    
    async def _rewrite_query(self, query: str) -> Dict[str, Any]:
        """
        任务2：重写用户问题并进行分类
        
        Args:
            query: 用户查询
            
        Returns:
            Dict: 包含重写后的查询和分类信息
                - rewritten_query: 重写后的查询字符串
                - query_classification: 分类信息字典
                    - query_type: "aggregate"（统计类）或 "detail"（明细类）
                    - has_time: 是否包含时间信息（bool）
                    - time_type: 时间统计类型（"single"/"dual"/"multi"/None）
        """
        # 使用 init_prompts 中的 prompt 生成函数
        prompt = get_rewrite_query_prompt(
            query=query,
            schema=self.schema
        )

        
        try:
            llm_start = time.time()
            model = os.getenv("OPENAI_MODEL")
            response = await self.llm_api.open_api_chat_async_json(prompt, model, temperature=0)
            
            # 解析JSON响应
            content = response.strip()
            
            # 如果包含markdown代码块，提取JSON
            if "```json" in content:
                json_start = content.find("```json") + 7
                json_end = content.find("```", json_start)
                if json_end > json_start:
                    content = content[json_start:json_end].strip()
            elif "```" in content:
                json_start = content.find("```") + 3
                json_end = content.find("```", json_start)
                if json_end > json_start:
                    content = content[json_start:json_end].strip()
            
            # 尝试找到JSON数组或对象的开始和结束位置
            if not content.startswith('[') and not content.startswith('{'):
                for start_char in ['[', '{']:
                    start_idx = content.find(start_char)
                    if start_idx != -1:
                        content = content[start_idx:]
                        break
            
            # 解析JSON
            try:
                outputs_data = json.loads(content)
            except json.JSONDecodeError as e:
                logger.warning(f"任务2 JSON解析失败，尝试修复: {e}")
                if content.rfind(']') > content.rfind('}'):
                    end_idx = content.rfind(']') + 1
                    content = content[:end_idx]
                elif content.rfind('}') > -1:
                    end_idx = content.rfind('}') + 1
                    content = content[:end_idx]
                try:
                    outputs_data = json.loads(content)
                except json.JSONDecodeError:
                    logger.error(f"任务2 JSON解析最终失败，原始内容: {content[:200]}...")
                    return query
            
            # 提取rewritten_query和分类信息
            if isinstance(outputs_data, dict) and "rewritten_query" in outputs_data:
                rewritten_query = outputs_data.get("rewritten_query", query)
                # 提取分类信息
                query_classification = outputs_data.get("query_classification", {})
                if not isinstance(query_classification, dict):
                    query_classification = {}
                
                # 确保分类信息格式正确
                result = {
                    "rewritten_query": rewritten_query,
                    "query_classification": {
                        "query_type": query_classification.get("query_type", "detail"),
                        "has_time": query_classification.get("has_time", False),
                        "time_type": query_classification.get("time_type")
                    }
                }
            else:
                # 如果解析失败，返回默认值
                result = {
                    "rewritten_query": query,
                    "query_classification": {
                        "query_type": "detail",
                        "has_time": False,
                        "time_type": None
                    }
                }
            
            return result
            
        except Exception as e:
            logger.error(f"任务2（查询重写）失败: {e}")
            # 返回默认值
            return {
                "rewritten_query": query,
                "query_classification": {
                    "query_type": "detail",
                    "has_time": False,
                    "time_type": None
                }
            }
    
    async def route(self, query: str) -> Tuple[List[Dict[str, str]], str, Dict[str, Any]]:
        """
        路由标签识别：识别标签，抽取实体
        
        功能：使用LLM识别用户查询中涉及的节点类型和实体。
        将两个任务（实体抽取和查询重写）拆分为独立的异步调用，并行执行以提高效率。
        
        Args:
            query: 用户查询
            
        Returns:
            tuple: (路由结果列表, 重写后的查询, 分类信息)
                路由结果列表格式如 [{"label": "Drug", "entity": "xxx"}]
                重写后的查询字符串，如果未获取到则返回原始query
                分类信息字典，包含 query_type, has_time, time_type
        """
        try:
            # ========== 并行执行两个任务 ==========
            total_start = time.time()
            logger.info(f"[时间统计] RouteLabel: 开始并行执行两个任务（实体抽取 + 查询重写+分类）...")
            
            # 使用 asyncio.gather 并行执行两个任务
            route_results, rewrite_result = await asyncio.gather(
                self._extract_entities(query),
                self._rewrite_query(query)
            )
            logger.info(route_results) 
            logger.info(111111111111111111111111111111111111111)           
            # 从重写结果中提取信息
            if isinstance(rewrite_result, dict):
                rewritten_query = rewrite_result.get("rewritten_query", query)
                query_classification = rewrite_result.get("query_classification", {
                    "query_type": "detail",
                    "has_time": False,
                    "time_type": None
                })
            else:
                # 兼容旧格式（如果返回的是字符串）
                rewritten_query = rewrite_result if isinstance(rewrite_result, str) else query
                query_classification = {
                    "query_type": "detail",
                    "has_time": False,
                    "time_type": None
                }
            
            total_elapsed = (time.time() - total_start) * 1000
            logger.info(f"[时间统计] RouteLabel: 两个任务并行执行完成，总耗时: {total_elapsed:.2f}ms")
            logger.info(f"最终结果 - 入口节点标签与实体: {route_results}")
            logger.info(f"最终结果 - 原始查询: '{query}' -> 重写后查询: '{rewritten_query}'")
            logger.info(f"最终结果 - 分类信息: {query_classification}")
            
            return route_results, rewritten_query, query_classification
            
        except Exception as e:
            logger.error(f"路由标签识别失败: {e}")
            return [], query, {
                "query_type": "detail",
                "has_time": False,
                "time_type": None
            }

