"""OpenTelemetry integration for automatic I/O capture.

Integrates with OpenTelemetry to automatically capture HTTP,
database, and Redis operations with minimal code changes.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

# Check for OpenTelemetry availability
try:
    from opentelemetry import trace
    from opentelemetry.sdk.trace import TracerProvider, ReadableSpan
    from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult

    OTEL_AVAILABLE = True
except ImportError:
    OTEL_AVAILABLE = False
    trace = None
    TracerProvider = None
    SpanExporter = object
    SpanExportResult = None
    ReadableSpan = None


@dataclass
class CapturedSpan:
    """A captured span from OpenTelemetry."""

    name: str
    kind: str
    start_time: float  # Unix timestamp in seconds
    end_time: float  # Unix timestamp in seconds
    duration_ms: float
    attributes: Dict[str, Any] = field(default_factory=dict)
    status: str = "UNSET"
    correlation_id: Optional[str] = None
    trace_id: Optional[str] = None
    span_id: Optional[str] = None
    parent_span_id: Optional[str] = None

    def __str__(self) -> str:
        return f"Span({self.name}, {self.duration_ms:.1f}ms, {self.status})"

    def is_http(self) -> bool:
        """Check if this is an HTTP span."""
        return (
            "http" in self.name.lower()
            or self.attributes.get("http.method") is not None
        )

    def is_database(self) -> bool:
        """Check if this is a database span."""
        return (
            any(k in self.name.lower() for k in ["sql", "query", "postgres", "mysql"])
            or self.attributes.get("db.system") is not None
        )

    def is_redis(self) -> bool:
        """Check if this is a Redis span."""
        return (
            "redis" in self.name.lower()
            or self.attributes.get("db.system") == "redis"
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "name": self.name,
            "kind": self.kind,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": self.duration_ms,
            "attributes": self.attributes,
            "status": self.status,
            "correlation_id": self.correlation_id,
            "trace_id": self.trace_id,
            "span_id": self.span_id,
        }


class ProbeSpanExporter(SpanExporter):
    """Custom exporter that feeds spans to the probe.

    This exporter captures OpenTelemetry spans and converts them
    to CapturedSpan objects for the probe.
    """

    def __init__(self, on_span: Optional[Callable[[CapturedSpan], None]] = None):
        """Initialize the exporter.

        Args:
            on_span: Callback to call for each captured span
        """
        self.on_span = on_span or (lambda s: None)
        self.captured_spans: List[CapturedSpan] = []

    def export(self, spans) -> "SpanExportResult":
        """Export spans to the probe.

        Args:
            spans: List of OpenTelemetry spans

        Returns:
            SpanExportResult.SUCCESS
        """
        from autodebug_probe.correlation import get_correlation_id

        for span in spans:
            captured = self._convert_span(span, get_correlation_id())
            self.captured_spans.append(captured)
            self.on_span(captured)

        if SpanExportResult:
            return SpanExportResult.SUCCESS
        return None

    def _convert_span(
        self,
        span: "ReadableSpan",
        correlation_id: str,
    ) -> CapturedSpan:
        """Convert OpenTelemetry span to CapturedSpan.

        Args:
            span: OpenTelemetry ReadableSpan
            correlation_id: Current correlation ID

        Returns:
            CapturedSpan instance
        """
        # Convert timestamps (nanoseconds to seconds)
        start_time = span.start_time / 1e9 if span.start_time else 0
        end_time = span.end_time / 1e9 if span.end_time else 0
        duration_ms = (span.end_time - span.start_time) / 1e6 if span.end_time and span.start_time else 0

        # Get span context
        trace_id = None
        span_id = None
        parent_span_id = None

        if span.context:
            trace_id = format(span.context.trace_id, "032x") if span.context.trace_id else None
            span_id = format(span.context.span_id, "016x") if span.context.span_id else None

        if span.parent:
            parent_span_id = format(span.parent.span_id, "016x") if span.parent.span_id else None

        # Get status
        status = "UNSET"
        if span.status:
            status = span.status.status_code.name if hasattr(span.status.status_code, "name") else str(span.status.status_code)

        return CapturedSpan(
            name=span.name,
            kind=span.kind.name if hasattr(span.kind, "name") else str(span.kind),
            start_time=start_time,
            end_time=end_time,
            duration_ms=duration_ms,
            attributes=dict(span.attributes) if span.attributes else {},
            status=status,
            correlation_id=correlation_id,
            trace_id=trace_id,
            span_id=span_id,
            parent_span_id=parent_span_id,
        )

    def shutdown(self) -> None:
        """Shutdown the exporter."""
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush any pending spans.

        Args:
            timeout_millis: Timeout in milliseconds

        Returns:
            True (always succeeds for in-memory exporter)
        """
        return True


class OpenTelemetryCapture:
    """Manages OpenTelemetry auto-instrumentation for the probe.

    Initializes OpenTelemetry with auto-instrumentation for common
    libraries and captures spans for the probe.
    """

    def __init__(self, on_span: Optional[Callable[[CapturedSpan], None]] = None):
        """Initialize the OpenTelemetry capture.

        Args:
            on_span: Optional callback for each captured span
        """
        self.on_span = on_span or (lambda s: None)
        self.exporter: Optional[ProbeSpanExporter] = None
        self._initialized = False
        self._instrumented_libraries: List[str] = []

    @property
    def available(self) -> bool:
        """Check if OpenTelemetry is available."""
        return OTEL_AVAILABLE

    def initialize(self) -> bool:
        """Initialize OpenTelemetry with auto-instrumentation.

        Returns:
            True if successfully initialized, False if OpenTelemetry unavailable
        """
        if not OTEL_AVAILABLE:
            return False

        if self._initialized:
            return True

        # Create custom exporter
        self.exporter = ProbeSpanExporter(self.on_span)

        # Set up tracer provider with our exporter
        from opentelemetry.sdk.trace.export import SimpleSpanProcessor

        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(self.exporter))
        trace.set_tracer_provider(provider)

        # Auto-instrument common libraries
        self._instrument_libraries()

        self._initialized = True
        return True

    def _instrument_libraries(self) -> None:
        """Auto-instrument common HTTP/DB libraries."""
        # HTTP clients
        self._try_instrument("opentelemetry.instrumentation.requests", "RequestsInstrumentor")
        self._try_instrument("opentelemetry.instrumentation.httpx", "HTTPXClientInstrumentor")
        self._try_instrument("opentelemetry.instrumentation.aiohttp_client", "AioHttpClientInstrumentor")

        # Databases (sync)
        self._try_instrument("opentelemetry.instrumentation.sqlalchemy", "SQLAlchemyInstrumentor")
        self._try_instrument("opentelemetry.instrumentation.psycopg2", "Psycopg2Instrumentor")
        self._try_instrument("opentelemetry.instrumentation.pymysql", "PyMySQLInstrumentor")

        # Databases (async) - Key for modern Python async support
        self._try_instrument("opentelemetry.instrumentation.asyncpg", "AsyncPGInstrumentor")

        # Redis (supports both sync and async)
        self._try_instrument("opentelemetry.instrumentation.redis", "RedisInstrumentor")

    def _try_instrument(self, module: str, instrumentor_name: str) -> bool:
        """Try to instrument a library.

        Args:
            module: Module containing the instrumentor
            instrumentor_name: Name of the instrumentor class

        Returns:
            True if successfully instrumented
        """
        # Validate module name to prevent arbitrary code execution via __import__
        import re
        if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_.]*$", module):
            return False
        if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", instrumentor_name):
            return False
        try:
            mod = __import__(module, fromlist=[instrumentor_name])
            instrumentor_cls = getattr(mod, instrumentor_name)
            instrumentor_cls().instrument()
            self._instrumented_libraries.append(instrumentor_name)
            return True
        except (ImportError, Exception):
            return False

    def get_all_spans(self) -> List[CapturedSpan]:
        """Get all captured spans.

        Returns:
            List of all captured spans
        """
        if not self.exporter:
            return []
        return list(self.exporter.captured_spans)

    def get_http_spans(self) -> List[CapturedSpan]:
        """Get captured HTTP spans.

        Returns:
            List of HTTP spans
        """
        return [s for s in self.get_all_spans() if s.is_http()]

    def get_db_spans(self) -> List[CapturedSpan]:
        """Get captured database spans.

        Returns:
            List of database spans
        """
        return [s for s in self.get_all_spans() if s.is_database()]

    def get_redis_spans(self) -> List[CapturedSpan]:
        """Get captured Redis spans.

        Returns:
            List of Redis spans
        """
        return [s for s in self.get_all_spans() if s.is_redis()]

    def get_spans_by_correlation_id(self, correlation_id: str) -> List[CapturedSpan]:
        """Get spans with a specific correlation ID.

        Args:
            correlation_id: Correlation ID to filter by

        Returns:
            List of matching spans
        """
        return [s for s in self.get_all_spans() if s.correlation_id == correlation_id]

    def clear_spans(self) -> None:
        """Clear all captured spans."""
        if self.exporter:
            self.exporter.captured_spans.clear()


# Singleton instance for convenience
_global_capture: Optional[OpenTelemetryCapture] = None


def get_global_capture() -> OpenTelemetryCapture:
    """Get or create the global OpenTelemetry capture.

    Returns:
        Global OpenTelemetryCapture instance
    """
    global _global_capture
    if _global_capture is None:
        _global_capture = OpenTelemetryCapture()
    return _global_capture


def initialize_otel() -> bool:
    """Initialize global OpenTelemetry capture.

    Returns:
        True if successfully initialized
    """
    return get_global_capture().initialize()


class HikaflowSpanExporter(SpanExporter):
    """Exports OpenTelemetry spans to Hikaflow transcript format.

    This exporter writes spans to JSONL files separated by type:
    - deps_sql.jsonl for database spans (asyncpg, psycopg2, etc.)
    - deps_redis.jsonl for Redis spans
    - deps_http.jsonl for HTTP spans (aiohttp, requests, etc.)

    This follows the OpenTelemetry instrumentation approach from Plan 51,
    particularly for async libraries (asyncpg, aiohttp, redis.asyncio).
    """

    def __init__(self, run_id: str, output_dir: Optional[str] = None):
        """Initialize the exporter.

        Args:
            run_id: The run ID for this capture session
            output_dir: Output directory (defaults to .hikaflow/runs/{run_id})
        """
        import os

        self.run_id = run_id
        self.output_dir = output_dir or f".hikaflow/runs/{run_id}"

        # Ensure output directory exists
        os.makedirs(self.output_dir, exist_ok=True)

        # Open file handles for each type
        self._sql_file = open(f"{self.output_dir}/deps_sql.jsonl", "a")
        self._redis_file = open(f"{self.output_dir}/deps_redis.jsonl", "a")
        self._http_file = open(f"{self.output_dir}/deps_http.jsonl", "a")

        self._closed = False

    def export(self, spans) -> "SpanExportResult":
        """Export spans to JSONL files.

        Args:
            spans: List of OpenTelemetry spans

        Returns:
            SpanExportResult.SUCCESS
        """
        import json

        for span in spans:
            record = self._span_to_record(span)
            if record is None:
                continue

            record_type = record.get("type", "")
            json_line = json.dumps(record) + "\n"

            if record_type == "SQL":
                self._sql_file.write(json_line)
            elif record_type == "REDIS":
                self._redis_file.write(json_line)
            elif record_type == "HTTP":
                self._http_file.write(json_line)

        # Flush after each batch
        self._sql_file.flush()
        self._redis_file.flush()
        self._http_file.flush()

        if SpanExportResult:
            return SpanExportResult.SUCCESS
        return None

    def _span_to_record(self, span: "ReadableSpan") -> Optional[Dict[str, Any]]:
        """Convert an OpenTelemetry span to a transcript record.

        Args:
            span: OpenTelemetry ReadableSpan

        Returns:
            Dict suitable for JSONL transcript, or None if not relevant
        """
        attrs = dict(span.attributes) if span.attributes else {}
        duration_ms = (span.end_time - span.start_time) / 1e6 if span.end_time and span.start_time else 0
        timestamp = span.start_time / 1e9 if span.start_time else 0

        # Extract trace context
        trace_id = None
        span_id = None
        if span.context:
            trace_id = format(span.context.trace_id, "032x") if span.context.trace_id else None
            span_id = format(span.context.span_id, "016x") if span.context.span_id else None

        # Determine span type and build record
        db_system = attrs.get("db.system", "")
        http_method = attrs.get("http.method") or attrs.get("http.request.method")

        if db_system in ("postgresql", "asyncpg", "psycopg2", "mysql", "sqlite"):
            # SQL database span
            return {
                "type": "SQL",
                "timestamp": timestamp,
                "query": attrs.get("db.statement", ""),
                "params": attrs.get("db.parameters", []),
                "duration_ms": duration_ms,
                "db_system": db_system,
                "db_name": attrs.get("db.name"),
                "trace_id": trace_id,
                "span_id": span_id,
            }

        elif db_system == "redis":
            # Redis span
            return {
                "type": "REDIS",
                "timestamp": timestamp,
                "command": attrs.get("db.statement", span.name),
                "duration_ms": duration_ms,
                "trace_id": trace_id,
                "span_id": span_id,
            }

        elif http_method:
            # HTTP span
            return {
                "type": "HTTP",
                "timestamp": timestamp,
                "method": http_method,
                "url": attrs.get("http.url") or attrs.get("url.full", ""),
                "status_code": attrs.get("http.status_code") or attrs.get("http.response.status_code"),
                "duration_ms": duration_ms,
                "trace_id": trace_id,
                "span_id": span_id,
            }

        # Not a span type we export
        return None

    def shutdown(self) -> None:
        """Shutdown the exporter and close files."""
        if self._closed:
            return

        self._sql_file.close()
        self._redis_file.close()
        self._http_file.close()
        self._closed = True

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush all files.

        Args:
            timeout_millis: Timeout in milliseconds (unused)

        Returns:
            True
        """
        if not self._closed:
            self._sql_file.flush()
            self._redis_file.flush()
            self._http_file.flush()
        return True


def install_async_instrumentation(run_id: str, output_dir: Optional[str] = None) -> bool:
    """Install OpenTelemetry instrumentation for async libraries.

    This is a convenience function that sets up OpenTelemetry with
    instrumentation for asyncpg, aiohttp, and redis.asyncio, exporting
    to Hikaflow transcript format.

    Args:
        run_id: The run ID for this capture session
        output_dir: Optional output directory

    Returns:
        True if successfully installed, False if OpenTelemetry unavailable
    """
    if not OTEL_AVAILABLE:
        return False

    from opentelemetry.sdk.trace.export import BatchSpanProcessor

    # Create provider with Hikaflow exporter
    provider = TracerProvider()
    exporter = HikaflowSpanExporter(run_id, output_dir)
    provider.add_span_processor(BatchSpanProcessor(exporter))
    trace.set_tracer_provider(provider)

    # Instrument async libraries
    instrumented = []

    # asyncpg - async PostgreSQL
    try:
        from opentelemetry.instrumentation.asyncpg import AsyncPGInstrumentor
        AsyncPGInstrumentor().instrument()
        instrumented.append("asyncpg")
    except ImportError:
        pass

    # aiohttp client
    try:
        from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
        AioHttpClientInstrumentor().instrument()
        instrumented.append("aiohttp")
    except ImportError:
        pass

    # Redis (includes redis.asyncio support)
    try:
        from opentelemetry.instrumentation.redis import RedisInstrumentor
        RedisInstrumentor().instrument()
        instrumented.append("redis")
    except ImportError:
        pass

    return len(instrumented) > 0
