"""
Optimization metrics for DSPy evaluation.

Provides evaluation metrics for agent task success, quality assessment,
and composite scoring for MIPROv2 optimization.

Requirements: SEC-06
"""

from typing import Optional, Any, Callable, Dict

try:
    import dspy

    HAS_DSPY = True
except ImportError:
    HAS_DSPY = False
    dspy = None


def task_success_metric(example: Any, pred: Any, trace: Any = None) -> float:
    """
    Metric for evaluating agent task success.

    This is the primary metric for MIPROv2 optimization. It evaluates
    whether an agent's prediction represents successful task completion.

    Args:
        example: The ground truth example (contains expected values).
        pred: The prediction made by the agent.
        trace: Optional execution trace for detailed analysis.

    Returns:
        Float score between 0.0 and 1.0.
    """
    if not HAS_DSPY:
        return 0.0

    # Check explicit success flag on prediction
    if hasattr(pred, "success"):
        return 1.0 if pred.success else 0.0

    # Check if prediction has an answer that matches example
    if hasattr(example, "answer") and hasattr(pred, "answer"):
        return 1.0 if pred.answer == example.answer else 0.0

    # Check task_output field (from ExecutionTrace pattern)
    if hasattr(pred, "task_output") and hasattr(example, "task_output"):
        pred_output = pred.task_output
        example_output = example.task_output
        if isinstance(pred_output, dict) and isinstance(example_output, dict):
            # Check for success field in output
            if "success" in pred_output:
                return 1.0 if pred_output.get("success") else 0.0
            # Check for matching result
            if "result" in pred_output and "result" in example_output:
                return 1.0 if pred_output["result"] == example_output["result"] else 0.0

    # Default: non-null prediction is considered successful
    return 1.0 if pred is not None else 0.0


def quality_metric(example: Any, pred: Any, trace: Any = None) -> float:
    """
    Metric for evaluating output quality beyond simple success.

    Evaluates qualitative aspects like completeness, relevance, and format.

    Args:
        example: The ground truth example.
        pred: The prediction made by the agent.
        trace: Optional execution trace for detailed analysis.

    Returns:
        Float score between 0.0 and 1.0.
    """
    if not HAS_DSPY:
        return 0.0

    score = 0.0

    # Check completeness (has required fields)
    if hasattr(pred, "task_output"):
        output = pred.task_output
        if isinstance(output, dict):
            required_fields = ["result", "status"]
            present = sum(1 for f in required_fields if f in output)
            score += (present / len(required_fields)) * 0.4

    # Check for context relevance
    if hasattr(pred, "context") and pred.context:
        score += 0.3

    # Check for proper format (JSON-like structure)
    if hasattr(pred, "task_output"):
        try:
            if isinstance(pred.task_output, dict):
                score += 0.3
        except (TypeError, AttributeError):
            pass

    return min(1.0, score)


def latency_metric(
    example: Any,
    pred: Any,
    trace: Any = None,
    target_ms: int = 5000,
    max_ms: int = 30000,
) -> float:
    """
    Metric for evaluating response latency.

    Lower latency scores higher. Uses a linear scale from target to max.

    Args:
        example: The ground truth example.
        pred: The prediction made by the agent.
        trace: Optional execution trace with timing info.
        target_ms: Target latency in milliseconds (full score).
        max_ms: Maximum acceptable latency in milliseconds (zero score).

    Returns:
        Float score between 0.0 and 1.0.
    """
    # Extract duration from trace if available
    duration_ms = None
    if trace is not None and hasattr(trace, "duration_ms"):
        val = trace.duration_ms
        if isinstance(val, (int, float)):
            duration_ms = val
    if duration_ms is None and hasattr(pred, "duration_ms"):
        val = pred.duration_ms
        if isinstance(val, (int, float)):
            duration_ms = val
    if duration_ms is None and hasattr(example, "duration_ms"):
        val = example.duration_ms
        if isinstance(val, (int, float)):
            duration_ms = val

    if duration_ms is None:
        # No timing info available, return neutral score
        return 0.5

    if duration_ms <= target_ms:
        return 1.0
    if duration_ms >= max_ms:
        return 0.0

    # Linear interpolation between target and max
    ratio = (max_ms - duration_ms) / (max_ms - target_ms)
    return max(0.0, min(1.0, ratio))


def composite_metric(
    example: Any,
    pred: Any,
    trace: Any = None,
    weights: Optional[Dict[str, float]] = None,
) -> float:
    """
    Composite metric combining success, quality, and latency.

    Provides a balanced evaluation for optimization that considers
    multiple aspects of agent performance.

    Args:
        example: The ground truth example.
        pred: The prediction made by the agent.
        trace: Optional execution trace for detailed analysis.
        weights: Custom weights for each component.
                 Default: {'success': 0.5, 'quality': 0.3, 'latency': 0.2}

    Returns:
        Float score between 0.0 and 1.0 (weighted average).
    """
    if weights is None:
        weights = {"success": 0.5, "quality": 0.3, "latency": 0.2}

    # Calculate individual scores
    success_score = task_success_metric(example, pred, trace)
    quality_score = quality_metric(example, pred, trace)
    latency_score = latency_metric(example, pred, trace)

    # Apply weights
    total = 0.0
    weight_sum = 0.0

    if "success" in weights:
        total += success_score * weights["success"]
        weight_sum += weights["success"]

    if "quality" in weights:
        total += quality_score * weights["quality"]
        weight_sum += weights["quality"]

    if "latency" in weights:
        total += latency_score * weights["latency"]
        weight_sum += weights["latency"]

    # Normalize by actual weight sum
    if weight_sum > 0:
        return total / weight_sum

    return 0.0


class MetricRegistry:
    """
    Registry for custom optimization metrics.

    Allows registration and retrieval of metric functions by name.
    Supports both built-in metrics and user-defined custom metrics.
    """

    _metrics: Dict[str, Callable] = {}

    @classmethod
    def register(cls, name: str, metric: Callable) -> None:
        """
        Register a metric function under a given name.

        Args:
            name: Unique identifier for the metric.
            metric: Callable that takes (example, pred, trace) and returns float.

        Raises:
            ValueError: If name is empty or metric is not callable.
        """
        if not name:
            raise ValueError("Metric name cannot be empty")
        if not callable(metric):
            raise ValueError("Metric must be callable")

        cls._metrics[name] = metric

    @classmethod
    def get(cls, name: str) -> Optional[Callable]:
        """
        Retrieve a registered metric by name.

        Args:
            name: The metric identifier.

        Returns:
            The metric function if found, None otherwise.
        """
        # Check custom metrics first
        if name in cls._metrics:
            return cls._metrics[name]

        # Check built-in metrics
        built_in = {
            "task_success": task_success_metric,
            "quality": quality_metric,
            "latency": latency_metric,
            "composite": composite_metric,
        }

        return built_in.get(name)

    @classmethod
    def list_metrics(cls) -> list:
        """
        List all available metric names.

        Returns:
            List of registered metric names (both custom and built-in).
        """
        built_in = ["task_success", "quality", "latency", "composite"]
        custom = list(cls._metrics.keys())
        return sorted(set(built_in + custom))

    @classmethod
    def unregister(cls, name: str) -> bool:
        """
        Remove a custom metric from the registry.

        Args:
            name: The metric identifier to remove.

        Returns:
            True if metric was removed, False if not found.
        """
        if name in cls._metrics:
            del cls._metrics[name]
            return True
        return False

    @classmethod
    def clear(cls) -> None:
        """Clear all custom metrics from the registry."""
        cls._metrics.clear()


# Auto-register built-in metrics
MetricRegistry.register("task_success", task_success_metric)
MetricRegistry.register("quality", quality_metric)
MetricRegistry.register("latency", latency_metric)
MetricRegistry.register("composite", composite_metric)
