"""
OpenTelemetry middleware for agentbyte.

Provides automatic instrumentation following OpenTelemetry Gen-AI semantic conventions.
Enable with: AGENTBYTE_ENABLE_OTEL=true

References:
    - https://opentelemetry.io/docs/specs/semconv/gen-ai/
    - https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/
    - https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-metrics/

Usage:
    # Manual: add to agent middlewares list
    from agentbyte.middleware import OTelMiddleware
    agent = Agent(..., middlewares=[OTelMiddleware()])

    # Automatic: set env vars before importing agentbyte
    import os
    os.environ["AGENTBYTE_ENABLE_OTEL"] = "true"
    os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = "http://localhost:4318"
    os.environ["OTEL_SERVICE_NAME"] = "my-agent-app"
    from agentbyte import Agent  # OTelMiddleware prepended automatically
"""

import logging
import os
import time
from typing import Any, Optional

from .base import BaseMiddleware, MiddlewareContext

logger = logging.getLogger(__name__)

# Gracefully handle missing OpenTelemetry libraries.
# Install with: uv sync --group otel
try:
    from opentelemetry import context as otel_context
    from opentelemetry import metrics, trace
    from opentelemetry.exporter.otlp.proto.http.metric_exporter import (
        OTLPMetricExporter,
    )
    from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
    from opentelemetry.sdk.metrics import MeterProvider
    from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
    from opentelemetry.sdk.resources import Resource
    from opentelemetry.sdk.trace import TracerProvider
    from opentelemetry.sdk.trace.export import BatchSpanProcessor
    from opentelemetry.trace import Status, StatusCode

    OTEL_AVAILABLE = True
except ImportError:
    OTEL_AVAILABLE = False
    # Define placeholders so names always exist without TYPE_CHECKING-only imports.
    otel_context = None
    metrics = None
    trace = None
    OTLPMetricExporter = None
    OTLPSpanExporter = None
    MeterProvider = None
    PeriodicExportingMetricReader = None
    Resource = None
    TracerProvider = None
    BatchSpanProcessor = None
    Status = None
    StatusCode = None


_TELEMETRY_TRACER: Any = None
_TELEMETRY_METER: Any = None
_TELEMETRY_INITIALIZED = False


# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------


def _is_enabled() -> bool:
    """Check if OpenTelemetry is enabled via environment variable."""
    return os.getenv("AGENTBYTE_ENABLE_OTEL", "false").lower() in ("true", "1", "yes")


def _should_capture_content() -> bool:
    """
    Check if content capture is enabled via environment variable.

    Following OpenTelemetry Gen-AI semantic conventions, content attributes
    (gen_ai.input.messages, gen_ai.output.messages, gen_ai.tool.parameters,
    gen_ai.tool.result) are opt-in due to potential sensitive information (PII).

    Returns:
        bool: True only when explicitly enabled, False by default (production-safe).
    """
    return os.getenv("AGENTBYTE_OTEL_CAPTURE_CONTENT", "false").lower() in (
        "true",
        "1",
        "yes",
    )


def _setup_telemetry() -> "tuple[Any, Any]":
    """
    Set up OpenTelemetry TracerProvider and optional MeterProvider.

    Reads configuration from standard OTEL environment variables:
        OTEL_EXPORTER_OTLP_ENDPOINT  (default: http://localhost:4318)
        OTEL_SERVICE_NAME             (default: agentbyte)
        OTEL_METRICS_ENABLED          (default: false — Jaeger traces only)

    Returns:
        (tracer, meter) — meter is None when metrics are disabled or unavailable.
    """
    if not OTEL_AVAILABLE:
        logger.warning(
            "OpenTelemetry enabled but libraries not installed. "
            "Install with: uv sync --group otel"
        )
        return None, None

    global _TELEMETRY_TRACER, _TELEMETRY_METER, _TELEMETRY_INITIALIZED
    if _TELEMETRY_INITIALIZED:
        return _TELEMETRY_TRACER, _TELEMETRY_METER

    try:
        endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318")
        service_name = os.getenv("OTEL_SERVICE_NAME", "agentbyte")
        metrics_enabled = os.getenv("OTEL_METRICS_ENABLED", "false").lower() in (
            "true",
            "1",
            "yes",
        )

        resource = Resource.create(
            {
                "service.name": service_name,
                "gen_ai.system": "agentbyte",
            }
        )

        # Always set up tracing
        trace_provider = TracerProvider(resource=resource)
        trace_exporter = OTLPSpanExporter(endpoint=f"{endpoint}/v1/traces")
        trace_provider.add_span_processor(BatchSpanProcessor(trace_exporter))
        trace.set_tracer_provider(trace_provider)
        tracer = trace.get_tracer("agentbyte")

        # Metrics are optional — Jaeger and many backends support traces only
        meter: Any = None
        if metrics_enabled:
            try:
                metric_reader = PeriodicExportingMetricReader(
                    OTLPMetricExporter(endpoint=f"{endpoint}/v1/metrics")
                )
                meter_provider = MeterProvider(
                    resource=resource, metric_readers=[metric_reader]
                )
                metrics.set_meter_provider(meter_provider)
                meter = metrics.get_meter("agentbyte")
                logger.info(
                    "OpenTelemetry initialized: service=%s, endpoint=%s, metrics=enabled",
                    service_name,
                    endpoint,
                )
            except Exception as exc:
                logger.debug(
                    "Metrics setup failed (backend may not support metrics): %s", exc
                )
                logger.info(
                    "OpenTelemetry initialized: service=%s, endpoint=%s, metrics=disabled",
                    service_name,
                    endpoint,
                )
        else:
            logger.info(
                "OpenTelemetry initialized: service=%s, endpoint=%s, "
                "metrics=disabled (set OTEL_METRICS_ENABLED=true to enable)",
                service_name,
                endpoint,
            )

        _TELEMETRY_TRACER = tracer
        _TELEMETRY_METER = meter
        _TELEMETRY_INITIALIZED = True
        return tracer, meter

    except Exception as exc:
        logger.error("Failed to initialize OpenTelemetry: %s", exc)
        return None, None


# ---------------------------------------------------------------------------
# OTelMiddleware
# ---------------------------------------------------------------------------


class OTelMiddleware(BaseMiddleware):
    """
    OpenTelemetry middleware for agentbyte agents.

    Automatically instruments agent operations with distributed traces and optional
    metrics following OpenTelemetry Gen-AI semantic conventions.

    Instrumented operations:
        model_call  → span named "chat <model>"   with Gen-AI attributes + token metrics
        tool_call   → span named "tool <name>"    with tool success/failure attributes

    Always-captured attributes (metadata-safe, no PII exposure):
        gen_ai.system, gen_ai.operation.name, gen_ai.agent.name,
        gen_ai.request.model, gen_ai.usage.input_tokens, gen_ai.usage.output_tokens,
        gen_ai.usage.cost_estimate_usd, gen_ai.tool.name, gen_ai.tool.success,
        gen_ai.session.id (when available)

    Opt-in content attributes (set AGENTBYTE_OTEL_CAPTURE_CONTENT=true):
        gen_ai.input.messages, gen_ai.output.messages,
        gen_ai.tool.parameters, gen_ai.tool.result

    Environment variables:
        AGENTBYTE_ENABLE_OTEL            Enable telemetry (default: false)
        AGENTBYTE_OTEL_CAPTURE_CONTENT   Capture prompts/completions (default: false)
        OTEL_EXPORTER_OTLP_ENDPOINT      OTLP HTTP endpoint (default: http://localhost:4318)
        OTEL_SERVICE_NAME                Service name in traces (default: agentbyte)
        OTEL_METRICS_ENABLED             Enable metrics export (default: false)
    """

    def __init__(self) -> None:
        self._enabled = _is_enabled()
        self._capture_content = _should_capture_content()

        if not self._enabled:
            self._tracer: Any = None
            self._meter: Any = None
            self._token_histogram: Any = None
            self._duration_histogram: Any = None
            return

        self._tracer, self._meter = _setup_telemetry()

        # Disable entirely if tracer setup failed
        if not self._tracer:
            self._enabled = False
            self._meter = None
            self._token_histogram = None
            self._duration_histogram = None
            return

        # Create Gen-AI semantic convention metric instruments (optional)
        if self._meter:
            self._token_histogram = self._meter.create_histogram(
                name="gen_ai.client.token.usage",
                unit="{token}",
                description="Number of tokens used per request",
            )
            self._duration_histogram = self._meter.create_histogram(
                name="gen_ai.client.operation.duration",
                unit="s",
                description="Duration of AI operation in seconds",
            )
        else:
            self._token_histogram = None
            self._duration_histogram = None

    # ------------------------------------------------------------------
    # BaseMiddleware interface
    # ------------------------------------------------------------------

    async def process_request(self, context: MiddlewareContext) -> MiddlewareContext:
        """Start a span before the operation executes."""
        if not self._enabled or not self._tracer:
            return context

        try:
            span_name = self._build_span_name(context)
            span = self._tracer.start_span(span_name)

            # Attach span to the OTel context so nested operations become children
            ctx_token = otel_context.attach(trace.set_span_in_context(span))

            # Always-captured Gen-AI semantic convention attributes
            span.set_attribute("gen_ai.system", "agentbyte")
            span.set_attribute("gen_ai.operation.name", context.operation)
            span.set_attribute("gen_ai.agent.name", context.agent_name)

            if context.agent_context is not None and context.agent_context.session_id:
                span.set_attribute(
                    "gen_ai.session.id", context.agent_context.session_id
                )

            if context.operation == "model_call":
                model_name = self._get_model_name(context)
                span.set_attribute("gen_ai.request.model", model_name)

                # Opt-in: capture full prompt messages (may contain PII)
                if self._capture_content:
                    messages = (
                        context.data.get("messages", [])
                        if isinstance(context.data, dict)
                        else []
                    )
                    if messages:
                        try:
                            import json

                            span.set_attribute(
                                "gen_ai.input.messages",
                                json.dumps(self._format_input_messages(messages)),
                            )
                        except Exception as exc:
                            logger.debug("Failed to capture input messages: %s", exc)

            elif context.operation == "embedding_call":
                model_name = self._get_model_name(context)
                span.set_attribute("gen_ai.request.model", model_name)
                if isinstance(context.data, dict):
                    input_count = context.data.get("input_count", 1)
                    span.set_attribute("gen_ai.embedding.input_count", input_count)

            elif context.operation == "tool_call":
                tool_name = self._get_tool_name(context)
                span.set_attribute("gen_ai.tool.name", tool_name)

                # Opt-in: capture tool input arguments (may contain sensitive data)
                if self._capture_content:
                    parameters = (
                        context.data.get("parameters", {})
                        if isinstance(context.data, dict)
                        else {}
                    )
                    if parameters:
                        try:
                            import json

                            span.set_attribute(
                                "gen_ai.tool.parameters", json.dumps(parameters)
                            )
                        except Exception as exc:
                            logger.debug("Failed to capture tool parameters: %s", exc)

            # Stash span and timing for process_response / process_error
            context.metadata["_otel_span"] = span
            context.metadata["_otel_token"] = ctx_token
            context.metadata["_otel_start"] = time.time()

        except Exception as exc:
            logger.debug("Error starting OTel span: %s", exc)

        return context

    async def process_response(self, context: MiddlewareContext, result: Any) -> Any:
        """Record metrics and end the span with OK status."""
        if not self._enabled:
            return result

        span = context.metadata.get("_otel_span")
        if not span:
            return result

        try:
            duration_s = time.time() - context.metadata.get("_otel_start", time.time())

            # Duration histogram (always)
            if self._duration_histogram:
                self._duration_histogram.record(
                    duration_s, {"gen_ai.operation.name": context.operation}
                )

            if context.operation == "model_call" and hasattr(result, "usage"):
                usage = result.usage
                if hasattr(usage, "tokens_input"):
                    span.set_attribute("gen_ai.usage.input_tokens", usage.tokens_input)
                    if self._token_histogram:
                        self._token_histogram.record(
                            usage.tokens_input,
                            {
                                "gen_ai.token.type": "input",
                                "gen_ai.operation.name": context.operation,
                            },
                        )
                if hasattr(usage, "tokens_output"):
                    span.set_attribute(
                        "gen_ai.usage.output_tokens", usage.tokens_output
                    )
                    if self._token_histogram:
                        self._token_histogram.record(
                            usage.tokens_output,
                            {
                                "gen_ai.token.type": "output",
                                "gen_ai.operation.name": context.operation,
                            },
                        )
                if hasattr(usage, "cost_estimate") and usage.cost_estimate is not None:
                    span.set_attribute(
                        "gen_ai.usage.cost_estimate_usd", usage.cost_estimate
                    )

                if hasattr(usage, "total_tokens"):
                    span.set_attribute("gen_ai.usage.total_tokens", usage.total_tokens)

                finish_reason = getattr(result, "finish_reason", None)
                if finish_reason is None and hasattr(result, "metadata"):
                    metadata = getattr(result, "metadata", {})
                    if isinstance(metadata, dict):
                        finish_reason = metadata.get("finish_reason")

                if finish_reason is not None:
                    span.set_attribute(
                        "gen_ai.response.finish_reason", str(finish_reason)
                    )

                # Opt-in: capture model output messages
                if self._capture_content and hasattr(result, "message"):
                    try:
                        import json

                        span.set_attribute(
                            "gen_ai.output.messages",
                            json.dumps([self._format_output_message(result.message)]),
                        )
                    except Exception as exc:
                        logger.debug("Failed to capture output messages: %s", exc)

                # NEW: Capture parsed tool call arguments as span attributes
                # This addresses the limitation where args: {} appears empty in the output message
                # By capturing the parsed arguments, developers can see what the LLM actually decided
                # without jumping to individual tool spans.
                if hasattr(result, "message") and hasattr(result.message, "tool_calls"):
                    try:
                        import json

                        for i, tool_call in enumerate(result.message.tool_calls):
                            call_id = getattr(tool_call, "call_id", f"tool_call_{i}")
                            # ToolCallRequest uses 'parameters', not 'arguments'
                            parameters = getattr(tool_call, "parameters", {})
                            # Capture: gen_ai.tool_call.{call_id}.arguments
                            span.set_attribute(
                                f"gen_ai.tool_call.{call_id}.arguments",
                                json.dumps(parameters)
                                if isinstance(parameters, dict)
                                else str(parameters),
                            )
                    except Exception as exc:
                        logger.debug("Failed to capture parsed tool arguments: %s", exc)

            elif context.operation == "embedding_call":
                usage = getattr(result, "usage", None)
                if usage is not None:
                    if hasattr(usage, "tokens_input"):
                        span.set_attribute(
                            "gen_ai.usage.input_tokens", usage.tokens_input
                        )
                        if self._token_histogram:
                            self._token_histogram.record(
                                usage.tokens_input,
                                {
                                    "gen_ai.token.type": "input",
                                    "gen_ai.operation.name": context.operation,
                                },
                            )
                    if (
                        hasattr(usage, "cost_estimate")
                        and usage.cost_estimate is not None
                    ):
                        span.set_attribute(
                            "gen_ai.usage.cost_estimate_usd", usage.cost_estimate
                        )
                # Record actual number of embeddings returned
                if hasattr(result, "embeddings"):
                    span.set_attribute(
                        "gen_ai.embedding.output_count", len(result.embeddings)
                    )
                elif hasattr(result, "embedding"):
                    span.set_attribute("gen_ai.embedding.output_count", 1)
                if hasattr(result, "model"):
                    span.set_attribute("gen_ai.response.model", result.model)

            elif context.operation == "tool_call":
                if hasattr(result, "success"):
                    span.set_attribute("gen_ai.tool.success", result.success)

                # Opt-in: capture tool result
                if self._capture_content and hasattr(result, "result"):
                    try:
                        span.set_attribute("gen_ai.tool.result", str(result.result))
                    except Exception as exc:
                        logger.debug("Failed to capture tool result: %s", exc)

            span.set_status(Status(StatusCode.OK))

        except Exception as exc:
            logger.debug("Error recording OTel response telemetry: %s", exc)
        finally:
            self._end_span(context)

        return result

    async def process_error(
        self, context: MiddlewareContext, error: Exception
    ) -> Optional[Any]:
        """Record error status on the span and re-raise."""
        if not self._enabled:
            raise error

        span = context.metadata.get("_otel_span")
        if span:
            try:
                span.set_status(Status(StatusCode.ERROR, str(error)))
                span.set_attribute("error.type", type(error).__name__)
                span.set_attribute("error.message", str(error))
            except Exception as exc:
                logger.debug("Error recording OTel error telemetry: %s", exc)
            finally:
                self._end_span(context)

        raise error

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _build_span_name(self, context: MiddlewareContext) -> str:
        """Build Gen-AI convention span name."""
        if context.operation == "model_call":
            return f"chat {self._get_model_name(context)}"
        if context.operation == "tool_call":
            return f"tool {self._get_tool_name(context)}"
        if context.operation == "embedding_call":
            return f"embed {self._get_model_name(context)}"
        return f"{context.operation} {context.agent_name}"

    def _get_model_name(self, context: MiddlewareContext) -> str:
        """Extract model name injected by Agent into metadata."""
        return context.metadata.get("model", "unknown")

    def _get_tool_name(self, context: MiddlewareContext) -> str:
        """Extract tool name from data dict (agentbyte passes data as dict)."""
        if isinstance(context.data, dict):
            return context.data.get("tool_name", "unknown")
        # Fallback: dataclass / object with tool_name attribute
        return getattr(context.data, "tool_name", "unknown")

    def _end_span(self, context: MiddlewareContext) -> None:
        """Detach OTel context token and end the span."""
        try:
            ctx_token = context.metadata.get("_otel_token")
            if ctx_token is not None:
                otel_context.detach(ctx_token)
            span = context.metadata.get("_otel_span")
            if span:
                span.end()
        except Exception as exc:
            logger.debug("Error ending OTel span: %s", exc)

    def _format_input_messages(self, messages: list) -> list:
        """
        Convert agentbyte messages to Gen-AI semantic convention format.

        Format: [{"role": "user"|"assistant", "parts": [{"type": "text", "content": "..."}]}]
        Spec: https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-input-messages.json
        """
        formatted = []
        for msg in messages:
            role = "user"
            if hasattr(msg, "source") and msg.source != "user":
                role = "assistant"

            parts = []
            if hasattr(msg, "content") and msg.content:
                parts.append({"type": "text", "content": str(msg.content)})

            if hasattr(msg, "tool_calls") and msg.tool_calls:
                for tc in msg.tool_calls:
                    parts.append(
                        {
                            "type": "tool_call",
                            "id": getattr(tc, "call_id", ""),
                            "name": getattr(tc, "tool_name", ""),
                            "arguments": getattr(tc, "arguments", {}),
                        }
                    )

            if parts:
                formatted.append({"role": role, "parts": parts})

        return formatted

    def _format_output_message(self, message: Any) -> dict:
        """
        Convert a model response message to Gen-AI semantic convention format.

        Format: {"role": "assistant", "parts": [...]}
        Spec: https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-output-messages.json
        """
        parts = []
        if hasattr(message, "content") and message.content:
            parts.append({"type": "text", "content": str(message.content)})

        if hasattr(message, "tool_calls") and message.tool_calls:
            for tc in message.tool_calls:
                parts.append(
                    {
                        "type": "tool_call",
                        "id": getattr(tc, "call_id", ""),
                        "name": getattr(tc, "tool_name", ""),
                        "arguments": getattr(tc, "arguments", {}),
                    }
                )

        return {"role": "assistant", "parts": parts}


# ---------------------------------------------------------------------------
# Auto-instrumentation
# ---------------------------------------------------------------------------


def auto_instrument() -> None:
    """
    Auto-instrument agentbyte agents with OpenTelemetry.

    Patches Agent.__init__ to prepend :class:`OTelMiddleware` to every agent's
    middleware chain so that per-call child spans (``chat <model>``, ``tool <name>``)
    are emitted automatically.  No run/run_stream monkey-patching.

    Called automatically at import time when ``AGENTBYTE_ENABLE_OTEL=true``.
    Transparent to calling code — no agent construction or invocation changes
    are required.
    """
    if not _is_enabled():
        return

    try:
        from agentbyte.agents.agent import Agent

        if getattr(Agent.__init__, "__agentbyte_otel_instrumented__", False):
            return

        original_init = Agent.__init__

        def _instrumented_init(
            self: Any, *args: Any, middlewares: Any = None, **kwargs: Any
        ) -> None:
            middlewares = list(middlewares) if middlewares else []
            if not any(isinstance(m, OTelMiddleware) for m in middlewares):
                middlewares.insert(0, OTelMiddleware())
            original_init(self, *args, middlewares=middlewares, **kwargs)

        _instrumented_init.__agentbyte_otel_instrumented__ = True  # type: ignore[attr-defined]
        _instrumented_init.__agentbyte_original_init__ = original_init  # type: ignore[attr-defined]

        Agent.__init__ = _instrumented_init  # type: ignore[method-assign]

        logger.info("agentbyte: OpenTelemetry auto-instrumentation enabled")

    except Exception as exc:
        logger.error("agentbyte: Failed to enable OTel auto-instrumentation: %s", exc)
