"""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, Tuple

from llm.providers import BaseLLMProvider
from eda_agent.infra.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", "")
        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"
            modality, operations = self._extract_binding_behavior(
                norm.get("action", {}),
                norm_type,
            )
            binding = NormBinding(
                norm_id=str(norm["id"]),
                norm_type=NormType(norm_type),
                condition_met=False,
                modality=modality,
                operations=operations,
            )
            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: List[str] = []
        for field in state.active_fields:
            try:
                field_ids.append(str(uuid.UUID(str(field))))
            except (ValueError, TypeError):
                continue
        if not field_ids:
            raise ValueError("active_fields must contain at least one valid UUID field ID")

        query_mode, nl_query = self._resolve_bind_query(context)
        role_tag_filters = self._build_retrieve_tag_filters(state.agent_role, context)
        base_payload = await self._retrieve_norms_payload(
            context=context,
            query_mode=query_mode,
            nl_query=nl_query,
            agent_role=state.agent_role,
            tag_filters=role_tag_filters,
            field_ids=field_ids,
        )
        normalized_agent_id = self._normalize_agent_id_or_empty(state.agent_id)
        agent_payload = {"norms": [], "scores": [], "total": 0, "usage_token": "", "meta": {}}
        if normalized_agent_id:
            # Agent-specific channel: avoid forcing role tag to reduce false negatives.
            agent_payload = await self._retrieve_norms_payload(
                context=context,
                query_mode=query_mode,
                nl_query=nl_query,
                agent_role=None,
                agent_ids=[normalized_agent_id],
                tag_filters=self._build_retrieve_tag_filters("", context),
                field_ids=field_ids,
            )
        norms = self._merge_norm_hits(
            base_norms=base_payload.get("norms", []) or [],
            agent_norms=agent_payload.get("norms", []) or [],
        )
        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),
            "canonical_nl_prompt_used": sum(
                1 for n in prompt_norms if n.get("prompt_source") == "canonical_nl"
            ),
            "fallback_prompt_used": sum(
                1 for n in prompt_norms if n.get("prompt_source") == "fallback_minimal"
            ),
            "llm_hint_count": len(llm_bind_hints),
            "query_mode": query_mode,
            "has_nl_query": bool(nl_query),
            "base_retrieve_meta": base_payload.get("meta", {}),
            "agent_retrieve_meta": agent_payload.get("meta", {}),
            "base_total": int(base_payload.get("total", len(base_payload.get("norms", []) or []))),
            "agent_total": int(agent_payload.get("total", len(agent_payload.get("norms", []) or []))),
            "agent_specific_enabled": bool(normalized_agent_id),
            "usage_token_base": str(base_payload.get("usage_token", "")).strip(),
            "usage_token_agent": str(agent_payload.get("usage_token", "")).strip(),
        }

        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"
            modality, operations = self._extract_binding_behavior(
                norm.get("action", {}),
                norm_type,
            )
            consequence_canonical = self._canonical(norm.get("consequence", {}))
            sanction_payload = consequence_canonical.get("on_violation")
            binding = NormBinding(
                norm_id=norm_id,
                norm_type=NormType(norm_type),
                condition_met=rule_condition_met,
                modality=modality,
                operations=operations,
                sanction=sanction_payload,
            )
            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 = []
        permission_explains = []
        
        action_str = json.dumps(action, ensure_ascii=False)
        for binding in state.norm_bindings:
            if not binding.condition_met:
                continue
            modality = str(binding.modality or "").strip().lower()
            operations = self._normalize_string_list(binding.operations)
            if modality == "must_not":
                for prohibited in operations:
                    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 modality == "must":
                for required in operations:
                    if required and required not in action_str:
                        violations.append(
                            {
                                "norm_id": binding.norm_id,
                                "type": "obligation",
                                "description": f"Missing required action: {required}",
                                "boundary": "soft",
                            }
                        )
            elif modality == "may":
                in_scope = True
                if operations:
                    in_scope = any(op and op in action_str for op in operations)
                permission_explains.append(
                    {
                        "norm_id": binding.norm_id,
                        "type": "permission",
                        "allowed_operations": operations,
                        "in_scope": in_scope,
                        "description": (
                            "Permission scope declared by norm."
                            if in_scope
                            else "Action may be outside declared permission scope."
                        ),
                    }
                )
            else:
                # Unknown modality: keep runtime permissive and traceable.
                permission_explains.append(
                    {
                        "norm_id": binding.norm_id,
                        "type": "unknown_modality",
                        "allowed_operations": operations,
                        "in_scope": True,
                        "description": f"Unsupported modality '{modality}', treated as advisory.",
                    }
                )

        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(hard) == 0 and len(soft) == 0,
            "violations": violations,
            "hard_boundary_violations": hard,
            "soft_boundary_violations": soft,
            "flexible_boundary_violations": flexible,
            "permission_explains": permission_explains,
        }
    
    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:
        """检查条件是否满足"""
        canonical = self._canonical(condition)
        if not canonical:
            return True

        trigger_type = str(canonical.get("trigger_type", "")).strip().lower()
        if trigger_type in {"", "event"}:
            return True
        if trigger_type == "state":
            return bool(context)
        if trigger_type == "threshold":
            numeric_signals = [
                value for value in (context or {}).values() if isinstance(value, (int, float))
            ]
            return bool(numeric_signals)
        if trigger_type == "schedule":
            for key in ("schedule", "cron", "time", "deadline", "at"):
                value = context.get(key) if isinstance(context, dict) else None
                if isinstance(value, str) and value.strip():
                    return True
            return False
        return False

    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 _normalize_string_list(self, value: Any) -> List[str]:
        """Normalize mixed string/list payload into a list of non-empty strings."""
        if isinstance(value, list):
            return [str(item).strip() for item in value if str(item).strip()]
        if isinstance(value, str):
            item = value.strip()
            return [item] if item else []
        return []

    def _canonical(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """Return canonical branch for norm content zone fields."""
        if not isinstance(payload, dict):
            return {}
        canonical = payload.get("canonical", {})
        return canonical if isinstance(canonical, dict) else {}

    def _extract_binding_behavior(
        self,
        action_payload: Dict[str, Any],
        norm_type: str,
    ) -> tuple[str, List[str]]:
        canonical = self._canonical(action_payload)
        operations = self._normalize_string_list(canonical.get("operations", []))
        modality = str(canonical.get("modality", "")).strip().lower()
        if modality in {"must", "must_not", "may"}:
            return modality, operations

        # Defensive fallback for legacy/invalid payloads still in storage.
        if str(norm_type).strip().lower() == "prohibition":
            return "must_not", operations
        if str(norm_type).strip().lower() == "obligation":
            return "must", operations
        if str(norm_type).strip().lower() == "permission":
            return "may", operations
        return "must", operations

    def _build_retrieve_tag_filters(
        self, agent_role: str, context: Dict[str, Any]
    ) -> Dict[str, List[str]]:
        """Build lightweight retrieval filters from role + common context tags."""
        filters: Dict[str, List[str]] = {}
        role = (agent_role or "").strip()
        if role:
            filters["role"] = [role]

        for category in ("scenario", "domain", "tool", "skill", "capability"):
            raw = context.get(category)
            values: List[str] = []
            if isinstance(raw, str):
                text = raw.strip()
                if text:
                    values.append(text)
            elif isinstance(raw, list):
                for item in raw:
                    text = str(item).strip()
                    if text and text not in values:
                        values.append(text)
            if values:
                filters[category] = values
        return filters

    def _build_bind_prompt_norm(self, norm: Dict[str, Any]) -> Dict[str, Any]:
        """Build bind-stage prompt material from canonical+nl norm structure."""
        norm_id = str(norm.get("id", "")).strip()
        norm_type = str(norm.get("norm_type") or norm.get("type") or "obligation").strip()
        subject_line = self._zone_prompt_text("subject", norm.get("subject", {}))
        context_line = self._zone_prompt_text("context", norm.get("context", {}))
        condition_line = self._zone_prompt_text("condition", norm.get("condition", {}))
        action_line = self._zone_prompt_text("action", norm.get("action", {}))
        consequence_line = self._zone_prompt_text("consequence", norm.get("consequence", {}))
        prompt_lines = [
            f"NORM[{norm_id}]",
            f"TYPE: {norm_type}",
            f"SUBJECT: {subject_line}",
            f"CONTEXT: {context_line}",
            f"CONDITION: {condition_line}",
            f"ACTION: {action_line}",
            f"CONSEQUENCE: {consequence_line}",
        ]
        prompt_text = "\n".join(prompt_lines).strip()
        if prompt_text:
            return {
                "norm_id": norm_id,
                "prompt_source": "canonical_nl",
                "prompt_text": prompt_text,
            }
        return {
            "norm_id": norm_id,
            "prompt_source": "fallback_minimal",
            "prompt_text": (
                f"NORM[{norm.get('id')}]\n"
                f"TYPE: {norm.get('norm_type') or norm.get('type') or 'obligation'}"
            ),
        }

    async def _retrieve_norms_payload(
        self,
        *,
        context: Dict[str, Any],
        query_mode: str,
        nl_query: str,
        agent_role: Optional[str] = None,
        agent_ids: Optional[List[str]] = None,
        tag_filters: Optional[Dict[str, List[str]]] = None,
        field_ids: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        # Prefer richer retrieval payload (norms/scores/meta/usage_token) when available.
        if hasattr(self.normbase, "retrieve_norms_with_meta"):
            data = await self.normbase.retrieve_norms_with_meta(
                context=context,
                query_mode=query_mode,
                nl_query=nl_query or None,
                agent_role=agent_role,
                agent_ids=agent_ids,
                tag_filters=tag_filters,
                field_ids=field_ids,
            )
            return data or {"norms": [], "scores": [], "total": 0, "usage_token": "", "meta": {}}
        norms = await self.normbase.retrieve_norms(
            context=context,
            query_mode=query_mode,
            nl_query=nl_query or None,
            agent_role=agent_role,
            agent_ids=agent_ids,
            tag_filters=tag_filters,
            field_ids=field_ids,
        )
        return {
            "norms": norms,
            "scores": [1.0 for _ in norms],
            "total": len(norms),
            "usage_token": "",
            "meta": {},
        }

    def _resolve_bind_query(self, context: Dict[str, Any]) -> Tuple[str, str]:
        raw_mode = str((context or {}).get("query_mode", "")).strip().lower()
        mode = raw_mode if raw_mode in {"structured", "semantic", "hybrid"} else ""
        nl_query = self._extract_text_signal(context)
        if mode == "semantic":
            # Keep request valid under semantic mode contract.
            return "semantic", nl_query or "bind norms for current context"
        if mode == "structured":
            return "structured", ""
        if nl_query:
            return "hybrid", nl_query
        return "structured", ""

    def _extract_text_signal(self, context: Dict[str, Any]) -> str:
        if not isinstance(context, dict):
            return ""
        for key in ("nl_query", "query", "message", "description", "content", "task"):
            value = context.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
        return ""

    def _normalize_agent_id_or_empty(self, value: Any) -> str:
        text = str(value or "").strip()
        if not text:
            return ""
        try:
            return str(uuid.UUID(text))
        except ValueError:
            return ""

    def _merge_norm_hits(
        self,
        *,
        base_norms: List[Dict[str, Any]],
        agent_norms: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        merged: List[Dict[str, Any]] = []
        seen: set[str] = set()
        for norm in list(agent_norms) + list(base_norms):
            if not isinstance(norm, dict):
                continue
            norm_id = str(norm.get("id", "")).strip()
            if not norm_id or norm_id in seen:
                continue
            seen.add(norm_id)
            merged.append(norm)
        return merged

    def _zone_prompt_text(self, zone_name: str, zone_payload: Dict[str, Any]) -> str:
        zone = zone_payload if isinstance(zone_payload, dict) else {}
        canonical = self._canonical(zone)
        nl = zone.get("nl", {}) if isinstance(zone, dict) else {}
        default_lang = ""
        nl_text = ""
        if isinstance(nl, dict):
            default_lang = str(nl.get("default_lang", "")).strip()
            texts = nl.get("texts", {})
            if isinstance(texts, dict):
                if default_lang and isinstance(texts.get(default_lang), str):
                    nl_text = texts.get(default_lang, "").strip()
                if not nl_text:
                    for value in texts.values():
                        if isinstance(value, str) and value.strip():
                            nl_text = value.strip()
                            break
        canonical_text = json.dumps(canonical, ensure_ascii=False, sort_keys=True)
        if nl_text and default_lang:
            return f"[{zone_name}.nl:{default_lang}] {nl_text} | canonical={canonical_text}"
        if nl_text:
            return f"[{zone_name}.nl] {nl_text} | canonical={canonical_text}"
        return f"canonical={canonical_text}"

    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_norms (canonical+nl text) and 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 {}
