"""
Shared real-time remediation engine.

Framework-agnostic module that classifies tool errors, queries the Aigie platform
for matching remediation patterns/flows, and builds corrective guidance text.

Each SDK integration (Strands, Claude Agent SDK, LangGraph) uses this engine
and applies the guidance through its own framework-specific mechanism.
"""

import logging
import time
from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


@dataclass
class RemediationConfig:
    """Configuration for real-time remediation behavior."""
    enable_realtime_remediation: bool = False
    remediation_mode: str = "recommendation"  # "recommendation" | "autonomous"
    remediation_query_timeout: float = 2.0

    def __post_init__(self):
        if self.remediation_mode not in ("recommendation", "autonomous"):
            raise ValueError("remediation_mode must be 'recommendation' or 'autonomous'")
        if self.remediation_query_timeout <= 0:
            raise ValueError("remediation_query_timeout must be positive")


@dataclass
class RemediationResult:
    """Result of a remediation evaluation."""
    applied: bool = False
    error_type: str = ""
    strategy: str = ""
    success_rate: float = 0.0
    guidance_text: str = ""
    query_ms: float = 0.0
    tool_name: str = ""
    span_id: str = ""
    trace_id: str = ""
    mode: str = "recommendation"
    action_type: str = ""       # "retry", "fallback", "skip", "" (guidance-only)
    action_params: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


class ErrorClassifier:
    """Classify error messages into platform error types."""

    ERROR_TYPE_KEYWORDS: Dict[str, List[str]] = {
        "RATE_LIMIT_ERROR": ["rate limit", "quota", "429", "throttl"],
        "TIMEOUT_ERROR": ["timeout", "timed out", "deadline"],
        "CONNECTION_ERROR": ["connection", "econnrefused", "unavailable", "knowledge base"],
        "API_KEY_ERROR": ["api key", "unauthorized", "authentication"],
        "PERMISSION_ERROR": ["permission", "forbidden", "403"],
        "NOT_FOUND_ERROR": ["not found", "404"],
        "VALIDATION_ERROR": ["validation", "invalid", "schema"],
    }

    _CLASSIFY_RULES = [
        (("rate limit", "429", "too many requests", "quota"), "RATE_LIMIT_ERROR"),
        (("timeout", "timed out", "deadline exceeded"), "TIMEOUT_ERROR"),
        (("connection", "econnrefused", "econnreset", "network"), "CONNECTION_ERROR"),
        (("api key", "api_key", "unauthorized", "401", "authentication"), "API_KEY_ERROR"),
        (("permission", "forbidden", "403"), "PERMISSION_ERROR"),
        (("not found", "notfound", "404"), "NOT_FOUND_ERROR"),
        (("validation", "invalid", "schema", "400"), "VALIDATION_ERROR"),
    ]

    def classify(self, error_msg: str, tool_name: str = "") -> Optional[str]:
        """Classify an error message into a platform error type."""
        if not error_msg:
            return None
        msg_lower = error_msg.lower()
        for keywords, error_type in self._CLASSIFY_RULES:
            if any(k in msg_lower for k in keywords):
                return error_type
        return "UNKNOWN_ERROR"


class RemediationEngine:
    """Core remediation engine — queries platform, builds guidance, caches patterns.

    Framework-agnostic: each integration adapter calls ``evaluate()`` and then
    decides how to inject the resulting guidance into its own event model.
    """

    def __init__(
        self,
        api_url: str,
        api_key: str,
        query_timeout: float = 2.0,
    ):
        self._api_url = api_url
        self._api_key = api_key
        self._query_timeout = query_timeout
        self._classifier = ErrorClassifier()

        # Internal state (reset per trace via ``reset()``)
        self._pattern_cache: Dict[str, List[Dict[str, Any]]] = {}
        self._results: List[RemediationResult] = []
        self._applied_count = 0

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def evaluate(
        self,
        error_msg: str,
        tool_name: str,
        span_id: str,
        trace_id: str,
        mode: str = "recommendation",
    ) -> Optional[RemediationResult]:
        """Classify error → query pattern → build guidance.

        Returns a ``RemediationResult`` for the adapter to apply, or ``None``
        if no matching pattern was found.
        """
        error_type = self._classifier.classify(error_msg, tool_name)
        if not error_type:
            return None

        start_ts = time.monotonic()
        pattern = await self._query_pattern(error_type)
        query_ms = (time.monotonic() - start_ts) * 1000

        if not pattern:
            logger.debug(f"[AIGIE] No remediation pattern for {error_type}")
            return None

        strategy = pattern.get("strategy", "unknown")
        success_rate = pattern.get("success_rate", 0)
        guidance_text = self._build_guidance(error_type, pattern)

        # Derive autonomous action from strategy
        action_type = ""
        action_params: Dict[str, Any] = {}
        if mode == "autonomous":
            if strategy == "retry_with_context":
                action_type = "retry"
                # Extract retry params from pattern steps if available
                max_retries = 3
                backoff_seconds = 1
                for step in (pattern.get("steps_detail") or []):
                    params = step.get("parameters", {})
                    if params.get("max_retries"):
                        max_retries = int(params["max_retries"])
                    if params.get("backoff_seconds"):
                        backoff_seconds = int(params["backoff_seconds"])
                action_params = {"max_retries": max_retries, "backoff_seconds": backoff_seconds}
            elif strategy == "fallback_tool":
                action_type = "fallback"
            elif strategy == "skip_step":
                action_type = "skip"

        result = RemediationResult(
            applied=False,
            error_type=error_type,
            strategy=strategy,
            success_rate=success_rate,
            guidance_text=guidance_text,
            query_ms=query_ms,
            tool_name=tool_name,
            span_id=span_id,
            trace_id=trace_id,
            mode=mode,
            action_type=action_type,
            action_params=action_params,
        )
        self._results.append(result)
        return result

    def mark_applied(self, result: RemediationResult) -> None:
        """Mark a result as applied (called by adapter after injection)."""
        result.applied = True
        self._applied_count += 1

    def reset(self) -> None:
        """Clear cache and stats — call at the end of each trace."""
        self._pattern_cache.clear()
        self._results.clear()
        self._applied_count = 0

    @property
    def applied_count(self) -> int:
        return self._applied_count

    @property
    def results(self) -> List[RemediationResult]:
        return list(self._results)

    # ------------------------------------------------------------------
    # Platform queries
    # ------------------------------------------------------------------

    async def _query_pattern(self, error_type: str) -> Optional[Dict[str, Any]]:
        """Query the platform for a remediation pattern matching this error type.

        Tries two sources:
        1. Learning system patterns (GET /v1/remediation/patterns)
        2. Remediation flows (GET /v1/remediation/flows) — matched by name/description
        """
        if error_type in self._pattern_cache:
            patterns = self._pattern_cache[error_type]
            return patterns[0] if patterns else None

        if not self._api_url:
            return None

        headers: Dict[str, str] = {}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
            headers["X-API-Key"] = self._api_key

        try:
            import httpx
            async with httpx.AsyncClient(timeout=self._query_timeout) as client:
                # Source 1: Learning system patterns
                try:
                    resp = await client.get(
                        f"{self._api_url}/v1/remediation/patterns",
                        params={"error_type": error_type, "min_success_rate": "0.5", "limit": "5"},
                        headers=headers,
                    )
                    if resp.status_code == 200:
                        data = resp.json()
                        patterns = data.get("patterns", [])
                        if patterns:
                            self._pattern_cache[error_type] = patterns
                            return patterns[0]
                except Exception:
                    pass

                # Source 2: Remediation flows — match by name/description keywords
                try:
                    resp = await client.get(f"{self._api_url}/v1/remediation/flows", headers=headers)
                    if resp.status_code == 200:
                        data = resp.json()
                        flows = data if isinstance(data, list) else data.get("flows", [])
                        keywords = self._classifier.ERROR_TYPE_KEYWORDS.get(
                            error_type, [error_type.lower()]
                        )
                        matched = self._match_flow(flows, keywords)
                        if matched:
                            self._pattern_cache[error_type] = [matched]
                            return matched
                except Exception:
                    pass

                self._pattern_cache[error_type] = []
                return None
        except Exception as e:
            logger.debug(f"[AIGIE] Remediation query failed: {e}")
            self._pattern_cache[error_type] = []
            return None

    # ------------------------------------------------------------------
    # Flow matching
    # ------------------------------------------------------------------

    @staticmethod
    def _match_flow(
        flows: List[Dict[str, Any]], keywords: List[str]
    ) -> Optional[Dict[str, Any]]:
        """Match a remediation flow to error keywords and convert to pattern format."""
        for flow in flows:
            name = (flow.get("name", "") or "").lower()
            desc = (flow.get("description", "") or "").lower()
            text = f"{name} {desc}"
            if any(kw in text for kw in keywords):
                steps = flow.get("steps", [])
                strategy = "retry_with_context"
                method_parts: List[str] = []
                for step in steps:
                    action = step.get("action", "")
                    if "retry" in action:
                        strategy = "retry_with_context"
                    elif "fallback" in action:
                        strategy = "fallback_tool"
                    elif "escalat" in action:
                        strategy = "escalate"
                    params = step.get("parameters", {})
                    if params.get("max_retries"):
                        method_parts.append(
                            f"Retry up to {params['max_retries']}x with "
                            f"{params.get('backoff_seconds', 5)}s backoff"
                        )
                    if params.get("fallback_tool"):
                        method_parts.append(f"Fallback to {params['fallback_tool']}")
                    if params.get("message"):
                        method_parts.append(f"Escalate: {params['message'][:80]}")
                return {
                    "flow_id": flow.get("id"),
                    "flow_name": flow.get("name"),
                    "strategy": strategy,
                    "method": ". ".join(method_parts) if method_parts else (flow.get("description", "") or "")[:200],
                    "success_rate": flow.get("success_rate", 0) or 0,
                    "steps": len(steps),
                }
        return None

    # ------------------------------------------------------------------
    # Guidance building
    # ------------------------------------------------------------------

    @staticmethod
    def _build_guidance(error_type: str, pattern: Dict[str, Any]) -> str:
        """Build the ``[Kytte AI Real-Time Correction]`` guidance text."""
        strategy = pattern.get("strategy", "unknown")
        success_rate = pattern.get("success_rate", 0)
        method = pattern.get("method", "")

        parts = [
            "\n[Kytte AI Real-Time Correction]",
            f"Error type: {error_type} (known pattern, {success_rate:.0%} fix rate)",
            f"Recommended strategy: {strategy}",
        ]
        if method:
            parts.append(f"Fix method: {method}")

        _strategy_actions = {
            "retry_with_context": "Action: Retry this tool call. If it fails again, try an alternative approach.",
            "fallback_model": "Action: Use a different data source or approach instead of retrying.",
            "inject_instruction": "Action: Modify your approach based on the error and continue.",
            "skip_step": "Action: Skip this step and proceed with available data.",
        }
        parts.append(
            _strategy_actions.get(strategy, "Action: Adapt your approach to work around this error.")
        )
        return "\n".join(parts)
