"""
Fix Interceptors - Hook into LLM/tool calls to apply fixes.

Interceptors are registered with the runtime and called
during agent execution to apply relevant fixes.
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, TypeVar

from risicare.runtime.applier import ApplyResult, FixApplier
from risicare.runtime.cache import FixCache
from risicare.runtime.config import ActiveFix, FixRuntimeConfig

logger = logging.getLogger(__name__)

T = TypeVar("T")


@dataclass
class InterceptContext:
    """Context for an intercept operation."""

    # Operation identification
    operation_type: str  # "llm_call", "tool_call", "agent_delegate"
    operation_name: str  # Model name, tool name, agent name

    # Session context for A/B testing
    session_id: Optional[str] = None
    trace_id: Optional[str] = None
    span_id: Optional[str] = None

    # Error context (if retrying after error)
    error_code: Optional[str] = None
    error_message: Optional[str] = None
    attempt: int = 1

    # Applied fixes tracking
    applied_fixes: List[ApplyResult] = field(default_factory=list)


class FixInterceptor(ABC):
    """
    Base class for fix interceptors.

    Interceptors are called at different points in the execution:
    - pre_call: Before the LLM/tool call
    - post_call: After the call (for output validation)
    - on_error: When an error occurs (for retry/fallback)
    """

    @abstractmethod
    def pre_call(
        self,
        context: InterceptContext,
        messages: Optional[List[Dict[str, Any]]] = None,
        params: Optional[Dict[str, Any]] = None,
    ) -> tuple[Optional[List[Dict[str, Any]]], Optional[Dict[str, Any]]]:
        """
        Called before an LLM/tool call.

        Args:
            context: Intercept context.
            messages: LLM messages (if applicable).
            params: Call parameters.

        Returns:
            Modified (messages, params).
        """
        pass

    @abstractmethod
    def post_call(
        self,
        context: InterceptContext,
        response: Any,
    ) -> tuple[Any, bool]:
        """
        Called after an LLM/tool call.

        Args:
            context: Intercept context.
            response: Call response.

        Returns:
            (modified_response, should_continue)
        """
        pass

    @abstractmethod
    def on_error(
        self,
        context: InterceptContext,
        error: Exception,
    ) -> tuple[bool, Optional[Dict[str, Any]]]:
        """
        Called when an error occurs.

        Args:
            context: Intercept context.
            error: The exception.

        Returns:
            (should_retry, modified_params)
        """
        pass


class DefaultFixInterceptor(FixInterceptor):
    """
    Default interceptor implementation.

    Applies fixes based on error codes and operation context.
    """

    def __init__(
        self,
        config: Optional[FixRuntimeConfig] = None,
        cache: Optional[FixCache] = None,
        applier: Optional[FixApplier] = None,
    ) -> None:
        """
        Initialize the interceptor.

        Args:
            config: Runtime configuration.
            cache: Fix cache.
            applier: Fix applier.
        """
        self._config = config or FixRuntimeConfig()
        self._cache = cache or FixCache(self._config)
        self._applier = applier or FixApplier(self._config, self._cache)

    def pre_call(
        self,
        context: InterceptContext,
        messages: Optional[List[Dict[str, Any]]] = None,
        params: Optional[Dict[str, Any]] = None,
    ) -> tuple[Optional[List[Dict[str, Any]]], Optional[Dict[str, Any]]]:
        """Apply pre-call fixes."""
        if not self._config.enabled:
            return messages, params

        modified_messages = messages
        modified_params = params or {}

        # If retrying after error, check for applicable fixes
        if context.error_code:
            fix = self._applier.get_fix_for_error(
                context.error_code,
                context.session_id,
            )

            if fix:
                logger.debug(f"Applying fix {fix.fix_id} for {context.error_code}")

                if fix.fix_type == "prompt" and messages:
                    modified_messages, result = self._applier.apply_prompt_fix(
                        fix, messages
                    )
                    context.applied_fixes.append(result)
                    self._applier.log_application(result)

                elif fix.fix_type == "parameter":
                    modified_params, result = self._applier.apply_parameter_fix(
                        fix, modified_params
                    )
                    context.applied_fixes.append(result)
                    self._applier.log_application(result)

                elif fix.fix_type == "fallback":
                    modified_params, result = self._applier.apply_fallback_fix(
                        fix, modified_params
                    )
                    context.applied_fixes.append(result)
                    self._applier.log_application(result)

        # Apply input guards if configured
        if messages and self._config.enabled:
            for msg in modified_messages or []:
                content = msg.get("content", "")
                if isinstance(content, str):
                    for fix in self._cache.get_all():
                        if fix.fix_type == "guard":
                            _, passed, result = self._applier.apply_guard_fix(
                                fix, content, is_input=True
                            )
                            if result.applied:
                                context.applied_fixes.append(result)
                                self._applier.log_application(result)

        return modified_messages, modified_params

    def post_call(
        self,
        context: InterceptContext,
        response: Any,
    ) -> tuple[Any, bool]:
        """Apply post-call fixes (output validation)."""
        if not self._config.enabled:
            return response, True

        # Apply output guards
        for fix in self._cache.get_all():
            if fix.fix_type != "guard":
                continue

            # Extract response content
            content = self._extract_response_content(response)
            if not content:
                continue

            _, passed, result = self._applier.apply_guard_fix(
                fix, content, is_input=False
            )

            if result.applied:
                context.applied_fixes.append(result)
                self._applier.log_application(result)

            if not passed:
                logger.warning(f"Output guard failed for fix {fix.fix_id}")
                # Could trigger retry here if configured

        return response, True

    def on_error(
        self,
        context: InterceptContext,
        error: Exception,
    ) -> tuple[bool, Optional[Dict[str, Any]]]:
        """Handle errors and decide on retry."""
        if not self._config.enabled:
            return False, None

        # Classify the error
        error_code = self._classify_error(error)
        context.error_code = error_code
        context.error_message = str(error)

        # Check for retry fix
        fix = self._applier.get_fix_for_error(error_code, context.session_id)
        if not fix:
            return False, None

        if fix.fix_type == "retry":
            config = fix.config
            max_retries = config.get("max_retries", 3)

            if context.attempt <= max_retries:
                logger.debug(
                    f"Retry fix {fix.fix_id}: attempt {context.attempt}/{max_retries}"
                )

                result = ApplyResult(
                    applied=True,
                    fix_id=fix.fix_id,
                    deployment_id=fix.deployment_id,
                    fix_type="retry",
                    modifications={"attempt": context.attempt},
                )
                context.applied_fixes.append(result)
                self._applier.log_application(result)

                return True, None

        # Check for fallback fix
        if fix.fix_type == "fallback":
            modified_params, result = self._applier.apply_fallback_fix(fix, {})
            if result.applied:
                context.applied_fixes.append(result)
                self._applier.log_application(result)
                return True, modified_params

        return False, None

    def _classify_error(self, error: Exception) -> str:
        """
        Classify an error into an error code.

        This is a simplified classification. In production,
        this would use the full taxonomy from risicare-core.
        """
        error_type = type(error).__name__
        error_msg = str(error).lower()

        # Timeout errors
        if "timeout" in error_type.lower() or "timeout" in error_msg:
            return "TOOL.EXECUTION.TIMEOUT"

        # Rate limit errors
        if "rate" in error_msg and "limit" in error_msg:
            return "TOOL.EXECUTION.RATE_LIMIT"

        # Connection errors
        if "connection" in error_type.lower() or "connection" in error_msg:
            return "TOOL.EXECUTION.CONNECTION_ERROR"

        # Authentication errors
        if "auth" in error_msg or "401" in error_msg or "403" in error_msg:
            return "TOOL.EXECUTION.AUTH_ERROR"

        # JSON parsing errors
        if "json" in error_type.lower() or "json" in error_msg:
            return "OUTPUT.FORMAT.JSON_INVALID"

        # Generic execution error
        return "TOOL.EXECUTION.FAILURE"

    def _extract_response_content(self, response: Any) -> Optional[str]:
        """Extract text content from response."""
        if isinstance(response, str):
            return response

        if isinstance(response, dict):
            # OpenAI-style response
            choices = response.get("choices", [])
            if choices:
                message = choices[0].get("message", {})
                return message.get("content")

        if hasattr(response, "content"):
            content = response.content
            if isinstance(content, str):
                return content
            if isinstance(content, list) and content:
                first = content[0]
                if hasattr(first, "text"):
                    return first.text

        return None


class CompositeInterceptor(FixInterceptor):
    """
    Combines multiple interceptors.

    Calls each interceptor in sequence.
    """

    def __init__(self, interceptors: Optional[List[FixInterceptor]] = None) -> None:
        """Initialize with list of interceptors."""
        self._interceptors = interceptors or []

    def add(self, interceptor: FixInterceptor) -> None:
        """Add an interceptor."""
        self._interceptors.append(interceptor)

    def pre_call(
        self,
        context: InterceptContext,
        messages: Optional[List[Dict[str, Any]]] = None,
        params: Optional[Dict[str, Any]] = None,
    ) -> tuple[Optional[List[Dict[str, Any]]], Optional[Dict[str, Any]]]:
        """Call all interceptors' pre_call in sequence."""
        for interceptor in self._interceptors:
            messages, params = interceptor.pre_call(context, messages, params)
        return messages, params

    def post_call(
        self,
        context: InterceptContext,
        response: Any,
    ) -> tuple[Any, bool]:
        """Call all interceptors' post_call in sequence."""
        for interceptor in self._interceptors:
            response, should_continue = interceptor.post_call(context, response)
            if not should_continue:
                return response, False
        return response, True

    def on_error(
        self,
        context: InterceptContext,
        error: Exception,
    ) -> tuple[bool, Optional[Dict[str, Any]]]:
        """Call interceptors until one handles the error."""
        for interceptor in self._interceptors:
            should_retry, params = interceptor.on_error(context, error)
            if should_retry:
                return True, params
        return False, None
