"""EDA Agent reasoning engine implementation."""
from __future__ import annotations

import json
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from llm.providers import BaseLLMProvider
from eda_agent.services.normbase_client import NormbaseClient
from eda_agent.state import (
    Belief,
    BeliefSource,
    EDAState,
    NormBinding,
    NormType,
    ValueWeight,
)


class EDAEngine:
    """EDA Agent 推理引擎"""
    
    def __init__(self, normbase_client: NormbaseClient, llm_provider: BaseLLMProvider):
        """
        初始化推理引擎
        
        Args:
            normbase_client: Normbase 客户端
            llm_provider: LLM 提供商
        """
        self.normbase = normbase_client
        self.llm = llm_provider
        self.perceive_system_prompt = self._load_prompt(
            "perceive_system.md",
            "You are a structured extractor for EDA propositions.",
        )
        self.bind_system_prompt = self._load_prompt("bind_system.md", "")
        self.rank_system_prompt = self._load_prompt("rank_system.md", "")
        self.last_bind_debug: Dict[str, Any] = {}
    
    async def initialize(
        self,
        agent_id: str,
        agent_role: str,
        value_config: List[ValueWeight],
        active_fields: Optional[List[str]] = None,
        trust_score: float = 0.8,
    ) -> EDAState:
        """
        初始化 Agent 状态
        
        Args:
            agent_id: Agent ID
            agent_role: Agent 角色
            value_config: 价值权重配置
        
        Returns:
            初始化的 EDA 状态
        """
        state = EDAState(
            agent_id=agent_id,
            agent_role=agent_role,
            value_weights=value_config,
            active_fields=active_fields or [],
            trust_score=trust_score,
        )
        
        # 从 Normbase 获取角色默认规范
        default_norms = await self.normbase.get_default_norms(agent_role)
        for norm in default_norms:
            norm_type = norm.get("norm_type") or norm.get("type") or "obligation"
            binding = NormBinding(
                norm_id=str(norm["id"]),
                norm_type=NormType(norm_type),
                condition_met=False,
                action_required=norm.get("action_required", "")
                or norm.get("action", {}).get("required", ""),
                action_prohibited=norm.get("action_prohibited", [])
                or norm.get("action", {}).get("prohibited", []),
            )
            state.norm_bindings.append(binding)
        
        return state
    
    async def perceive(
        self,
        state: EDAState,
        input_data: Dict[str, Any],
    ) -> EDAState:
        """
        感知阶段：处理输入，更新信念
        
        Args:
            state: 当前状态
            input_data: 输入数据
        
        Returns:
            更新后的状态
        """
        # 1. 提取关键命题（使用 LLM）
        propositions = await self._extract_propositions(input_data)
        
        # 2. 更新信念库
        for prop in propositions:
            belief = Belief(
                id=prop["id"],
                proposition=prop["content"],
                confidence=prop["confidence"],
                source=BeliefSource(prop["source"]),
                norm_refs=prop.get("norm_refs", [])
            )
            state.beliefs[belief.id] = belief
        
        state.updated_at = datetime.now()
        return state
    
    async def bind_norms(
        self,
        state: EDAState,
        context: Dict[str, Any],
    ) -> EDAState:
        """
        D 组件：从 Normbase 检索并绑定适用规范
        
        Args:
            state: 当前状态
            context: 上下文信息
        
        Returns:
            更新后的状态（含新的规范绑定）
        """
        field_ids: Optional[List[str]] = None
        if state.active_fields:
            # Normbase field_ids expects UUID strings; active_fields may contain business tags.
            parsed: List[str] = []
            for field in state.active_fields:
                try:
                    parsed.append(str(uuid.UUID(str(field))))
                except (ValueError, TypeError):
                    continue
            if parsed:
                field_ids = parsed

        # 如果有可用 field_ids，则限制检索范围；否则做角色+上下文检索
        if field_ids:
            norms = await self.normbase.retrieve_norms(
                context=context,
                agent_role=state.agent_role,
                field_ids=field_ids,
            )
        else:
            # 如果没有激活的场，则按原方式检索所有规范
            norms = await self.normbase.retrieve_norms(
                context=context,
                agent_role=state.agent_role,
                top_k=20,
            )

        prompt_norms = [self._build_bind_prompt_norm(norm) for norm in norms]
        llm_bind_hints = await self._llm_bind_hints(
            prompt_norms=prompt_norms,
            context=context,
            agent_role=state.agent_role,
        )
        self.last_bind_debug = {
            "total_norms": len(prompt_norms),
            "formatted_prompt_used": sum(
                1 for n in prompt_norms if n.get("prompt_source") == "formatted_prompt"
            ),
            "fallback_prompt_used": sum(
                1 for n in prompt_norms if n.get("prompt_source") == "fallback_template"
            ),
            "llm_hint_count": len(llm_bind_hints),
        }

        for norm in norms:
            # 检查是否已绑定
            norm_id = str(norm["id"])
            existing = [b for b in state.norm_bindings if b.norm_id == norm_id]
            if existing:
                continue

            # 结构化字段仍是最终判定依据；LLM 仅做相关性提示
            rule_condition_met = self._check_condition(norm.get("condition", {}), context)
            llm_hint = llm_bind_hints.get(norm_id)
            if llm_hint is False and not rule_condition_met:
                continue

            # 创建新的规范绑定
            norm_type = norm.get("norm_type") or norm.get("type") or "obligation"
            binding = NormBinding(
                norm_id=norm_id,
                norm_type=NormType(norm_type),
                condition_met=rule_condition_met,
                action_required=norm.get("action", {}).get("required", ""),
                action_prohibited=norm.get("action", {}).get("prohibited", []),
                sanction=norm.get("consequence", {}).get("on_violation"),
            )
            state.norm_bindings.append(binding)
        
        state.updated_at = datetime.now()
        return state
    
    async def rank_values(
        self,
        state: EDAState,
        candidate_actions: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """
        A 组件：按价值权重排序候选行动
        
        Args:
            state: 当前状态
            candidate_actions: 候选行动列表
        
        Returns:
            排序后的行动列表（含得分）
        """
        # 构建权重映射
        weight_map = {vw.value_name: vw.weight for vw in state.value_weights}
        
        # 计算每个行动的价值得分
        scored_actions = []
        for action in candidate_actions:
            score = 0.0
            for value_name, satisfaction in action.get("value_satisfactions", {}).items():
                weight = weight_map.get(value_name, 0.5)
                score += weight * satisfaction
            
            scored_actions.append({
                "action": action,
                "score": score,
                "explanation": {
                    "weights": weight_map,
                    "value_satisfactions": action.get("value_satisfactions", {}),
                },
            })
        
        # 按得分排序
        scored_actions.sort(key=lambda x: x["score"], reverse=True)
        return scored_actions
    
    async def decide(
        self,
        state: EDAState,
        candidate_actions: List[Dict[str, Any]],
    ) -> Optional[Dict[str, Any]]:
        """
        决策阶段：选择最终行动
        
        Args:
            state: 当前状态
            candidate_actions: 候选行动
        
        Returns:
            选中的行动
        """
        ranked = await self.rank_values(state, candidate_actions)
        if ranked:
            state.current_intention = json.dumps(ranked[0]["action"], ensure_ascii=False)
            return ranked[0]["action"]
        return None
    
    async def check_compliance(
        self,
        state: EDAState,
        action: Dict[str, Any],
    ) -> Dict[str, Any]:
        """
        合规校验：检查行动是否违反绑定规范
        
        Args:
            state: 当前状态
            action: 待执行行动
        
        Returns:
            校验结果
        """
        violations = []
        
        action_str = json.dumps(action, ensure_ascii=False)
        for binding in state.norm_bindings:
            # 检查禁止性规范
            if binding.norm_type == NormType.PROHIBITION:
                if binding.action_prohibited:
                    for prohibited in binding.action_prohibited:
                        if prohibited and prohibited in action_str:
                            violations.append({
                                "norm_id": binding.norm_id,
                                "type": "prohibition",
                                "description": f"Action contains prohibited: {prohibited}",
                                "boundary": "hard",
                            })
            
            # 检查义务性规范
            elif binding.norm_type == NormType.OBLIGATION:
                if binding.action_required and binding.condition_met:
                    if binding.action_required and binding.action_required not in action_str:
                        violations.append({
                            "norm_id": binding.norm_id,
                            "type": "obligation",
                            "description": f"Missing required action: {binding.action_required}",
                            "boundary": "soft",
                        })

        hard = [v for v in violations if v.get("boundary") == "hard"]
        soft = [v for v in violations if v.get("boundary") == "soft"]
        flexible = [v for v in violations if v.get("boundary") == "flexible"]

        return {
            "approved": len(violations) == 0,
            "violations": violations,
            "hard_boundary_violations": hard,
            "soft_boundary_violations": soft,
            "flexible_boundary_violations": flexible,
        }
    
    async def execute_action(
        self,
        state: EDAState,
        action: Dict[str, Any],
    ) -> Dict[str, Any]:
        """
        执行行动并更新状态
        
        Args:
            state: 当前状态
            action: 待执行行动
        
        Returns:
            执行结果
        """
        # 1. 行动前校验
        compliance = await self.check_compliance(state, action)
        if compliance["hard_boundary_violations"]:
            return {
                "success": False,
                "reason": "hard_boundary_violation",
                "violations": compliance["violations"]
            }
        if compliance["soft_boundary_violations"]:
            return {
                "success": False,
                "reason": "soft_boundary_violation",
                "violations": compliance["violations"],
                "suggestion": "Provide alternative compliant action",
            }
        
        # 2. 执行行动（由调用方执行）
        # ...
        
        # 3. 更新状态
        state.current_intention = None
        state.updated_at = datetime.now()
        
        return {"success": True, "compliance": compliance}
    
    # ========== Private Methods ==========
    
    async def _extract_propositions(self, input_data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """从输入中提取命题（LLM 辅助）"""
        if self.llm:
            prompt = (
                "Extract key propositions from input as JSON list with keys: "
                "id, content, confidence, source, norm_refs. "
                f"Input: {json.dumps(input_data, ensure_ascii=False)}"
            )
            try:
                response = await self.llm.chat(
                    [
                        {
                            "role": "system",
                            "content": self.perceive_system_prompt,
                        },
                        {"role": "user", "content": prompt},
                    ]
                )
                parsed = json.loads(response)
                if isinstance(parsed, list) and parsed:
                    normalized = []
                    for item in parsed:
                        normalized.append(
                            {
                                "id": item.get("id", str(uuid.uuid4())),
                                "content": item.get("content", ""),
                                "confidence": float(item.get("confidence", 0.7)),
                                "source": item.get("source", BeliefSource.PERCEPTION.value),
                                "norm_refs": item.get("norm_refs", []),
                            }
                        )
                    return normalized
            except Exception:
                pass

        propositions = []
        for key, value in input_data.items():
            propositions.append(
                {
                    "id": str(uuid.uuid4()),
                    "content": f"{key}={value}",
                    "confidence": 0.7,
                    "source": BeliefSource.PERCEPTION.value,
                    "norm_refs": [],
                }
            )
        return propositions
    
    def _check_condition(self, condition: Dict[str, Any], context: Dict[str, Any]) -> bool:
        """检查条件是否满足"""
        if not condition:
            return True

        trigger = condition.get("trigger")
        if trigger:
            if isinstance(trigger, str) and context.get(trigger) not in [True, "true", 1]:
                # Allow trigger text matching in any context value
                text = json.dumps(context, ensure_ascii=False)
                if trigger not in text:
                    return False

        preconditions = condition.get("preconditions", [])
        for expr in preconditions:
            if not isinstance(expr, str):
                continue
            if "==" in expr:
                key, expected = [s.strip() for s in expr.split("==", 1)]
                if str(context.get(key)) != expected.strip("\"'"):
                    return False
            elif "!=" in expr:
                key, expected = [s.strip() for s in expr.split("!=", 1)]
                if str(context.get(key)) == expected.strip("\"'"):
                    return False
            else:
                key = expr.strip()
                if context.get(key) in [None, False, "", 0]:
                    return False

        return True

    def _load_prompt(self, filename: str, fallback: str) -> str:
        """Load prompt from eda_agent/prompts with fallback string."""
        prompt_path = Path(__file__).resolve().parent / "prompts" / filename
        try:
            return prompt_path.read_text(encoding="utf-8").strip() or fallback
        except Exception:
            return fallback

    def _build_bind_prompt_norm(self, norm: Dict[str, Any]) -> Dict[str, Any]:
        """Build bind-stage prompt material with formatted_prompt-first strategy."""
        formatted = norm.get("formatted_prompt")
        if isinstance(formatted, str) and formatted.strip():
            return {
                "norm_id": str(norm["id"]),
                "prompt_source": "formatted_prompt",
                "prompt_text": formatted.strip(),
            }
        return {
            "norm_id": str(norm["id"]),
            "prompt_source": "fallback_template",
            "prompt_text": (
                f"NORM[{norm.get('id')}]\n"
                f"TYPE: {norm.get('norm_type') or norm.get('type') or 'obligation'}\n"
                f"CONDITION: {json.dumps(norm.get('condition', {}), ensure_ascii=False)}\n"
                f"ACTION: {json.dumps(norm.get('action', {}), ensure_ascii=False)}\n"
                f"CONSEQUENCE: {json.dumps(norm.get('consequence', {}), ensure_ascii=False)}"
            ),
        }

    async def _llm_bind_hints(
        self,
        prompt_norms: List[Dict[str, Any]],
        context: Dict[str, Any],
        agent_role: str,
    ) -> Dict[str, bool]:
        """Use bind-stage prompt with formatted_prompt inputs; return lightweight hints."""
        if not self.llm or not prompt_norms:
            return {}
        try:
            payload = {
                "agent_role": agent_role,
                "context": context,
                "norms": prompt_norms,
            }
            response = await self.llm.chat(
                [
                    {"role": "system", "content": self.bind_system_prompt},
                    {
                        "role": "user",
                        "content": (
                            "Assess norm relevance and condition status for this context. "
                            "Return JSON array with fields: norm_id, condition_met, reason.\n"
                            f"{json.dumps(payload, ensure_ascii=False)}"
                        ),
                    },
                ]
            )
            parsed = json.loads(response)
            if not isinstance(parsed, list):
                return {}
            hints: Dict[str, bool] = {}
            for item in parsed:
                norm_id = str(item.get("norm_id", "")).strip()
                if not norm_id:
                    continue
                hints[norm_id] = bool(item.get("condition_met", False))
            return hints
        except Exception:
            return {}
