# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
OpenTelemetry Wrapper for Familiar (v1.4.0)

Provides distributed tracing, metrics, and observability:
- OpenTelemetry-compatible tracing with spans
- LLM-specific instrumentation (tokens, latency, cost)
- Metric collection and export
- Context propagation across async boundaries
- Multiple exporter support (Console, OTLP, Jaeger)
- Decorators for easy instrumentation

Usage:
    from familiar.core.observability import (
        init_telemetry, get_tracer, trace_llm_call,
        Span, create_span, with_span
    )

    # Initialize with OTLP exporter
    init_telemetry(
        service_name="my-agent",
        exporter="otlp",
        endpoint="http://localhost:4317"
    )

    # Manual span creation
    with create_span("process_request") as span:
        span.set_attribute("user_id", "123")
        result = process(request)
        span.set_attribute("result_size", len(result))

    # Decorator-based tracing
    @with_span("my_function")
    async def my_function(x):
        return x * 2

    # LLM-specific tracing
    with trace_llm_call("gpt-4", prompt) as span:
        response = await llm.generate(prompt)
        span.record_completion(response, tokens=150, cost=0.003)
"""

import asyncio
import functools
import json
import logging
import threading
import time
import traceback
import uuid
from abc import ABC, abstractmethod
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from enum import Enum
from typing import (
    Any,
    AsyncGenerator,
    Callable,
    Dict,
    Generator,
    List,
    Optional,
    TypeVar,
    Union,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


# ============================================================
# ENUMS AND TYPES
# ============================================================


class SpanKind(str, Enum):
    """Type of span (follows OpenTelemetry conventions)."""

    INTERNAL = "internal"  # Internal operation
    SERVER = "server"  # Server-side of RPC
    CLIENT = "client"  # Client-side of RPC
    PRODUCER = "producer"  # Message producer
    CONSUMER = "consumer"  # Message consumer


class SpanStatus(str, Enum):
    """Status of a span."""

    UNSET = "unset"
    OK = "ok"
    ERROR = "error"


class ExporterType(str, Enum):
    """Supported exporter types."""

    CONSOLE = "console"  # Print to console
    OTLP = "otlp"  # OpenTelemetry Protocol
    JAEGER = "jaeger"  # Jaeger
    ZIPKIN = "zipkin"  # Zipkin
    MEMORY = "memory"  # In-memory (for testing)
    NONE = "none"  # No export


# ============================================================
# SPAN CONTEXT
# ============================================================


@dataclass
class SpanContext:
    """
    Context for distributed tracing.

    Contains trace_id and span_id for correlation across services.
    """

    trace_id: str
    span_id: str
    parent_span_id: Optional[str] = None
    trace_flags: int = 1  # Sampled by default
    trace_state: Dict[str, str] = field(default_factory=dict)

    @classmethod
    def create_root(cls) -> "SpanContext":
        """Create a new root context (new trace)."""
        return cls(
            trace_id=uuid.uuid4().hex,
            span_id=uuid.uuid4().hex[:16],
        )

    def create_child(self) -> "SpanContext":
        """Create a child context (same trace, new span)."""
        return SpanContext(
            trace_id=self.trace_id,
            span_id=uuid.uuid4().hex[:16],
            parent_span_id=self.span_id,
            trace_flags=self.trace_flags,
            trace_state=self.trace_state.copy(),
        )

    def to_dict(self) -> dict:
        return {
            "trace_id": self.trace_id,
            "span_id": self.span_id,
            "parent_span_id": self.parent_span_id,
            "trace_flags": self.trace_flags,
        }

    def to_traceparent(self) -> str:
        """Convert to W3C traceparent header format."""
        return f"00-{self.trace_id}-{self.span_id}-{self.trace_flags:02x}"

    @classmethod
    def from_traceparent(cls, header: str) -> Optional["SpanContext"]:
        """Parse W3C traceparent header."""
        try:
            parts = header.split("-")
            if len(parts) >= 4 and parts[0] == "00":
                return cls(
                    trace_id=parts[1],
                    span_id=parts[2],
                    trace_flags=int(parts[3], 16),
                )
        except Exception:
            pass
        return None


# ============================================================
# SPAN EVENT
# ============================================================


@dataclass
class SpanEvent:
    """An event that occurred during a span's lifetime."""

    name: str
    timestamp: float  # Unix timestamp
    attributes: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "timestamp": self.timestamp,
            "attributes": self.attributes,
        }


# ============================================================
# SPAN
# ============================================================


@dataclass
class Span:
    """
    A span representing a unit of work in a trace.

    Spans can have:
    - Attributes (key-value metadata)
    - Events (timestamped logs)
    - Status (ok/error)
    - Links to other spans
    """

    name: str
    context: SpanContext
    kind: SpanKind = SpanKind.INTERNAL

    # Timing
    start_time: float = 0.0  # Unix timestamp
    end_time: float = 0.0

    # Status
    status: SpanStatus = SpanStatus.UNSET
    status_message: str = ""

    # Data
    attributes: Dict[str, Any] = field(default_factory=dict)
    events: List[SpanEvent] = field(default_factory=list)
    links: List[SpanContext] = field(default_factory=list)

    # Resource info
    service_name: str = ""

    # Internal state
    _ended: bool = False
    _tracer: Optional["Tracer"] = None

    def __post_init__(self):
        if self.start_time == 0.0:
            self.start_time = time.time()

    def set_attribute(self, key: str, value: Any) -> "Span":
        """Set a span attribute."""
        if not self._ended:
            self.attributes[key] = value
        return self

    def set_attributes(self, attributes: Dict[str, Any]) -> "Span":
        """Set multiple attributes."""
        if not self._ended:
            self.attributes.update(attributes)
        return self

    def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> "Span":
        """Add an event to the span."""
        if not self._ended:
            self.events.append(
                SpanEvent(name=name, timestamp=time.time(), attributes=attributes or {})
            )
        return self

    def add_link(self, context: SpanContext) -> "Span":
        """Add a link to another span."""
        if not self._ended:
            self.links.append(context)
        return self

    def set_status(self, status: SpanStatus, message: str = "") -> "Span":
        """Set the span status."""
        if not self._ended:
            self.status = status
            self.status_message = message
        return self

    def record_exception(self, exception: Exception, escaped: bool = True) -> "Span":
        """Record an exception as an event."""
        if not self._ended:
            self.add_event(
                "exception",
                {
                    "exception.type": type(exception).__name__,
                    "exception.message": str(exception),
                    "exception.stacktrace": traceback.format_exc(),
                    "exception.escaped": escaped,
                },
            )
            self.set_status(SpanStatus.ERROR, str(exception))
        return self

    def end(self, end_time: Optional[float] = None) -> None:
        """End the span."""
        if self._ended:
            return

        self.end_time = end_time or time.time()
        self._ended = True

        if self._tracer:
            self._tracer._on_span_end(self)

    @property
    def duration_ms(self) -> float:
        """Get span duration in milliseconds."""
        end = self.end_time or time.time()
        return (end - self.start_time) * 1000

    @property
    def trace_id(self) -> str:
        return self.context.trace_id

    @property
    def span_id(self) -> str:
        return self.context.span_id

    @property
    def parent_span_id(self) -> Optional[str]:
        return self.context.parent_span_id

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "trace_id": self.trace_id,
            "span_id": self.span_id,
            "parent_span_id": self.parent_span_id,
            "kind": self.kind.value,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": self.duration_ms,
            "status": self.status.value,
            "status_message": self.status_message,
            "attributes": self.attributes,
            "events": [e.to_dict() for e in self.events],
            "service_name": self.service_name,
        }

    def __enter__(self) -> "Span":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if exc_val:
            self.record_exception(exc_val)
        elif self.status == SpanStatus.UNSET:
            self.set_status(SpanStatus.OK)
        self.end()

    async def __aenter__(self) -> "Span":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if exc_val:
            self.record_exception(exc_val)
        elif self.status == SpanStatus.UNSET:
            self.set_status(SpanStatus.OK)
        self.end()

    # --- Backwards-compatibility helpers ---

    @classmethod
    def create_compat(cls, name: str, trace_id: str = None, **extra_attrs) -> "Span":
        """Create a span with minimal arguments (backwards compatibility)."""
        if trace_id:
            context = SpanContext(trace_id=trace_id, span_id=uuid.uuid4().hex[:16])
        else:
            context = SpanContext.create_root()
        span = cls(name=name, context=context)
        span.attributes.update(extra_attrs)
        return span

    def finish(self) -> None:
        """End the span (backwards-compat alias for end())."""
        self.end()

    def set_error(self, exception: Exception, status: str = "error") -> None:
        """Record an error on the span (backwards compatibility)."""
        self.record_exception(exception)
        self.set_attribute("error_status", status)

    @property
    def metadata(self) -> Dict[str, Any]:
        """Alias for attributes (backwards compatibility)."""
        return self.attributes


# ============================================================
# LLM SPAN (Specialized for LLM calls)
# ============================================================


class LLMSpan(Span):
    """
    Specialized span for LLM API calls.

    Tracks LLM-specific attributes like tokens, model, cost.
    """

    def __init__(
        self, name: str, context: SpanContext, model: str = "", provider: str = "", **kwargs
    ):
        super().__init__(name=name, context=context, kind=SpanKind.CLIENT, **kwargs)
        self.set_attributes(
            {
                "llm.model": model,
                "llm.provider": provider,
                "llm.request_type": "completion",
            }
        )

    def set_prompt(self, prompt: str, truncate: int = 500) -> "LLMSpan":
        """Record the prompt (truncated for safety)."""
        self.set_attribute("llm.prompt", prompt[:truncate] if len(prompt) > truncate else prompt)
        self.set_attribute("llm.prompt_length", len(prompt))
        return self

    def set_completion(self, completion: str, truncate: int = 500) -> "LLMSpan":
        """Record the completion (truncated for safety)."""
        self.set_attribute(
            "llm.completion", completion[:truncate] if len(completion) > truncate else completion
        )
        self.set_attribute("llm.completion_length", len(completion))
        return self

    def record_tokens(
        self, prompt_tokens: int = 0, completion_tokens: int = 0, total_tokens: Optional[int] = None
    ) -> "LLMSpan":
        """Record token usage."""
        self.set_attributes(
            {
                "llm.token_count.prompt": prompt_tokens,
                "llm.token_count.completion": completion_tokens,
                "llm.token_count.total": total_tokens or (prompt_tokens + completion_tokens),
            }
        )
        return self

    def record_cost(self, cost_usd: float) -> "LLMSpan":
        """Record cost in USD."""
        self.set_attribute("llm.cost_usd", cost_usd)
        return self

    def record_completion(
        self,
        completion: str,
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
        cost_usd: float = 0.0,
    ) -> "LLMSpan":
        """Convenience method to record all completion details."""
        self.set_completion(completion)
        self.record_tokens(prompt_tokens, completion_tokens)
        if cost_usd > 0:
            self.record_cost(cost_usd)
        return self

    def set_temperature(self, temperature: float) -> "LLMSpan":
        """Record temperature setting."""
        self.set_attribute("llm.temperature", temperature)
        return self

    def set_max_tokens(self, max_tokens: int) -> "LLMSpan":
        """Record max_tokens setting."""
        self.set_attribute("llm.max_tokens", max_tokens)
        return self


# ============================================================
# SPAN EXPORTER
# ============================================================


class SpanExporter(ABC):
    """Abstract base for span exporters."""

    @abstractmethod
    def export(self, spans: List[Span]) -> bool:
        """Export spans. Returns True on success."""
        pass

    def shutdown(self) -> None:
        """Clean shutdown."""
        pass


class ConsoleSpanExporter(SpanExporter):
    """Export spans to console (for development)."""

    def __init__(self, pretty: bool = True):
        self.pretty = pretty

    def export(self, spans: List[Span]) -> bool:
        for span in spans:
            if self.pretty:
                self._print_pretty(span)
            else:
                print(json.dumps(span.to_dict()))
        return True

    def _print_pretty(self, span: Span) -> None:
        status_icon = (
            "✓" if span.status == SpanStatus.OK else "✗" if span.status == SpanStatus.ERROR else "○"
        )
        indent = "  " if span.parent_span_id else ""

        print(f"{indent}{status_icon} {span.name} ({span.duration_ms:.1f}ms)")
        print(f"{indent}  trace_id: {span.trace_id[:8]}... span_id: {span.span_id}")

        if span.attributes:
            for key, value in list(span.attributes.items())[:5]:
                print(f"{indent}  {key}: {str(value)[:50]}")

        if span.events:
            for event in span.events[:3]:
                print(f"{indent}  📍 {event.name}")


class MemorySpanExporter(SpanExporter):
    """Export spans to memory (for testing)."""

    def __init__(self):
        self.spans: List[Span] = []
        self._lock = threading.Lock()

    def export(self, spans: List[Span]) -> bool:
        with self._lock:
            self.spans.extend(spans)
        return True

    def get_spans(self) -> List[Span]:
        with self._lock:
            return list(self.spans)

    def clear(self) -> None:
        with self._lock:
            self.spans.clear()

    def find_by_name(self, name: str) -> List[Span]:
        return [s for s in self.spans if s.name == name]

    def find_by_trace(self, trace_id: str) -> List[Span]:
        return [s for s in self.spans if s.trace_id == trace_id]


class OTLPSpanExporter(SpanExporter):
    """
    Export spans via OpenTelemetry Protocol (OTLP).

    Requires: pip install opentelemetry-exporter-otlp
    """

    def __init__(
        self,
        endpoint: str = "http://localhost:4317",
        headers: Optional[Dict[str, str]] = None,
        insecure: bool = True,
    ):
        self.endpoint = endpoint
        self.headers = headers or {}
        self.insecure = insecure
        self._exporter = None
        self._init_exporter()

    def _init_exporter(self):
        """Initialize the OTLP exporter if available."""
        try:
            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
                OTLPSpanExporter as OTLP,
            )

            self._exporter = OTLP(
                endpoint=self.endpoint,
                insecure=self.insecure,
                headers=self.headers,
            )
        except ImportError:
            logger.warning("OTLP exporter not available.")

    def export(self, spans: List[Span]) -> bool:
        if not self._exporter:
            # Fallback to console
            for span in spans:
                logger.info(f"Span: {span.name} ({span.duration_ms:.1f}ms)")
            return True

        # Convert to OpenTelemetry format
        try:
            # This is simplified - real implementation would create proper OT spans
            for span in spans:
                logger.debug(f"Exporting span: {span.name}")

            return True
        except Exception as e:
            logger.error(f"OTLP export failed: {e}")
            return False

    def shutdown(self) -> None:
        if self._exporter:
            self._exporter.shutdown()


class JSONFileExporter(SpanExporter):
    """Export spans to a JSON file."""

    def __init__(self, filepath: str):
        self.filepath = filepath
        self._lock = threading.Lock()

    def export(self, spans: List[Span]) -> bool:
        with self._lock:
            try:
                # Append to file
                with open(self.filepath, "a") as f:
                    for span in spans:
                        f.write(json.dumps(span.to_dict()) + "\n")
                return True
            except Exception as e:
                logger.error(f"File export failed: {e}")
                return False


# ============================================================
# METRIC TYPES
# ============================================================


@dataclass
class MetricPoint:
    """A single metric measurement."""

    name: str
    value: float
    timestamp: float = 0.0
    labels: Dict[str, str] = field(default_factory=dict)

    def __post_init__(self):
        if self.timestamp == 0.0:
            self.timestamp = time.time()


class Counter:
    """A monotonically increasing counter."""

    def __init__(self, name: str, description: str = "", unit: str = ""):
        self.name = name
        self.description = description
        self.unit = unit
        self._value = 0.0
        self._lock = threading.Lock()
        self._labels_values: Dict[str, float] = {}

    def add(self, value: float = 1, labels: Optional[Dict[str, str]] = None) -> None:
        """Add to the counter."""
        with self._lock:
            self._value += value
            if labels:
                key = json.dumps(labels, sort_keys=True)
                self._labels_values[key] = self._labels_values.get(key, 0) + value

    def get(self, labels: Optional[Dict[str, str]] = None) -> float:
        """Get current value."""
        with self._lock:
            if labels:
                key = json.dumps(labels, sort_keys=True)
                return self._labels_values.get(key, 0)
            return self._value


class Gauge:
    """A metric that can go up or down."""

    def __init__(self, name: str, description: str = "", unit: str = ""):
        self.name = name
        self.description = description
        self.unit = unit
        self._value = 0.0
        self._lock = threading.Lock()

    def set(self, value: float) -> None:
        """Set the gauge value."""
        with self._lock:
            self._value = value

    def inc(self, value: float = 1) -> None:
        """Increment the gauge."""
        with self._lock:
            self._value += value

    def dec(self, value: float = 1) -> None:
        """Decrement the gauge."""
        with self._lock:
            self._value -= value

    def get(self) -> float:
        """Get current value."""
        with self._lock:
            return self._value


class Histogram:
    """A metric that tracks value distributions."""

    DEFAULT_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10)

    def __init__(
        self, name: str, description: str = "", unit: str = "", buckets: tuple = DEFAULT_BUCKETS
    ):
        self.name = name
        self.description = description
        self.unit = unit
        self.buckets = buckets
        self._max_values = 10000  # Sliding window to prevent unbounded growth
        self._values: List[float] = []
        self._sum = 0.0
        self._count = 0
        self._lock = threading.Lock()

    def observe(self, value: float) -> None:
        """Record a value."""
        with self._lock:
            self._values.append(value)
            self._sum += value
            self._count += 1
            # Prune oldest values if exceeding max window size
            if len(self._values) > self._max_values:
                removed = self._values[: len(self._values) - self._max_values]
                self._values = self._values[-self._max_values :]
                self._sum -= sum(removed)

    def get_stats(self) -> Dict[str, float]:
        """Get histogram statistics."""
        with self._lock:
            if not self._values:
                return {"count": 0, "sum": 0, "mean": 0, "p50": 0, "p95": 0, "p99": 0}

            sorted_values = sorted(self._values)
            n = len(sorted_values)

            return {
                "count": self._count,
                "sum": self._sum,
                "mean": self._sum / self._count,
                "min": sorted_values[0],
                "max": sorted_values[-1],
                "p50": sorted_values[int(n * 0.5)],
                "p95": sorted_values[int(n * 0.95)] if n >= 20 else sorted_values[-1],
                "p99": sorted_values[int(n * 0.99)] if n >= 100 else sorted_values[-1],
            }


# ============================================================
# TRACER
# ============================================================


class Tracer:
    """
    Main tracer for creating and managing spans.

    Handles span creation, context propagation, and export.
    """

    def __init__(
        self,
        service_name: str = "familiar",
        exporter: Optional[SpanExporter] = None,
        sample_rate: float = 1.0,  # 1.0 = sample everything
        max_spans_per_trace: int = 1000,
    ):
        self.service_name = service_name
        self.exporter = exporter or ConsoleSpanExporter()
        self.sample_rate = sample_rate
        self.max_spans_per_trace = max_spans_per_trace

        # Context storage (thread-local and async context var)
        self._context_var = None
        self._thread_local = threading.local()

        # Pending spans for batch export
        self._pending_spans: List[Span] = []
        self._pending_lock = threading.Lock()

        # Metrics
        self._metrics: Dict[str, Union[Counter, Gauge, Histogram]] = {}

        # Initialize context variable for async
        try:
            import contextvars

            self._context_var = contextvars.ContextVar("current_span", default=None)
        except ImportError:
            pass

    def create_span(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        parent: Optional[Union[Span, SpanContext]] = None,
        attributes: Optional[Dict[str, Any]] = None,
        links: Optional[List[SpanContext]] = None,
    ) -> Span:
        """Create a new span."""

        # Determine parent context
        if parent is None:
            parent = self.get_current_span()

        if isinstance(parent, Span):
            context = parent.context.create_child()
        elif isinstance(parent, SpanContext):
            context = parent.create_child()
        else:
            context = SpanContext.create_root()

        # Check sampling
        if self.sample_rate < 1.0:
            import random

            if random.random() > self.sample_rate:
                context.trace_flags = 0  # Not sampled

        span = Span(
            name=name,
            context=context,
            kind=kind,
            service_name=self.service_name,
            _tracer=self,
        )

        if attributes:
            span.set_attributes(attributes)

        if links:
            for link in links:
                span.add_link(link)

        return span

    def create_llm_span(
        self,
        name: str,
        model: str,
        provider: str = "",
        parent: Optional[Union[Span, SpanContext]] = None,
    ) -> LLMSpan:
        """Create an LLM-specific span."""

        if parent is None:
            parent = self.get_current_span()

        if isinstance(parent, Span):
            context = parent.context.create_child()
        elif isinstance(parent, SpanContext):
            context = parent.create_child()
        else:
            context = SpanContext.create_root()

        span = LLMSpan(
            name=name,
            context=context,
            model=model,
            provider=provider,
            service_name=self.service_name,
            _tracer=self,
        )

        return span

    @contextmanager
    def start_span(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> Generator[Span, None, None]:
        """Context manager for span lifecycle."""
        span = self.create_span(name, kind=kind, attributes=attributes)
        token = self._set_current_span(span)

        try:
            yield span
            if span.status == SpanStatus.UNSET:
                span.set_status(SpanStatus.OK)
        except Exception as e:
            span.record_exception(e)
            raise
        finally:
            self._reset_current_span(token)
            span.end()

    @asynccontextmanager
    async def start_span_async(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> AsyncGenerator[Span, None]:
        """Async context manager for span lifecycle."""
        span = self.create_span(name, kind=kind, attributes=attributes)
        token = self._set_current_span(span)

        try:
            yield span
            if span.status == SpanStatus.UNSET:
                span.set_status(SpanStatus.OK)
        except Exception as e:
            span.record_exception(e)
            raise
        finally:
            self._reset_current_span(token)
            span.end()

    def get_current_span(self) -> Optional[Span]:
        """Get the current active span."""
        if self._context_var:
            return self._context_var.get()
        return getattr(self._thread_local, "span", None)

    def _set_current_span(self, span: Span):
        """Set the current span and return a token for reset."""
        if self._context_var:
            return self._context_var.set(span)
        else:
            old = getattr(self._thread_local, "span", None)
            self._thread_local.span = span
            return old

    def _reset_current_span(self, token):
        """Reset to previous span."""
        if self._context_var:
            self._context_var.reset(token)
        else:
            self._thread_local.span = token

    def _on_span_end(self, span: Span) -> None:
        """Called when a span ends."""
        if span.context.trace_flags == 0:
            # Not sampled
            return

        with self._pending_lock:
            self._pending_spans.append(span)

            # Batch export when we have enough spans
            if len(self._pending_spans) >= 10:
                spans_to_export = self._pending_spans
                self._pending_spans = []
            else:
                spans_to_export = None

        if spans_to_export:
            self._export(spans_to_export)

    def _export(self, spans: List[Span]) -> None:
        """Export spans."""
        if self.exporter:
            try:
                self.exporter.export(spans)
            except Exception as e:
                logger.error(f"Span export failed: {e}")

    def flush(self) -> None:
        """Flush pending spans."""
        with self._pending_lock:
            if self._pending_spans:
                self._export(self._pending_spans)
                self._pending_spans = []

    def shutdown(self) -> None:
        """Shutdown the tracer."""
        self.flush()
        if self.exporter:
            self.exporter.shutdown()

    # --------------------------------------------------------
    # METRICS
    # --------------------------------------------------------

    def counter(self, name: str, description: str = "", unit: str = "") -> Counter:
        """Get or create a counter metric."""
        if name not in self._metrics:
            self._metrics[name] = Counter(name, description, unit)
        return self._metrics[name]

    def gauge(self, name: str, description: str = "", unit: str = "") -> Gauge:
        """Get or create a gauge metric."""
        if name not in self._metrics:
            self._metrics[name] = Gauge(name, description, unit)
        return self._metrics[name]

    def histogram(
        self,
        name: str,
        description: str = "",
        unit: str = "",
        buckets: tuple = Histogram.DEFAULT_BUCKETS,
    ) -> Histogram:
        """Get or create a histogram metric."""
        if name not in self._metrics:
            self._metrics[name] = Histogram(name, description, unit, buckets)
        return self._metrics[name]

    def get_metrics_summary(self) -> Dict[str, Any]:
        """Get summary of all metrics."""
        summary = {}
        for name, metric in self._metrics.items():
            if isinstance(metric, Counter):
                summary[name] = {"type": "counter", "value": metric.get()}
            elif isinstance(metric, Gauge):
                summary[name] = {"type": "gauge", "value": metric.get()}
            elif isinstance(metric, Histogram):
                summary[name] = {"type": "histogram", **metric.get_stats()}
        return summary


# ============================================================
# GLOBAL INSTANCE
# ============================================================

_tracer: Optional[Tracer] = None


def get_tracer() -> Tracer:
    """Get the global tracer instance."""
    global _tracer
    if _tracer is None:
        _tracer = Tracer()
    return _tracer


def set_tracer(tracer: Tracer) -> None:
    """Set the global tracer instance."""
    global _tracer
    _tracer = tracer


def init_telemetry(
    service_name: str = "familiar",
    exporter: Union[str, ExporterType] = ExporterType.CONSOLE,
    endpoint: str = "http://localhost:4317",
    sample_rate: float = 1.0,
    **kwargs,
) -> Tracer:
    """
    Initialize telemetry with the specified exporter.

    Args:
        service_name: Name of your service
        exporter: Type of exporter (console, otlp, memory, etc.)
        endpoint: Endpoint for OTLP/Jaeger exporters
        sample_rate: Fraction of traces to sample (0.0-1.0)

    Returns:
        Configured Tracer instance
    """
    global _tracer

    if isinstance(exporter, str):
        exporter = ExporterType(exporter)

    # Create exporter
    if exporter == ExporterType.CONSOLE:
        span_exporter = ConsoleSpanExporter(pretty=kwargs.get("pretty", True))
    elif exporter == ExporterType.MEMORY:
        span_exporter = MemorySpanExporter()
    elif exporter == ExporterType.OTLP:
        span_exporter = OTLPSpanExporter(
            endpoint=endpoint,
            headers=kwargs.get("headers"),
            insecure=kwargs.get("insecure", True),
        )
    elif exporter == ExporterType.NONE:
        span_exporter = None
    else:
        span_exporter = ConsoleSpanExporter()

    _tracer = Tracer(
        service_name=service_name,
        exporter=span_exporter,
        sample_rate=sample_rate,
    )

    logger.info(f"Telemetry initialized: service={service_name}, exporter={exporter.value}")

    return _tracer


# ============================================================
# DECORATORS
# ============================================================


def with_span(
    name: Optional[str] = None,
    kind: SpanKind = SpanKind.INTERNAL,
    attributes: Optional[Dict[str, Any]] = None,
):
    """
    Decorator to trace a function.

    Usage:
        @with_span("my_operation")
        def my_function(x, y):
            return x + y

        @with_span()  # Uses function name
        async def async_function():
            await something()
    """

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        span_name = name or func.__name__

        if asyncio.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                tracer = get_tracer()
                async with tracer.start_span_async(span_name, kind, attributes):
                    return await func(*args, **kwargs)

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args, **kwargs):
                tracer = get_tracer()
                with tracer.start_span(span_name, kind, attributes):
                    return func(*args, **kwargs)

            return sync_wrapper

    return decorator


def trace_llm_call(
    model: str,
    prompt: Optional[str] = None,
    provider: str = "",
):
    """
    Context manager for tracing LLM API calls.

    Usage:
        with trace_llm_call("gpt-4", prompt) as span:
            response = await openai.chat.completions.create(...)
            span.record_completion(
                response.choices[0].message.content,
                prompt_tokens=response.usage.prompt_tokens,
                completion_tokens=response.usage.completion_tokens,
            )
    """
    tracer = get_tracer()
    span = tracer.create_llm_span(
        name=f"llm.{model}",
        model=model,
        provider=provider,
    )

    if prompt:
        span.set_prompt(prompt)

    return span


# ============================================================
# CONVENIENCE FUNCTIONS
# ============================================================


def create_span(
    name: str,
    kind: SpanKind = SpanKind.INTERNAL,
    attributes: Optional[Dict[str, Any]] = None,
) -> Span:
    """Create a span using the global tracer."""
    return get_tracer().create_span(name, kind=kind, attributes=attributes)


def start_span(
    name: str,
    kind: SpanKind = SpanKind.INTERNAL,
    attributes: Optional[Dict[str, Any]] = None,
):
    """Context manager for spans using global tracer."""
    return get_tracer().start_span(name, kind, attributes)


def get_current_span() -> Optional[Span]:
    """Get the current active span."""
    return get_tracer().get_current_span()


def record_metric(name: str, value: float, labels: Optional[Dict[str, str]] = None):
    """Record a metric value."""
    tracer = get_tracer()
    counter = tracer.counter(name)
    counter.add(value, labels)


# ============================================================
# COST CALCULATION
# ============================================================

# Model pricing (per 1K tokens)
MODEL_PRICING = {
    # OpenAI
    "gpt-4": {"input": 0.03, "output": 0.06},
    "gpt-4-turbo": {"input": 0.01, "output": 0.03},
    "gpt-4o": {"input": 0.005, "output": 0.015},
    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
    "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
    # Anthropic
    "claude-3-opus": {"input": 0.015, "output": 0.075},
    "claude-3-sonnet": {"input": 0.003, "output": 0.015},
    "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
    "claude-3.5-sonnet": {"input": 0.003, "output": 0.015},
    "claude-opus-4": {"input": 0.015, "output": 0.075},
    "claude-sonnet-4": {"input": 0.003, "output": 0.015},
}


def calculate_cost(
    model: str,
    prompt_tokens: int,
    completion_tokens: int,
) -> float:
    """Calculate cost for an LLM call."""
    # Find matching model
    pricing = None
    for model_key, prices in MODEL_PRICING.items():
        if model_key in model.lower():
            pricing = prices
            break

    if not pricing:
        return 0.0

    input_cost = (prompt_tokens / 1000) * pricing["input"]
    output_cost = (completion_tokens / 1000) * pricing["output"]

    return input_cost + output_cost


def get_model_pricing(model: str) -> Optional[Dict[str, float]]:
    """Get pricing for a model."""
    for model_key, prices in MODEL_PRICING.items():
        if model_key in model.lower():
            return prices
    return None


# ============================================================
# TRACE CLASS (Backwards compatibility)
# ============================================================


class Trace:
    """
    Backwards-compatible Trace class.

    Wraps the new Span-based tracing.
    """

    def __init__(self, name: str, metadata: Optional[Dict] = None):
        self.name = name
        self.metadata = metadata or {}
        self._span: Optional[Span] = None
        self._child_spans: List[Span] = []
        self._tracer = get_tracer()

    def __enter__(self) -> "Trace":
        self._span = self._tracer.create_span(self.name, attributes=self.metadata)
        self._tracer._set_current_span(self._span)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._span:
            if exc_val:
                self._span.record_exception(exc_val)
            else:
                self._span.set_status(SpanStatus.OK)
            self._span.end()

    def add_metadata(self, key: str, value: Any):
        """Add metadata to the trace."""
        if self._span:
            self._span.set_attribute(key, value)

    def add_span(self, span: Span) -> None:
        """Add a child span to this trace."""
        self._child_spans.append(span)

    @property
    def trace_id(self) -> str:
        return self._span.trace_id if self._span else ""
