"""Base patcher class for ML framework instrumentation.

This module provides shared logic for patching ML frameworks with
OpenTelemetry instrumentation using wrapt for robust monkeypatching.
"""

import logging
import sys
import time
from typing import Any, Callable, Dict, Optional

import wrapt

from ..core.client import MCAClient

logger = logging.getLogger(__name__)

# Global registry to track which methods have been patched
# Format: {(class, method_name): original_method}
_patched_methods: Dict[tuple, Callable] = {}

# Track if we've logged a warning about missing client (avoid spam)
_missing_client_warned = False


class BasePatcher:
    """Base class for framework patchers.

    Provides shared logic for monkey-patching ML framework classes
    with OpenTelemetry instrumentation. Uses wrapt library for robust
    monkeypatching that preserves signatures and cooperates with other tools.
    """

    def __init__(
        self,
        framework_name: str,
        capture_input: bool = False,
        capture_output: bool = False,
        max_payload_size: int = 10000,
    ):
        """Initialize the patcher.

        Args:
            framework_name: Name of the framework being patched (e.g., "sklearn").
            capture_input: Whether to capture input data in spans.
            capture_output: Whether to capture output data in spans.
            max_payload_size: Maximum size in bytes for captured data (prevents memory bloat).
        """
        self.framework_name = framework_name
        self.capture_input = capture_input
        self.capture_output = capture_output
        self.max_payload_size = max_payload_size

    def patch_method(
        self,
        target_class: type,
        method_name: str,
        span_name_template: str,
        additional_attrs: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Patch a single method on a target class using wrapt.

        Args:
            target_class: The class to patch (e.g., sklearn.base.BaseEstimator).
            method_name: Name of the method to patch (e.g., "predict").
            span_name_template: Template for span name (e.g., "{framework}.{class_name}.{method}").
            additional_attrs: Additional attributes to include in telemetry.

        Raises:
            AttributeError: If method does not exist on target class.
        """
        # Check if already patched (prevent double-patching)
        patch_key = (target_class, method_name)
        if patch_key in _patched_methods:
            logger.debug(f"{target_class.__name__}.{method_name} already patched, skipping")
            return

        # Get original method
        if not hasattr(target_class, method_name):
            raise AttributeError(f"{target_class.__name__} has no method '{method_name}'")

        original_method = getattr(target_class, method_name)

        # Store original method for potential unpatching
        _patched_methods[patch_key] = original_method

        # Capture patcher instance and additional_attrs in closure
        patcher = self
        attrs = additional_attrs or {}

        # Create wrapper using wrapt.FunctionWrapper for robust signature preservation
        # wrapt.FunctionWrapper ensures proper introspection, descriptor handling,
        # and compatibility with other instrumentation tools (MLflow, wandb)
        # Note: wrapt passes args as tuple and kwargs as dict (already unpacked)
        def wrapper(wrapped, instance, args, kwargs):
            return patcher._instrumented_call(
                wrapped,
                instance,
                span_name_template,
                method_name,
                attrs,
                args,  # Pass as tuple (not unpacked)
                kwargs,  # Pass as dict (not unpacked)
            )

        # Apply wrapt wrapping - preserves signature and cooperates with other wrappers
        wrapped_method = wrapt.FunctionWrapper(original_method, wrapper)

        # Replace method on target class
        setattr(target_class, method_name, wrapped_method)

        logger.debug(f"Patched {target_class.__name__}.{method_name}")

    def _instrumented_call(
        self,
        original_method: Callable,
        instance: Any,
        span_name_template: str,
        method_name: str,
        additional_attrs: Dict[str, Any],
        args: tuple,
        kwargs: dict,
    ) -> Any:
        """Execute an instrumented method call with proper OTel semantics.

        Records timing, attributes, and exceptions on the active span context.
        Avoids creating disconnected spans for error handling.

        Args:
            original_method: The original unpatched method (wrapped by wrapt).
            instance: The instance on which the method is being called (can be None).
            span_name_template: Template for the span name.
            method_name: Name of the method being called.
            additional_attrs: Additional attributes for telemetry.
            args: Positional arguments tuple to the method.
            kwargs: Keyword arguments dict to the method.

        Returns:
            The result of the original method call.
        """
        global _missing_client_warned

        client = MCAClient.get_instance()

        # Graceful fallback if client not available (log warning once to avoid spam)
        if client is None:
            if not _missing_client_warned:
                logger.warning(
                    f"MCAClient not initialized when {method_name} was called. "
                    "Call MCAClient(...) before using instrumented models. "
                    "Predictions will execute without telemetry."
                )
                _missing_client_warned = True
            return original_method(*args, **kwargs)

        # Build span name - handle None instance for standalone functions
        if instance is not None:
            class_name = instance.__class__.__name__
        else:
            class_name = "function"  # Fallback for non-instance methods

        span_name = span_name_template.format(
            framework=self.framework_name, class_name=class_name, method=method_name
        )

        # Build telemetry attributes
        attrs = {
            "framework": self.framework_name,
            "model_class": class_name,
            "method": method_name,
        }
        attrs.update(additional_attrs)

        # Track timing manually for accurate latency metric
        start_time = time.perf_counter()

        # Story 8.5: Delegate to record_prediction() for unified schema
        # Extract raw input/output data and let record_prediction() handle serialization
        # Issue #8 Fix: Use same logic as decorator - single arg direct, else wrap
        input_for_capture = None
        if self.capture_input and (args or kwargs):
            if len(args) == 1 and not kwargs:
                # Single positional arg - pass directly
                input_for_capture = args[0]
            else:
                # Multi-args, kwargs-only, or mixed - wrap everything to preserve all data
                input_for_capture = {"args": args, "kwargs": kwargs}

        output_for_capture = None  # Set after method call

        # Execute with instrumentation
        try:
            with client.trace(span_name, **attrs):
                # Call the wrapped method (already bound to instance by wrapt)
                result = original_method(*args, **kwargs)

                # Calculate actual latency
                latency_seconds = time.perf_counter() - start_time

                # Capture output if requested
                if self.capture_output:
                    output_for_capture = result

                # Story 8.5: Delegate to record_prediction() with raw data
                # This ensures consistent schema (data, shape, truncated) across all instrumentation methods
                try:
                    client.record_prediction(
                        input_data=input_for_capture,
                        output=output_for_capture,
                        latency=latency_seconds,
                        **attrs,
                    )
                except Exception as metric_error:
                    logger.warning(f"Failed to record prediction metric: {metric_error}")

                return result

        except Exception as e:
            # Calculate latency even for failed predictions
            latency_seconds = time.perf_counter() - start_time

            # Add error info to attributes FIRST (before any span operations)
            attrs["error"] = str(e)[:500]  # Truncate long error messages
            attrs["error_type"] = type(e).__name__
            attrs["failed"] = True

            # Record exception on the ACTIVE span (not a new disconnected span)
            try:
                # Get the current active span from context
                from opentelemetry import trace
                from opentelemetry.trace import Status, StatusCode

                current_span = trace.get_current_span()
                if current_span and current_span.is_recording():
                    # Record exception and error status on active span
                    current_span.record_exception(e)
                    current_span.set_status(Status(StatusCode.ERROR, str(e)[:500]))

                    # Set error attributes on span
                    current_span.set_attribute("error", True)
                    current_span.set_attribute("error.type", type(e).__name__)
                    current_span.set_attribute("error.message", str(e)[:500])
            except Exception as span_error:
                # Log but don't crash if error recording fails
                logger.warning(f"Failed to record exception in span: {span_error}")

            # Story 8.5: Delegate to record_prediction() even for failed predictions
            # Record failed prediction metric with actual latency
            try:
                client.record_prediction(
                    input_data=input_for_capture,
                    output=None,  # No output on error
                    latency=latency_seconds,
                    **attrs,
                )
            except Exception as metric_error:
                logger.warning(f"Failed to record failed prediction metric: {metric_error}")

            # Re-raise the original exception (preserve stack trace)
            raise

    def _capture_data_safely(self, data: Any, data_type: str) -> Optional[Dict[str, Any]]:
        """Safely capture data with size limits to prevent memory bloat.

        CRITICAL: Checks object size BEFORE converting to string to avoid OOM.
        Uses consistent byte-based size measurement and truncation.

        Args:
            data: Input or output data to capture.
            data_type: Type label ("input" or "output") for attribute naming.

        Returns:
            Dictionary with captured data attributes, or None if capture fails/exceeds limit.
        """
        try:
            result: Dict[str, Any] = {}

            # Always capture shape (lightweight)
            shape_str = self._get_shape(data)
            if shape_str:
                result[f"{data_type}_shape"] = shape_str

            # Check object size BEFORE converting to string to prevent OOM
            # For large numpy arrays or DataFrames, sys.getsizeof() gives reasonable estimate
            object_size = sys.getsizeof(data)

            # If object itself is huge, skip string conversion entirely
            if object_size > self.max_payload_size * 10:
                result[f"{data_type}_sample"] = (
                    f"[Object too large to serialize: {object_size} bytes]"
                )
                result[f"{data_type}_truncated"] = True
                logger.debug(
                    f"Skipped {data_type} data serialization: object size {object_size} bytes "
                    f"exceeds safety threshold of {self.max_payload_size * 10} bytes"
                )
                return result

            # Convert to string (now safe - object size was checked)
            data_str = str(data)

            # Measure actual string byte size for consistent truncation
            # Use UTF-8 encoding length, not character count
            data_bytes = data_str.encode("utf-8", errors="replace")
            data_size_bytes = len(data_bytes)

            if data_size_bytes > self.max_payload_size:
                # Truncate by BYTES, not character count, for consistency
                # Decode truncated bytes back to string (handle partial UTF-8 chars)
                truncated_bytes = data_bytes[: self.max_payload_size]
                truncated_str = truncated_bytes.decode("utf-8", errors="ignore")

                result[f"{data_type}_sample"] = (
                    f"{truncated_str}...[truncated, {data_size_bytes} bytes total]"
                )
                result[f"{data_type}_truncated"] = True
                logger.debug(
                    f"Truncated {data_type} data: {data_size_bytes} bytes exceeds limit "
                    f"of {self.max_payload_size} bytes"
                )
            else:
                result[f"{data_type}_sample"] = data_str
                result[f"{data_type}_truncated"] = False

            return result

        except Exception as e:
            logger.debug(f"Failed to capture {data_type} data: {e}")
            return None

    def _get_shape(self, data: Any) -> Optional[str]:
        """Extract shape information from data.

        Args:
            data: Input or output data (numpy array, pandas DataFrame, etc.).

        Returns:
            String representation of shape, or None if shape cannot be determined.
        """
        try:
            # Handle numpy arrays, pandas DataFrames, scipy sparse matrices
            if hasattr(data, "shape"):
                return str(data.shape)
            # Handle lists
            elif isinstance(data, (list, tuple)):
                return f"({len(data)},)"
            else:
                return None
        except Exception:
            return None

    @staticmethod
    def unpatch_method(target_class: type, method_name: str) -> None:
        """Unpatch a previously patched method.

        This is primarily used for testing to restore original behavior.

        Args:
            target_class: The class that was patched.
            method_name: Name of the method to unpatch.
        """
        patch_key = (target_class, method_name)
        if patch_key in _patched_methods:
            original_method = _patched_methods.pop(patch_key)
            setattr(target_class, method_name, original_method)
            logger.debug(f"Unpatched {target_class.__name__}.{method_name}")
        else:
            logger.warning(f"{target_class.__name__}.{method_name} was not patched")


def unpatch_all() -> None:
    """Unpatch all previously patched methods.

    This is primarily used for testing cleanup.
    """
    for (target_class, method_name), original_method in list(_patched_methods.items()):
        setattr(target_class, method_name, original_method)
        logger.debug(f"Unpatched {target_class.__name__}.{method_name}")

    _patched_methods.clear()
    logger.info("All patched methods restored")
