"""Tracing helpers — no-op when OTel not installed"""

import logging
import time
from functools import wraps
from typing import Any, Callable

logger = logging.getLogger(__name__)

# Try importing OTel — if not available, provide no-op implementations
try:
    from opentelemetry import trace
    from opentelemetry.trace import SpanKind, Status, StatusCode

    _OTEL_AVAILABLE = True
except ImportError:
    _OTEL_AVAILABLE = False

    # Provide fallback stubs so module-level references work in tests
    class _StubEnum:
        INTERNAL = "INTERNAL"
        CLIENT = "CLIENT"
        SERVER = "SERVER"
        PRODUCER = "PRODUCER"
        CONSUMER = "CONSUMER"

    SpanKind = _StubEnum  # type: ignore[misc,assignment]

    class Status:  # type: ignore[no-redef]
        def __init__(self, *args, **kwargs):
            pass

    class StatusCode:  # type: ignore[no-redef]
        UNSET = "UNSET"
        OK = "OK"
        ERROR = "ERROR"


class _NoOpSpan:
    """No-op span context manager used when OTel is not installed."""

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def set_attribute(self, *args):
        pass

    def add_event(self, *args):
        pass

    def set_status(self, *args):
        pass

    def record_exception(self, *args):
        pass


def get_tracer(name: str = "fischer.agentkit"):
    """Get tracer — returns None if OTel not installed."""
    if _OTEL_AVAILABLE:
        return trace.get_tracer(name)
    return None


def start_span(
    name: str,
    kind: Any = None,
    attributes: dict | None = None,
):
    """Start a span — returns no-op span if OTel not installed.

    Returns a context manager that yields a span (or no-op).
    """
    if not _OTEL_AVAILABLE:
        return _NoOpSpan()
    tracer = get_tracer()
    if tracer is None:
        return _NoOpSpan()
    if kind is None:
        kind = SpanKind.INTERNAL
    span = tracer.start_span(name, kind=kind, attributes=attributes)
    return trace.use_span(span, end_on_exit=True)


def trace_agent(agent_name: str, agent_type: str = "react"):
    """Decorator: trace agent execution."""

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            if not _OTEL_AVAILABLE:
                return await func(*args, **kwargs)
            tracer = get_tracer()
            with tracer.start_as_current_span(
                "agent.execute",
                kind=SpanKind.INTERNAL,
                attributes={"agent.name": agent_name, "agent.type": agent_type},
            ) as span:
                start = time.monotonic()
                try:
                    result = await func(*args, **kwargs)
                    duration_ms = int((time.monotonic() - start) * 1000)
                    span.set_attribute("agent.result.success", True)
                    span.set_attribute("agent.duration_ms", duration_ms)
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    span.set_attribute("agent.result.success", False)
                    raise

        return wrapper

    return decorator


def trace_tool(tool_name: str):
    """Decorator: trace tool call."""

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            if not _OTEL_AVAILABLE:
                return await func(*args, **kwargs)
            tracer = get_tracer()
            with tracer.start_as_current_span(
                "tool.execute",
                kind=SpanKind.CLIENT,
                attributes={"tool.name": tool_name},
            ) as span:
                start = time.monotonic()
                try:
                    result = await func(*args, **kwargs)
                    duration_ms = int((time.monotonic() - start) * 1000)
                    span.set_attribute("tool.duration_ms", duration_ms)
                    span.set_attribute("tool.result.success", True)
                    return result
                except Exception as e:
                    duration_ms = int((time.monotonic() - start) * 1000)
                    span.set_attribute("tool.duration_ms", duration_ms)
                    span.set_attribute("tool.result.success", False)
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise

        return wrapper

    return decorator


def trace_llm(provider: str, model: str):
    """Decorator: trace LLM call — follows GenAI Semantic Conventions."""

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            if not _OTEL_AVAILABLE:
                return await func(*args, **kwargs)
            tracer = get_tracer()
            with tracer.start_as_current_span(
                "gen_ai.chat",
                kind=SpanKind.CLIENT,
                attributes={
                    "gen_ai.system": provider,
                    "gen_ai.operation.name": "chat",
                    "gen_ai.request.model": model,
                },
            ) as span:
                start = time.monotonic()
                try:
                    result = await func(*args, **kwargs)
                    duration_ms = int((time.monotonic() - start) * 1000)
                    span.set_attribute("gen_ai.duration_ms", duration_ms)
                    # Record token usage if available on the response
                    if hasattr(result, "usage") and result.usage is not None:
                        span.set_attribute(
                            "gen_ai.usage.input_tokens",
                            getattr(result.usage, "prompt_tokens", 0),
                        )
                        span.set_attribute(
                            "gen_ai.usage.output_tokens",
                            getattr(result.usage, "completion_tokens", 0),
                        )
                    return result
                except Exception as e:
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise

        return wrapper

    return decorator


def trace_pipeline_step(pipeline_name: str, step_name: str):
    """Decorator: trace pipeline step execution."""

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            if not _OTEL_AVAILABLE:
                return await func(*args, **kwargs)
            tracer = get_tracer()
            with tracer.start_as_current_span(
                "pipeline.step",
                kind=SpanKind.INTERNAL,
                attributes={
                    "pipeline.name": pipeline_name,
                    "step.name": step_name,
                },
            ) as span:
                start = time.monotonic()
                try:
                    result = await func(*args, **kwargs)
                    duration_ms = int((time.monotonic() - start) * 1000)
                    span.set_attribute("step.duration_ms", duration_ms)
                    span.set_attribute("step.status", "success")
                    return result
                except Exception as e:
                    duration_ms = int((time.monotonic() - start) * 1000)
                    span.set_attribute("step.duration_ms", duration_ms)
                    span.set_attribute("step.status", "error")
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise

        return wrapper

    return decorator
