"""
Schema Linker Agent - 表筛选专家
从大量数据表中识别与用户问题相关的表
"""

import logging
from typing import Dict, List, Tuple, Optional
import json

from camel.agents import ChatAgent
from camel.models import ChatModel
from config.prompts import SCHEMA_LINKER_SYSTEM

logger = logging.getLogger(__name__)


class SchemaLinkerAgent:
    """
    Schema Linker Agent

    职责：
    - 分析用户问题中的实体和意图
    - 从候选表中筛选真正相关的表
    - 提供选择理由
    """

    def __init__(
        self,
        model: ChatModel,
        system_message: Optional[str] = None
    ):
        """
        初始化Agent

        Args:
            model: CAMEL AI模型实例
            system_message: 系统提示词（默认使用SCHEMA_LINKER_SYSTEM）
        """
        self.system_message = system_message or SCHEMA_LINKER_SYSTEM
        self.agent = ChatAgent(
            system_message=self.system_message,
            model=model,
        )
        logger.info("[OK] SchemaLinkerAgent初始化完成")

    def select_tables(
        self,
        question: str,
        candidate_tables: List[str],
        table_metadata: Optional[Dict[str, str]] = None,
        max_tables: int = 5
    ) -> Tuple[List[str], str]:
        """
        选择相关表

        Args:
            question: 用户问题
            candidate_tables: 候选表列表（粗筛结果）
            table_metadata: 表元数据 {表名: 注释}
            max_tables: 最多返回表数量

        Returns:
            (相关表列表, 推理理由)
        """
        # 构造候选表信息字符串
        if table_metadata:
            table_list = "\n".join([
                f"- {tbl}: {table_metadata.get(tbl, '无描述')}"
                for tbl in candidate_tables
            ])
        else:
            table_list = "\n".join([f"- {tbl}" for tbl in candidate_tables])

        from config.prompts import SCHEMA_LINKER_USER
        prompt = SCHEMA_LINKER_USER.format(
            question=question,
            table_list=table_list
        )

        try:
            response = self.agent.step(prompt)
            content = response.msg.content.strip()

            # 解析JSON响应
            result = self._parse_json_response(content)

            relevant_tables = result.get("relevant_tables", [])
            reasoning = result.get("reasoning", "")

            # 限制数量
            relevant_tables = relevant_tables[:max_tables]

            logger.info(
                f"SchemaLinker选中 {len(relevant_tables)} 张表: {relevant_tables}"
            )
            return relevant_tables, reasoning

        except json.JSONDecodeError as e:
            logger.error(f"JSON解析失败: {e}, 原始内容: {content[:200]}")
            # 降级：返回前max_tables个候选表
            return candidate_tables[:max_tables], "JSON解析失败，使用粗筛结果"
        except Exception as e:
            logger.error(f"Agent调用失败: {e}")
            return candidate_tables[:max_tables], f"Agent错误: {str(e)}"

    def _parse_json_response(self, content: str) -> Dict:
        """
        解析Agent的JSON响应

        处理可能的markdown代码块包裹
        """
        import json

        # 尝试提取```json```块
        if "```json" in content:
            start = content.find("```json") + 7
            end = content.find("```", start)
            if end != -1:
                content = content[start:end].strip()
        elif "```" in content:
            start = content.find("```") + 3
            end = content.find("```", start)
            if end != -1:
                content = content[start:end].strip()

        return json.loads(content)

    def expand_by_relations(
        self,
        selected_tables: List[str],
        schema_manager,
        depth: int = 1
    ) -> List[str]:
        """
        通过外键关系扩展表（备选方案，也可在Orchestrator中完成）

        Args:
            selected_tables: 已选中的表
            schema_manager: Schema管理器
            depth: 递归深度

        Returns:
            扩展后的表列表
        """
        result = set(selected_tables)

        for tbl_name in selected_tables:
            table = schema_manager.get_table(tbl_name)
            if not table:
                continue

            # 外键指向的表
            for fk in table.foreign_keys:
                if fk.ref_table not in result:
                    result.add(fk.ref_table)

            # 引用当前表的表（反向外键）
            for other in schema_manager.get_tables():
                for fk in other.foreign_keys:
                    if fk.ref_table == tbl_name and other.name not in result:
                        result.add(other.name)

        return list(result)