"""Metrics collection using AWS Lambda Powertools.

This module provides CloudWatch metrics integration for monitoring
AI agent performance, token usage, and error rates.

Usage:
    from agent_factory.observability import get_metrics, record_ai_tokens

    metrics = get_metrics()
    metrics.add_metric(name="Requests", unit=MetricUnit.Count, value=1)

    record_ai_tokens(metrics, input_tokens=100, output_tokens=50, model="claude-3")

For non-Lambda environments, provides a no-op fallback.
"""

import os
from functools import lru_cache
from typing import Any, Optional

# Try to import Powertools Metrics
try:
    from aws_lambda_powertools import Metrics as PowertoolsMetrics
    from aws_lambda_powertools.metrics import MetricUnit

    HAS_POWERTOOLS = True
except ImportError:
    HAS_POWERTOOLS = False
    PowertoolsMetrics = None
    MetricUnit = None


class NoOpMetrics:
    """No-op metrics class for non-Lambda environments."""

    def add_metric(self, name: str, unit: Any = None, value: float = 1) -> None:
        """No-op metric recording."""
        pass

    def add_dimension(self, name: str, value: str) -> None:
        """No-op dimension adding."""
        pass

    def add_metadata(self, key: str, value: Any) -> None:
        """No-op metadata adding."""
        pass

    def flush_metrics(self) -> None:
        """No-op flush."""
        pass


@lru_cache(maxsize=8)
def get_metrics(
    namespace: Optional[str] = None,
    service: Optional[str] = None,
) -> Any:
    """Get a configured metrics instance for CloudWatch.

    Uses AWS Lambda Powertools Metrics when available, which automatically
    publishes metrics to CloudWatch. Falls back to a no-op implementation
    in non-Lambda environments.

    Args:
        namespace: CloudWatch namespace (defaults to env or "AgentFactory")
        service: Service name (defaults to env or "agent")

    Returns:
        Metrics instance (Powertools or NoOp)

    Example:
        metrics = get_metrics()
        metrics.add_metric(name="Invocations", unit=MetricUnit.Count, value=1)
    """
    if HAS_POWERTOOLS:
        return PowertoolsMetrics(
            namespace=namespace or os.getenv("POWERTOOLS_METRICS_NAMESPACE", "AgentFactory"),
            service=service or os.getenv("POWERTOOLS_SERVICE_NAME", "agent"),
        )
    else:
        return NoOpMetrics()


def record_invocation(
    metrics: Any,
    handler_name: str,
    *,
    success: bool = True,
) -> None:
    """Record a handler invocation metric.

    Args:
        metrics: Metrics instance
        handler_name: Name of the handler
        success: Whether the invocation was successful
    """
    metrics.add_dimension(name="Handler", value=handler_name)

    if HAS_POWERTOOLS:
        metrics.add_metric(
            name="Invocations",
            unit=MetricUnit.Count,
            value=1,
        )
        if success:
            metrics.add_metric(
                name="SuccessfulInvocations",
                unit=MetricUnit.Count,
                value=1,
            )
        else:
            metrics.add_metric(
                name="FailedInvocations",
                unit=MetricUnit.Count,
                value=1,
            )
    else:
        metrics.add_metric("Invocations", value=1)


def record_latency(
    metrics: Any,
    handler_name: str,
    duration_ms: float,
) -> None:
    """Record handler latency metric.

    Args:
        metrics: Metrics instance
        handler_name: Name of the handler
        duration_ms: Handler duration in milliseconds
    """
    metrics.add_dimension(name="Handler", value=handler_name)

    if HAS_POWERTOOLS:
        metrics.add_metric(
            name="Latency",
            unit=MetricUnit.Milliseconds,
            value=duration_ms,
        )
    else:
        metrics.add_metric("Latency", value=duration_ms)


def record_ai_tokens(
    metrics: Any,
    input_tokens: int,
    output_tokens: int,
    *,
    model: Optional[str] = None,
) -> None:
    """Record AI token usage metrics.

    Args:
        metrics: Metrics instance
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        model: Model name/ID
    """
    if model:
        metrics.add_dimension(name="Model", value=model)

    if HAS_POWERTOOLS:
        metrics.add_metric(
            name="InputTokens",
            unit=MetricUnit.Count,
            value=input_tokens,
        )
        metrics.add_metric(
            name="OutputTokens",
            unit=MetricUnit.Count,
            value=output_tokens,
        )
        metrics.add_metric(
            name="TotalTokens",
            unit=MetricUnit.Count,
            value=input_tokens + output_tokens,
        )
    else:
        metrics.add_metric("InputTokens", value=input_tokens)
        metrics.add_metric("OutputTokens", value=output_tokens)
        metrics.add_metric("TotalTokens", value=input_tokens + output_tokens)


def record_tool_execution(
    metrics: Any,
    tool_name: str,
    duration_ms: float,
    *,
    success: bool = True,
) -> None:
    """Record tool execution metrics.

    Args:
        metrics: Metrics instance
        tool_name: Name of the tool
        duration_ms: Execution duration in milliseconds
        success: Whether the execution was successful
    """
    metrics.add_dimension(name="Tool", value=tool_name)

    if HAS_POWERTOOLS:
        metrics.add_metric(
            name="ToolExecutions",
            unit=MetricUnit.Count,
            value=1,
        )
        metrics.add_metric(
            name="ToolLatency",
            unit=MetricUnit.Milliseconds,
            value=duration_ms,
        )
        if not success:
            metrics.add_metric(
                name="ToolErrors",
                unit=MetricUnit.Count,
                value=1,
            )
    else:
        metrics.add_metric("ToolExecutions", value=1)
        metrics.add_metric("ToolLatency", value=duration_ms)


def record_circuit_breaker_state(
    metrics: Any,
    breaker_name: str,
    state: str,
) -> None:
    """Record circuit breaker state change.

    Args:
        metrics: Metrics instance
        breaker_name: Name of the circuit breaker
        state: Current state (closed, open, half_open)
    """
    metrics.add_dimension(name="CircuitBreaker", value=breaker_name)
    metrics.add_dimension(name="State", value=state)

    if HAS_POWERTOOLS:
        metrics.add_metric(
            name="CircuitBreakerStateChange",
            unit=MetricUnit.Count,
            value=1,
        )
    else:
        metrics.add_metric("CircuitBreakerStateChange", value=1)
