"""Instrumentation decorators for MCA SDK.

This module provides decorators for automatic model instrumentation,
making it easy to add telemetry to prediction functions.
"""

import functools
import time
from typing import Callable, Optional, Any

from ..core.client import MCAClient


def instrument_model(
    client: MCAClient,
    metric_name: Optional[str] = None,
    capture_input: bool = False,
    capture_output: bool = False,
):
    """Decorator for automatic model instrumentation.

    Wraps a prediction function to automatically record metrics including
    latency, call counts, and optionally input/output data.

    Args:
        client: MCAClient instance for recording telemetry
        metric_name: Custom metric name prefix (default: function name)
        capture_input: Whether to capture input data (default: False for security)
        capture_output: Whether to capture output data (default: False for security)

    Returns:
        Decorated function with automatic instrumentation

    Examples:
        Basic usage (no data capture):
        >>> client = MCAClient(...)
        >>> @instrument_model(client)
        ... def predict(input_data):
        ...     return model.predict(input_data)

        With custom metric name:
        >>> @instrument_model(client, metric_name="readmission_prediction")
        ... def predict_readmission(patient_data):
        ...     return model.predict(patient_data)

        Capturing input/output (HIPAA: ensure PHI sanitized):
        >>> @instrument_model(client, capture_input=True, capture_output=True)
        ... def predict(input_data):
        ...     return model.predict(input_data)
    """

    def decorator(func: Callable) -> Callable:
        # Use function name if metric_name not provided
        name = metric_name or func.__name__

        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            import uuid

            start_time = time.time()

            # FIXED (Round 11 Issue #3): Generate prediction_id for correlation
            prediction_id = str(uuid.uuid4())

            # Create trace span
            # FIXED (Round 12 Issue #2): Use dynamic span name for differentiation
            with client.trace(f"{name}.predict") as span:
                # FIXED (Round 11 Issue #3): Add prediction_id to span
                span.set_attribute("prediction.id", prediction_id)

                try:
                    # Call the original function
                    result = func(*args, **kwargs)

                    # FIXED (Round 14 Issue #5): Detect generators/streaming responses
                    import inspect

                    if inspect.isgenerator(result) or inspect.isasyncgen(result):
                        # Wrap generator to measure full latency
                        def generator_wrapper():
                            try:
                                for item in result:
                                    yield item
                            finally:
                                # Measure latency when generator exhausted
                                latency = time.time() - start_time
                                span.set_attribute("status.code", "OK")
                                span.set_attribute("prediction.latency_ms", latency * 1000)
                                client.record_prediction(
                                    input_data=None,  # Don't capture for generators
                                    output=None,
                                    latency=latency,
                                    function=name,
                                    prediction_id=prediction_id,
                                )

                        return generator_wrapper()

                    # Calculate latency for non-generator results
                    latency = time.time() - start_time

                    # FIXED (Round 14 Issue #8): Safer kwarg capture - don't blindly capture sensitive data
                    input_data = None
                    output_data = None

                    if capture_input:
                        if args and kwargs:
                            # Serialize both safely
                            from ..utils.serialization import serialize_for_telemetry

                            input_data = {
                                "args": serialize_for_telemetry(args),
                                "kwargs": serialize_for_telemetry(kwargs),
                            }
                        elif args:
                            input_data = args[0] if len(args) == 1 else args
                        elif kwargs:
                            # FIXED (Round 14 Issue #8): Serialize kwargs safely instead of blind capture
                            from ..utils.serialization import serialize_for_telemetry

                            input_data = serialize_for_telemetry(kwargs)

                    if capture_output:
                        output_data = result

                    # FIXED (Round 11 Issue #5): Use OpenTelemetry semantic conventions
                    span.set_attribute("status.code", "OK")
                    span.set_attribute("prediction.latency_ms", latency * 1000)

                    # FIXED (Round 11 Issue #4): No 'or {}' mutation
                    # FIXED (Round 11 Issue #3): Pass prediction_id for correlation
                    client.record_prediction(
                        input_data=input_data,
                        output=output_data,
                        latency=latency,
                        function=name,
                        prediction_id=prediction_id,
                    )

                    return result

                except Exception as e:
                    # Calculate latency even on error
                    latency = time.time() - start_time

                    # FIXED (Round 11 Issue #5 & #6): Standard attributes + latency
                    span.set_attribute("status.code", "ERROR")
                    span.set_attribute("error.type", type(e).__name__)
                    span.set_attribute("error.message", str(e))
                    span.set_attribute("prediction.latency_ms", latency * 1000)

                    # FIXED (Round 11 Issue #3): Pass prediction_id
                    client.record_error(
                        e, latency=latency, function=name, prediction_id=prediction_id
                    )

                    # Re-raise the exception
                    raise

        return wrapper

    return decorator


def instrument_async_model(
    client: MCAClient,
    metric_name: Optional[str] = None,
    capture_input: bool = False,
    capture_output: bool = False,
):
    """Decorator for automatic async model instrumentation.

    Wraps an async prediction function to automatically record metrics.

    Args:
        client: MCAClient instance for recording telemetry
        metric_name: Custom metric name prefix (default: function name)
        capture_input: Whether to capture input data (default: False for security)
        capture_output: Whether to capture output data (default: False for security)

    Returns:
        Decorated async function with automatic instrumentation

    Example:
        >>> @instrument_async_model(client)
        ... async def predict_async(input_data):
        ...     result = await model.predict_async(input_data)
        ...     return result
    """

    def decorator(func: Callable) -> Callable:
        # Use function name if metric_name not provided
        name = metric_name or func.__name__

        @functools.wraps(func)
        async def wrapper(*args, **kwargs) -> Any:
            import uuid

            start_time = time.time()

            # FIXED (Round 11 Issue #3): Generate prediction_id for correlation
            prediction_id = str(uuid.uuid4())

            # Create trace span
            # FIXED (Round 12 Issue #2): Use dynamic span name for differentiation
            with client.trace(f"{name}.predict") as span:
                # FIXED (Round 11 Issue #3): Add prediction_id to span
                span.set_attribute("prediction.id", prediction_id)

                try:
                    # Call the original async function
                    result = await func(*args, **kwargs)

                    # FIXED (Round 14 Issue #5): Detect async generators/streaming responses
                    import inspect

                    if inspect.isasyncgen(result):
                        # Wrap async generator to measure full latency
                        async def async_generator_wrapper():
                            try:
                                async for item in result:
                                    yield item
                            finally:
                                # Measure latency when generator exhausted
                                latency = time.time() - start_time
                                span.set_attribute("status.code", "OK")
                                span.set_attribute("prediction.latency_ms", latency * 1000)
                                client.record_prediction(
                                    input_data=None,  # Don't capture for generators
                                    output=None,
                                    latency=latency,
                                    function=name,
                                    prediction_id=prediction_id,
                                )

                        return async_generator_wrapper()

                    # Calculate latency for non-generator results
                    latency = time.time() - start_time

                    # FIXED (Round 14 Issue #8): Safer kwarg capture - don't blindly capture sensitive data
                    input_data = None
                    output_data = None

                    if capture_input:
                        if args and kwargs:
                            # Serialize both safely
                            from ..utils.serialization import serialize_for_telemetry

                            input_data = {
                                "args": serialize_for_telemetry(args),
                                "kwargs": serialize_for_telemetry(kwargs),
                            }
                        elif args:
                            input_data = args[0] if len(args) == 1 else args
                        elif kwargs:
                            # FIXED (Round 14 Issue #8): Serialize kwargs safely instead of blind capture
                            from ..utils.serialization import serialize_for_telemetry

                            input_data = serialize_for_telemetry(kwargs)

                    if capture_output:
                        output_data = result

                    # FIXED (Round 11 Issue #5): Use OpenTelemetry semantic conventions
                    span.set_attribute("status.code", "OK")
                    span.set_attribute("prediction.latency_ms", latency * 1000)

                    # FIXED (Round 11 Issue #4): No 'or {}' mutation
                    # FIXED (Round 11 Issue #3): Pass prediction_id for correlation
                    client.record_prediction(
                        input_data=input_data,
                        output=output_data,
                        latency=latency,
                        function=name,
                        prediction_id=prediction_id,
                    )

                    return result

                except Exception as e:
                    # Calculate latency even on error
                    latency = time.time() - start_time

                    # FIXED (Round 11 Issue #5 & #6): Standard attributes + latency
                    span.set_attribute("status.code", "ERROR")
                    span.set_attribute("error.type", type(e).__name__)
                    span.set_attribute("error.message", str(e))
                    span.set_attribute("prediction.latency_ms", latency * 1000)

                    # FIXED (Round 11 Issue #3): Pass prediction_id
                    client.record_error(
                        e, latency=latency, function=name, prediction_id=prediction_id
                    )

                    # Re-raise the exception
                    raise

        return wrapper

    return decorator


def track_batch_prediction(client: MCAClient, batch_size_metric: Optional[str] = None):
    """Decorator for tracking batch prediction metrics.

    Automatically tracks batch size and per-item latency for batch
    prediction functions.

    Args:
        client: MCAClient instance
        batch_size_metric: DEPRECATED - No longer used, kept for backward compatibility

    Returns:
        Decorated function with batch tracking

    Example:
        >>> @track_batch_prediction(client)
        ... def predict_batch(inputs):
        ...     return [model.predict(x) for x in inputs]
    """
    # CRITICAL FIX (Issue #7): Make parameter optional with deprecation warning
    if batch_size_metric is not None:
        import warnings

        warnings.warn(
            "batch_size_metric parameter is deprecated and has no effect. "
            "Metric names are now determined by client configuration.",
            DeprecationWarning,
            stacklevel=2,
        )

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            # Start timing
            start_time = time.time()

            # Get batch from first argument
            batch = args[0] if args else []
            batch_size = len(batch) if hasattr(batch, "__len__") else 0

            try:
                # Call the original function
                result = func(*args, **kwargs)

                # Calculate metrics
                total_latency = time.time() - start_time
                per_item_latency = total_latency / batch_size if batch_size > 0 else 0

                # FIXED (Round 10 Issue #3): Use dynamic prefix based on client._metric_prefix
                prefix = f"emms_{client._metric_prefix}"

                # Record batch size
                client.record_metric(f"{prefix}_batch_size_count", batch_size)

                # FIXED (Round 14 Issue #4): Remove batch_size from labels to prevent cardinality explosion
                # Record total latency (no labels - cardinality safe)
                client.record_metric(f"{prefix}_batch_latency_seconds", total_latency)

                # Record per-item latency (no labels - cardinality safe)
                client.record_metric(f"{prefix}_per_item_latency_seconds", per_item_latency)

                return result

            except Exception as e:
                # FIXED (Issue #2): Use centralized record_error() instead of bypassing
                client.record_error(e, function="batch_predict")
                raise

        return wrapper

    return decorator
