"""Instrumentation decorators for MCA SDK.

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

import contextvars
import functools
import time
import uuid
import warnings
from typing import Callable, Optional, Any

from ..core.client import MCAClient

# FIX Issue #2: Thread-safe goal_id storage using ContextVar
_current_goal_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
    "current_goal_id", default=None
)


def _get_metric_prefix(client) -> str:
    """Get metric prefix from client, defaulting to 'model' for non-SDK objects."""
    prefix = getattr(client, "_metric_prefix", "model")
    return prefix if isinstance(prefix, str) else "model"


def instrument_model(
    client: MCAClient,
    metric_name: Optional[str] = None,
    capture_input: bool = True,
    capture_output: bool = True,
):
    """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("ml.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

                    # Set span attributes using function-name prefix convention
                    span.set_attribute(f"{name}.latency_seconds", latency)
                    span.set_attribute(f"{name}.status", "success")

                    # Pass input/output (default to {} when not capturing)
                    client.record_prediction(
                        input_data=input_data if input_data is not None else {},
                        output=output_data if output_data is not None else {},
                        latency=latency,
                        function=name,
                        prediction_id=prediction_id,
                    )

                    return result

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

                    # Set error span attributes using function-name prefix
                    span.set_attribute(f"{name}.status", "error")
                    span.set_attribute(f"{name}.error", str(e))

                    # Use centralized record_error for consistent dual recording
                    # (increments predictions_total{status="error"} + errors_total)
                    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 = True,
    capture_output: bool = True,
):
    """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("ml.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

                    # Set span attributes using function-name prefix convention
                    span.set_attribute(f"{name}.latency_seconds", latency)
                    span.set_attribute(f"{name}.status", "success")

                    # Pass input/output (default to {} when not capturing)
                    client.record_prediction(
                        input_data=input_data if input_data is not None else {},
                        output=output_data if output_data is not None else {},
                        latency=latency,
                        function=name,
                        prediction_id=prediction_id,
                    )

                    return result

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

                    # Set error span attributes using function-name prefix
                    span.set_attribute(f"{name}.status", "error")
                    span.set_attribute(f"{name}.error", str(e))

                    # Use centralized record_error for consistent dual recording
                    client.record_error(
                        e, latency=latency, function=name, prediction_id=prediction_id
                    )

                    # Re-raise the exception
                    raise

        return wrapper

    return decorator


def trace_agent_goal(
    client: MCAClient,
    description: str,
    goal_type: str = "general",
):
    """Decorator for automatic agent goal tracking (Story 1.19).

    Wraps an agent execution function to automatically track goal lifecycle,
    including start/end times and completion duration. The goal_id is stored
    in context and can be retrieved via client.get_current_goal_id() if needed.

    FIX Issue #2: Does NOT inject goal_id as parameter to preserve function signature.

    Args:
        client: MCAClient instance for recording telemetry
        description: Human-readable goal description
        goal_type: Category of goal (e.g., "research", "analysis", "planning")

    Returns:
        Decorated function with automatic goal tracking

    Example:
        >>> client = MCAClient(service_name="research-agent", model_type="agentic")
        >>>
        >>> @trace_agent_goal(client, description="Find diabetes treatments", goal_type="research")
        ... def execute_research():
        ...     # Retrieve goal_id from context if needed
        ...     goal_id = client.get_current_goal_id()
        ...     with client.track_tool("pubmed_search", goal_id=goal_id):
        ...         results = search_pubmed("diabetes")
        ...     return results
        >>>
        >>> result = execute_research()  # Signature unchanged
        >>> client.shutdown()
    """

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            # Start goal tracking
            goal_id = client.record_goal_started(goal_description=description, goal_type=goal_type)

            # FIX Issue #2: Thread-safe storage using ContextVar
            token = _current_goal_id.set(goal_id)

            try:
                # Call function with original signature (no injection)
                result = func(*args, **kwargs)

                # Mark goal as successful
                client.record_goal_completed(goal_id, status="success")

                return result

            except Exception:
                # Mark goal as failed
                client.record_goal_completed(goal_id, status="failure")
                raise
            finally:
                # Restore previous context value
                _current_goal_id.reset(token)

        return wrapper

    return decorator


def trace_agent_step(
    client: MCAClient,
    step_id: Optional[str] = None,
):
    """Decorator for automatic agent step tracking (Story 1.19).

    Wraps individual reasoning/action loops to track step-level execution.
    Creates a child span under the active goal span with start/end timestamps.

    FIX Issue #10: Explicit error handling to track step failures consistently
    with goal and tool tracking patterns.

    Args:
        client: MCAClient instance for recording telemetry
        step_id: Optional step identifier (defaults to auto-generated UUID)

    Returns:
        Decorated function with automatic step tracking

    Example:
        >>> @trace_agent_step(client, step_id="search-pubmed")
        ... def search_step():
        ...     goal_id = client.get_current_goal_id()
        ...     with client.track_tool("pubmed_search", goal_id=goal_id):
        ...         return search_pubmed("diabetes")
        >>>
        >>> result = search_step()
    """

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            # Generate step_id if not provided
            current_step_id = step_id or str(uuid.uuid4())

            # Create step span as child of current context (goal span)
            with client._tracer.start_as_current_span(
                "agent.step",
                attributes={
                    "genai.agent.step.id": current_step_id,
                },
            ) as span:
                try:
                    # Execute step
                    result = func(*args, **kwargs)

                    # FIX Issue #10: Explicitly mark step success
                    span.set_attribute("step.status", "success")

                    return result

                except Exception as e:
                    # FIX Issue #10: Explicitly mark step failure
                    from opentelemetry.trace import Status, StatusCode

                    span.set_status(Status(StatusCode.ERROR, str(e)[:200]))
                    span.set_attribute("step.status", "error")
                    span.set_attribute("error.type", type(e).__name__)
                    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: Custom metric name for batch size (deprecated; default: "{prefix}.batch_size_count")

    Returns:
        Decorated function with batch tracking

    Example:
        >>> @track_batch_prediction(client)
        ... def predict_batch(inputs):
        ...     return [model.predict(x) for x in inputs]
    """
    if batch_size_metric is not None:
        warnings.warn(
            "batch_size_metric parameter is deprecated. Use the default batch_size metric instead.",
            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

            # Resolve metric prefix (safe fallback for mock clients)
            prefix = _get_metric_prefix(client)

            # Batch size metric name (custom or default)
            size_metric = batch_size_metric if batch_size_metric else f"{prefix}.batch_size_count"

            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

                # Record batch size
                client.record_metric(size_metric, batch_size)

                # Record total latency with batch size attribute
                client.record_metric(
                    f"{prefix}.batch_latency_seconds",
                    total_latency,
                    batch_size=str(batch_size),
                )

                # Record per-item latency with batch size attribute
                client.record_metric(
                    f"{prefix}.per_item_latency_seconds",
                    per_item_latency,
                    batch_size=str(batch_size),
                )

                return result

            except Exception as e:
                import uuid

                batch_id = str(uuid.uuid4())
                # Use centralized record_error (no batch_size label in error attributes)
                client.record_error(e, prediction_id=batch_id)
                raise

        return wrapper

    return decorator
