"""
Fix Applier - Applies fixes to agent operations.

Handles different fix types:
- PromptFix: Modify system/user prompts
- ParameterFix: Adjust LLM parameters
- RetryFix: Add retry logic with backoff
- FallbackFix: Switch to fallback model/strategy
- GuardFix: Add input/output validation
"""

from __future__ import annotations

import hashlib
import logging
import random
import re
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar

from risicare.runtime.cache import FixCache
from risicare.runtime.config import ActiveFix, FixRuntimeConfig

logger = logging.getLogger(__name__)

T = TypeVar("T")

# Maximum length for fix content injected into prompts.
# Prevents unbounded prompt injection from malicious/corrupted fix configs.
MAX_FIX_CONTENT_LENGTH = 10_000


@dataclass
class ApplyResult:
    """Result of applying a fix."""

    applied: bool
    fix_id: Optional[str] = None
    deployment_id: Optional[str] = None
    fix_type: Optional[str] = None
    modifications: Dict[str, Any] = field(default_factory=dict)
    error: Optional[str] = None


class FixApplier:
    """
    Applies fixes to agent operations.

    Each fix type has a specific application method that modifies
    the operation in a targeted way.
    """

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

        Args:
            config: Runtime configuration.
            cache: Fix cache.
        """
        self._config = config or FixRuntimeConfig()
        self._cache = cache or FixCache(self._config)
        self._application_log: List[ApplyResult] = []

    def get_fix_for_error(
        self,
        error_code: str,
        session_id: Optional[str] = None,
    ) -> Optional[ActiveFix]:
        """
        Get applicable fix for an error code.

        Considers A/B testing bucket if session_id provided.

        Args:
            error_code: Error code to match.
            session_id: Optional session ID for A/B bucketing.

        Returns:
            ActiveFix if found and should be applied.
        """
        fix = self._cache.get(error_code)
        if not fix:
            return None

        # Check A/B test bucketing (deterministic across process restarts)
        if session_id and self._config.ab_testing_enabled:
            session_hash = int(hashlib.md5(session_id.encode()).hexdigest()[:8], 16)
            if not fix.should_apply(session_hash):
                logger.debug(f"Fix {fix.fix_id} not applied (A/B control group)")
                return None

        return fix

    def apply_prompt_fix(
        self,
        fix: ActiveFix,
        messages: List[Dict[str, Any]],
    ) -> Tuple[List[Dict[str, Any]], ApplyResult]:
        """
        Apply a prompt fix to messages.

        Supports:
        - prepend: Add content before system message
        - append: Add content after system message
        - replace: Replace matching pattern in messages
        - few_shot: Add example messages

        Args:
            fix: The fix to apply.
            messages: Current message list.

        Returns:
            Modified messages and result.
        """
        config = fix.config
        modification_type = config.get("modification_type", "append")
        target = config.get("target", "system")
        content = config.get("content", "")

        # Truncate oversized content to prevent prompt injection
        if len(content) > MAX_FIX_CONTENT_LENGTH:
            logger.warning(
                "Fix %s content truncated from %d to %d chars",
                fix.fix_id, len(content), MAX_FIX_CONTENT_LENGTH,
            )
            content = content[:MAX_FIX_CONTENT_LENGTH]

        result = ApplyResult(
            applied=False,
            fix_id=fix.fix_id,
            deployment_id=fix.deployment_id,
            fix_type="prompt",
        )

        if self._config.dry_run:
            result.modifications = {
                "type": modification_type,
                "target": target,
                "content_preview": content[:100] if content else "",
            }
            return messages, result

        modified = list(messages)

        if modification_type == "prepend":
            modified = self._prepend_to_target(modified, target, content)
            result.applied = True
            result.modifications = {"type": "prepend", "target": target}

        elif modification_type == "append":
            modified = self._append_to_target(modified, target, content)
            result.applied = True
            result.modifications = {"type": "append", "target": target}

        elif modification_type == "replace":
            pattern = config.get("pattern", "")
            if pattern:
                modified = self._replace_in_messages(modified, pattern, content)
                result.applied = True
                result.modifications = {"type": "replace", "pattern": pattern}

        elif modification_type == "few_shot":
            examples = config.get("examples", [])
            if examples:
                modified = self._add_few_shot_examples(modified, examples)
                result.applied = True
                result.modifications = {"type": "few_shot", "count": len(examples)}

        return modified, result

    def apply_parameter_fix(
        self,
        fix: ActiveFix,
        params: Dict[str, Any],
    ) -> Tuple[Dict[str, Any], ApplyResult]:
        """
        Apply a parameter fix to LLM parameters.

        Args:
            fix: The fix to apply.
            params: Current parameters.

        Returns:
            Modified parameters and result.
        """
        config = fix.config
        new_params = config.get("parameters", {})

        result = ApplyResult(
            applied=False,
            fix_id=fix.fix_id,
            deployment_id=fix.deployment_id,
            fix_type="parameter",
        )

        if self._config.dry_run:
            result.modifications = {"parameters": new_params}
            return params, result

        modified = dict(params)
        modified.update(new_params)

        result.applied = True
        result.modifications = {"parameters": new_params}

        return modified, result

    def apply_retry_fix(
        self,
        fix: ActiveFix,
        operation: Callable[[], T],
    ) -> Tuple[T, ApplyResult]:
        """
        Apply a retry fix to an operation.

        Args:
            fix: The fix to apply.
            operation: Function to retry.

        Returns:
            Operation result and apply result.
        """
        config = fix.config
        max_retries = config.get("max_retries", 3)
        initial_delay_ms = config.get("initial_delay_ms", 1000)
        max_delay_ms = config.get("max_delay_ms", 30000)
        exponential_base = config.get("exponential_base", 2.0)
        jitter = config.get("jitter", True)
        retry_on = config.get("retry_on", [])

        result = ApplyResult(
            applied=False,
            fix_id=fix.fix_id,
            deployment_id=fix.deployment_id,
            fix_type="retry",
        )

        if self._config.dry_run:
            result.modifications = {
                "max_retries": max_retries,
                "initial_delay_ms": initial_delay_ms,
            }
            return operation(), result

        last_error = None
        attempts = 0

        for attempt in range(max_retries + 1):
            attempts += 1
            try:
                op_result = operation()
                result.applied = True
                result.modifications = {"attempts": attempts}
                return op_result, result

            except Exception as e:
                last_error = e
                error_type = type(e).__name__

                # Check if we should retry this error
                if retry_on and error_type not in retry_on:
                    raise

                if attempt < max_retries:
                    # Calculate delay with exponential backoff
                    delay_ms = min(
                        initial_delay_ms * (exponential_base ** attempt),
                        max_delay_ms,
                    )

                    # Add jitter
                    if jitter:
                        delay_ms *= 0.5 + random.random()

                    logger.debug(
                        f"Retry {attempt + 1}/{max_retries} after {delay_ms:.0f}ms"
                    )
                    time.sleep(delay_ms / 1000)

        # All retries exhausted
        result.applied = True
        result.error = str(last_error)
        result.modifications = {"attempts": attempts, "exhausted": True}
        raise last_error  # type: ignore

    async def apply_retry_fix_async(
        self,
        fix: ActiveFix,
        operation: Callable[[], Any],
    ) -> Tuple[Any, ApplyResult]:
        """
        Apply a retry fix to an async operation.

        Args:
            fix: The fix to apply.
            operation: Async function to retry.

        Returns:
            Operation result and apply result.
        """
        import asyncio

        config = fix.config
        max_retries = config.get("max_retries", 3)
        initial_delay_ms = config.get("initial_delay_ms", 1000)
        max_delay_ms = config.get("max_delay_ms", 30000)
        exponential_base = config.get("exponential_base", 2.0)
        jitter = config.get("jitter", True)
        retry_on = config.get("retry_on", [])

        result = ApplyResult(
            applied=False,
            fix_id=fix.fix_id,
            deployment_id=fix.deployment_id,
            fix_type="retry",
        )

        if self._config.dry_run:
            result.modifications = {"max_retries": max_retries}
            return await operation(), result

        last_error = None
        attempts = 0

        for attempt in range(max_retries + 1):
            attempts += 1
            try:
                op_result = await operation()
                result.applied = True
                result.modifications = {"attempts": attempts}
                return op_result, result

            except Exception as e:
                last_error = e
                error_type = type(e).__name__

                if retry_on and error_type not in retry_on:
                    raise

                if attempt < max_retries:
                    delay_ms = min(
                        initial_delay_ms * (exponential_base ** attempt),
                        max_delay_ms,
                    )
                    if jitter:
                        delay_ms *= 0.5 + random.random()

                    await asyncio.sleep(delay_ms / 1000)

        result.applied = True
        result.error = str(last_error)
        result.modifications = {"attempts": attempts, "exhausted": True}
        raise last_error  # type: ignore

    def apply_fallback_fix(
        self,
        fix: ActiveFix,
        params: Dict[str, Any],
    ) -> Tuple[Dict[str, Any], ApplyResult]:
        """
        Apply a fallback fix (e.g., switch to different model).

        Args:
            fix: The fix to apply.
            params: Current parameters.

        Returns:
            Modified parameters and result.
        """
        config = fix.config
        fallback_type = config.get("fallback_type", "model")
        fallback_config = config.get("fallback_config", {})

        result = ApplyResult(
            applied=False,
            fix_id=fix.fix_id,
            deployment_id=fix.deployment_id,
            fix_type="fallback",
        )

        if self._config.dry_run:
            result.modifications = {
                "fallback_type": fallback_type,
                "fallback_config": fallback_config,
            }
            return params, result

        modified = dict(params)

        if fallback_type == "model":
            # Switch to fallback model
            fallback_model = fallback_config.get("model")
            if fallback_model:
                modified["model"] = fallback_model
                result.applied = True
                result.modifications = {"model": fallback_model}

        elif fallback_type == "default":
            # Use default response
            default_response = fallback_config.get("response")
            if default_response:
                modified["_fallback_response"] = default_response
                result.applied = True
                result.modifications = {"default_response": True}

        return modified, result

    def apply_guard_fix(
        self,
        fix: ActiveFix,
        content: str,
        is_input: bool = True,
    ) -> Tuple[str, bool, ApplyResult]:
        """
        Apply a guard fix (validation).

        Args:
            fix: The fix to apply.
            content: Content to validate.
            is_input: Whether this is input (vs output) validation.

        Returns:
            (content, passed_validation, result)
        """
        config = fix.config
        guard_type = config.get("guard_type", "output_validation")
        guard_config = config.get("guard_config", {})

        result = ApplyResult(
            applied=False,
            fix_id=fix.fix_id,
            deployment_id=fix.deployment_id,
            fix_type="guard",
        )

        # Check if guard applies to this direction
        if is_input and guard_type not in ("input_validation", "content_filter"):
            return content, True, result
        if not is_input and guard_type not in ("output_validation", "format_check"):
            return content, True, result

        if self._config.dry_run:
            result.modifications = {"guard_type": guard_type}
            return content, True, result

        passed = True

        if guard_type == "content_filter":
            # Filter blocked patterns (validate regex before use)
            blocked_patterns = guard_config.get("blocked_patterns", [])
            for pattern in blocked_patterns:
                compiled = self._safe_compile(pattern, re.IGNORECASE)
                if compiled and compiled.search(content):
                    passed = False
                    break

        elif guard_type == "format_check":
            # Check format (e.g., JSON valid)
            required_format = guard_config.get("format")
            if required_format == "json":
                import json as json_module
                try:
                    json_module.loads(content)
                except json_module.JSONDecodeError:
                    passed = False

        elif guard_type in ("input_validation", "output_validation"):
            # Check length constraints
            min_length = guard_config.get("min_length", 0)
            max_length = guard_config.get("max_length", float("inf"))
            if len(content) < min_length or len(content) > max_length:
                passed = False

        result.applied = True
        result.modifications = {"passed": passed, "guard_type": guard_type}

        return content, passed, result

    # ========================================================================
    # Helper Methods
    # ========================================================================

    @staticmethod
    def _safe_compile(
        pattern: str,
        flags: int = 0,
        max_length: int = 500,
    ) -> Optional[re.Pattern]:
        """Safely compile a regex pattern from untrusted input.

        Returns None if the pattern is invalid or too long.
        """
        if not pattern or len(pattern) > max_length:
            return None
        try:
            return re.compile(pattern, flags)
        except re.error:
            logger.warning(f"Invalid regex pattern (length={len(pattern)})")
            return None

    def _prepend_to_target(
        self,
        messages: List[Dict[str, Any]],
        target: str,
        content: str,
    ) -> List[Dict[str, Any]]:
        """Prepend content to target message."""
        result = []
        for msg in messages:
            if msg.get("role") == target:
                modified = dict(msg)
                modified["content"] = content + "\n\n" + modified.get("content", "")
                result.append(modified)
            else:
                result.append(msg)
        return result

    def _append_to_target(
        self,
        messages: List[Dict[str, Any]],
        target: str,
        content: str,
    ) -> List[Dict[str, Any]]:
        """Append content to target message."""
        result = []
        for msg in messages:
            if msg.get("role") == target:
                modified = dict(msg)
                modified["content"] = modified.get("content", "") + "\n\n" + content
                result.append(modified)
            else:
                result.append(msg)
        return result

    def _replace_in_messages(
        self,
        messages: List[Dict[str, Any]],
        pattern: str,
        replacement: str,
    ) -> List[Dict[str, Any]]:
        """Replace pattern in all messages."""
        compiled = self._safe_compile(pattern)
        if compiled is None:
            logger.warning(f"Skipping invalid regex pattern in prompt fix")
            return messages

        result = []
        for msg in messages:
            modified = dict(msg)
            content = modified.get("content", "")
            if isinstance(content, str):
                modified["content"] = compiled.sub(replacement, content)
            result.append(modified)
        return result

    def _add_few_shot_examples(
        self,
        messages: List[Dict[str, Any]],
        examples: List[Dict[str, str]],
    ) -> List[Dict[str, Any]]:
        """Add few-shot examples after system message."""
        result = []
        system_found = False

        for msg in messages:
            result.append(msg)
            if msg.get("role") == "system" and not system_found:
                system_found = True
                # Add examples after system message
                for example in examples:
                    if "user" in example:
                        result.append({"role": "user", "content": example["user"]})
                    if "assistant" in example:
                        result.append(
                            {"role": "assistant", "content": example["assistant"]}
                        )

        return result

    def log_application(self, result: ApplyResult) -> None:
        """Log a fix application result."""
        self._application_log.append(result)

        # Limit log size
        if len(self._application_log) > 1000:
            self._application_log = self._application_log[-500:]

    def get_application_log(self) -> List[ApplyResult]:
        """Get the application log."""
        return list(self._application_log)
