"""
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 asdict, dataclass, field
from typing import Any, Dict, List

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)
    confidence_band: str = ""  # "high", "medium", "low" — drives graduated decisions
    pattern_id: str = ""  # Platform pattern/flow ID for closed-loop learning

    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 = "") -> str | None:
        """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.

    When a RemediationLoop is available (from the main Aigie client), this engine
    delegates to it — gaining access to LLM judges, context aggregation, and the
    full runtime pipeline. Without a loop, falls back to direct HTTP queries.
    """

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

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

    def set_remediation_loop(self, loop: Any) -> None:
        """Set the RemediationLoop for delegated evaluation."""
        self._remediation_loop = loop

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

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

        Delegates to RemediationLoop when available for full pipeline access.
        Falls back to direct HTTP queries when RemediationLoop is not initialized.

        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)

        # Delegate to RemediationLoop if available (unified path)
        if self._remediation_loop and error_type:
            try:
                loop_result = await self._remediation_loop.process_span(
                    span_id=span_id,
                    trace_id=trace_id,
                    input_messages=[],
                    output_content="",
                    error=Exception(error_msg),
                )
                if loop_result and loop_result.success:
                    loop_confidence = loop_result.confidence or 0.0
                    loop_band = self._confidence_band(loop_confidence)
                    loop_strategy = (
                        loop_result.strategy_used.value if loop_result.strategy_used else "unknown"
                    )
                    loop_action, loop_params = "", {}
                    if loop_band == "high":
                        loop_action, loop_params = self._derive_action(loop_strategy, {})
                    elif loop_band == "medium" and mode == "autonomous":
                        mode = "recommendation"
                        loop_action, loop_params = self._derive_action(loop_strategy, {})
                    return RemediationResult(
                        applied=False,
                        error_type=error_type,
                        strategy=loop_strategy,
                        success_rate=loop_confidence,
                        guidance_text=loop_result.recommendation or "",
                        tool_name=tool_name,
                        span_id=span_id,
                        trace_id=trace_id,
                        mode=mode,
                        action_type=loop_action,
                        action_params=loop_params,
                        confidence_band=loop_band,
                    )
            except Exception as e:
                logger.debug(f"RemediationLoop delegation failed, falling back to HTTP: {e}")
        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)
        confidence = pattern.get("confidence", success_rate)
        guidance_text = self._build_guidance(error_type, pattern)

        # Confidence-graduated decisions:
        # - HIGH  (≥0.85): auto-apply fix in autonomous mode, strong recommend otherwise
        # - MEDIUM (0.60–0.85): recommend with action hint, never auto-apply
        # - LOW   (<0.60): guidance only, gather more data
        band = self._confidence_band(confidence)

        action_type = ""
        action_params: Dict[str, Any] = {}
        if band == "high":
            action_type, action_params = self._derive_action(strategy, pattern)
        elif band == "medium" and mode == "autonomous":
            # Downgrade to recommendation — not confident enough to auto-apply
            mode = "recommendation"
            action_type, action_params = self._derive_action(strategy, pattern)
        # LOW band: no action_type, guidance text only

        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,
            confidence_band=band,
            pattern_id=str(pattern.get("flow_id") or pattern.get("id") or ""),
        )
        self._results.append(result)
        return result

    def to_fix_action(self, result: RemediationResult):
        """Convert a RemediationResult to a structured FixAction for AutoFixApplicator.

        Returns None when the result has no actionable action_type or when
        the action_type doesn't map to a known FixActionType.
        """
        if not result.action_type:
            return None

        from ..interceptor.protocols import FixAction, FixActionType

        ACTION_MAP = {
            "retry": FixActionType.RETRY,
            "fallback": FixActionType.FALLBACK,
            "skip": FixActionType.TRUNCATE_CONTEXT,
            "truncate_context": FixActionType.TRUNCATE_CONTEXT,
            "modify_response": FixActionType.MODIFY_RESPONSE,
            "inject_instruction": FixActionType.INJECT_INSTRUCTION,
        }
        fix_type = ACTION_MAP.get(result.action_type)
        if not fix_type:
            return None

        return FixAction(
            action_type=fix_type,
            parameters=result.action_params,
            confidence=result.success_rate,
            source="backend",
            reason=result.guidance_text,
        )

    def mark_applied(self, result: RemediationResult, success: bool = True) -> None:
        """Mark a result as applied and report outcome for closed-loop learning.

        Args:
            result: The remediation result that was applied
            success: Whether the fix actually resolved the issue
        """
        result.applied = True
        self._applied_count += 1

        # Fire-and-forget outcome report to backend
        try:
            import asyncio

            loop = asyncio.get_running_loop()
            loop.create_task(self._report_fix_outcome(result, success))
        except RuntimeError:
            pass  # No event loop — skip reporting

    async def _report_fix_outcome(self, result: RemediationResult, success: bool) -> None:
        """Report fix outcome to backend for closed-loop pattern learning.

        POSTs to /v1/remediation/fix-outcome so the learning system can
        update pattern effectiveness scores and promote/demote autonomous patterns.
        """
        if not self._api_url:
            return
        try:
            import httpx

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

            payload = {
                "trace_id": result.trace_id or "",
                "span_id": result.span_id or "",
                "pattern_id": result.pattern_id or "",
                "flow_id": getattr(result, "flow_id", None) or "",
                "fix_action": result.action_type or result.strategy,
                "success": success,
                "latency_ms": result.query_ms,
                "error_before": result.error_type,
                "metadata": {
                    "tool_name": result.tool_name,
                    "strategy": result.strategy,
                    "confidence": result.success_rate,
                    "guidance_applied": bool(result.guidance_text),
                    "source": "sdk_realtime_outcome",
                },
            }
            async with httpx.AsyncClient(timeout=2.0) as client:
                await client.post(
                    f"{self._api_url}/v1/remediation/fix-outcome",
                    json=payload,
                    headers=headers,
                )
        except Exception:
            pass  # Non-blocking, best-effort

    # ------------------------------------------------------------------
    # Confidence-graduated helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _confidence_band(confidence: float) -> str:
        """Classify confidence into graduated decision bands.

        HIGH  (≥0.85): Safe to auto-apply — pattern proven across 10+ samples
                       with Wilson score lower bound above threshold.
        MEDIUM (0.60–0.85): Recommend but don't auto-apply — pattern shows
                            promise but needs more data or has edge cases.
        LOW   (<0.60): Guidance only — pattern is tentative, gather more data.
        """
        if confidence >= 0.85:
            return "high"
        if confidence >= 0.60:
            return "medium"
        return "low"

    @staticmethod
    def _derive_action(strategy: str, pattern: Dict[str, Any]) -> tuple:
        """Derive action_type and action_params from strategy.

        Extracted from evaluate() so both high and medium bands can use it.
        """
        action_type = ""
        action_params: Dict[str, Any] = {}
        if strategy == "retry_with_context":
            action_type = "retry"
            max_retries = 3
            backoff_seconds = 1
            for step in pattern.get("steps_detail") or []:
                params = step.get("parameters", {})
                try:
                    if params.get("max_retries"):
                        max_retries = int(float(params["max_retries"]))
                    if params.get("backoff_seconds"):
                        backoff_seconds = int(float(params["backoff_seconds"]))
                except (ValueError, TypeError):
                    pass  # Use defaults if pattern data is malformed
            action_params = {"max_retries": max_retries, "backoff_seconds": backoff_seconds}
        elif strategy == "fallback_tool":
            action_type = "fallback"
        elif strategy == "skip_step":
            action_type = "skip"
        return action_type, action_params

    async def report_result(self, result: RemediationResult, trace_id: str | None = None) -> None:
        """Report a remediation result to the platform for closed-loop learning.

        POSTs to /v1/remediation/results so the learning system can track
        pattern effectiveness and improve autonomous promotion decisions.
        All integrations (Strands, LangChain, CrewAI, etc.) share this path.
        """
        if not self._api_url:
            return
        try:
            import httpx

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

            payload = {
                "trace_id": trace_id or result.trace_id or "",
                "span_id": result.span_id or None,
                "strategy": result.strategy,
                "method": result.action_type or result.strategy,
                "success": result.applied,
                "confidence": result.success_rate,
                "latency_ms": result.query_ms,
                "fix_applied": result.applied,
                "mode": result.mode,
                "metadata": {
                    "error_type": result.error_type,
                    "tool_name": result.tool_name,
                    "action_type": result.action_type,
                    "action_params": result.action_params,
                    "source": "sdk_realtime",
                    "guidance_text": result.guidance_text[:500] if result.guidance_text else "",
                },
            }
            async with httpx.AsyncClient(timeout=2.0) as client:
                await client.post(
                    f"{self._api_url}/v1/remediation/results",
                    json=payload,
                    headers=headers,
                )
        except Exception:
            pass  # Non-blocking, best-effort

    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) -> Dict[str, Any] | None:
        """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

                # Do NOT cache empty/failed results — allow retry on next call
                return None
        except Exception as e:
            logger.debug(f"[AIGIE] Remediation query failed: {e}")
            # Do NOT cache failures — platform may recover on next attempt
            return None

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

    @staticmethod
    def _match_flow(flows: List[Dict[str, Any]], keywords: List[str]) -> Dict[str, Any] | None:
        """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 ``[Aigie Real-Time Correction]`` guidance text."""
        strategy = pattern.get("strategy", "unknown")
        success_rate = pattern.get("success_rate", 0)
        confidence = pattern.get("confidence", success_rate)
        method = pattern.get("method", "")
        band = RemediationEngine._confidence_band(confidence)

        confidence_label = {
            "high": "proven fix",
            "medium": "likely fix",
            "low": "possible fix",
        }[band]

        parts = [
            "\n[Aigie Real-Time Correction]",
            f"Error type: {error_type} ({confidence_label}, {success_rate:.0%} success 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)
