"""
Foil OpenTelemetry Integration

Provides a SpanProcessor that exports OpenTelemetry spans to Foil's OTLP endpoint.
This enables automatic instrumentation of LLM calls using OpenLLMetry (traceloop-sdk)
or other OpenTelemetry instrumentations.

Usage:
    # Option 1: Use Foil.init() convenience method (recommended)
    from foil.otel import Foil
    Foil.init(api_key='sk_live_...', agent_name='my-agent')
    # All OpenAI/Anthropic calls are now auto-traced!

    # Option 2: Use FoilSpanProcessor directly
    from foil.otel import FoilSpanProcessor
    from opentelemetry.sdk.trace import TracerProvider

    provider = TracerProvider()
    provider.add_span_processor(FoilSpanProcessor(api_key='sk_live_...'))
    trace.set_tracer_provider(provider)
"""

import os
import threading
import time
import requests
from typing import Any, Dict, List, Optional, Sequence
from weakref import WeakKeyDictionary

# SDK version
__version__ = "0.5.0"

# Default Foil OTLP endpoint
DEFAULT_ENDPOINT = "https://api.getfoil.ai/api/otlp/v1/traces"


def _is_debug_enabled() -> bool:
    """Check if debug logging is enabled via environment variable."""
    return os.environ.get("FOIL_DEBUG", "").lower() == "true"


def _debug_log(message: str, data: Optional[Dict[str, Any]] = None) -> None:
    """Log debug message if debug is enabled."""
    if _is_debug_enabled():
        data_str = f" {data}" if data else ""
        print(f"[Foil] {message}{data_str}")


class FoilSpanProcessor:
    """
    FoilSpanProcessor implements the OpenTelemetry SpanProcessor interface.
    It batches completed spans and exports them to Foil's OTLP endpoint.
    """

    def __init__(
        self,
        api_key: str,
        endpoint: Optional[str] = None,
        max_batch_size: int = 100,
        scheduled_delay_ms: int = 5000,
        export_timeout_ms: int = 30000,
        debug: bool = False,
    ):
        """
        Create a new FoilSpanProcessor.

        Args:
            api_key: Foil API key (required)
            endpoint: Custom OTLP endpoint URL
            max_batch_size: Maximum spans per batch (default: 100)
            scheduled_delay_ms: Batch export interval in milliseconds (default: 5000)
            export_timeout_ms: Export request timeout in milliseconds (default: 30000)
            debug: Enable debug logging
        """
        if not api_key:
            raise ValueError('Foil API key is required. Pass api_key="sk_live_..."')

        self.api_key = api_key
        self.endpoint = endpoint or os.environ.get("FOIL_OTLP_ENDPOINT") or DEFAULT_ENDPOINT
        self.max_batch_size = max_batch_size
        self.scheduled_delay_ms = scheduled_delay_ms
        self.export_timeout_ms = export_timeout_ms
        self.debug = debug or _is_debug_enabled()

        # Pending spans waiting to be exported
        self._pending_spans: List[Dict[str, Any]] = []
        self._lock = threading.Lock()

        # Export timer
        self._timer: Optional[threading.Timer] = None
        self._shutdown = False

        # Start the export timer
        self._start_timer()

        if self.debug:
            print(
                f"[Foil] SpanProcessor initialized "
                f"(endpoint={self.endpoint}, maxBatchSize={self.max_batch_size}, "
                f"scheduledDelayMs={self.scheduled_delay_ms})"
            )

    def on_start(self, span: Any, parent_context: Any = None) -> None:
        """
        Called when a span is started. We don't need to do anything here
        since we only export completed spans.
        """
        pass

    def on_end(self, span: Any) -> None:
        """
        Called when a span is ended. We collect the span for batch export.
        """
        if self._shutdown:
            return

        # Convert OTEL span to OTLP format
        otlp_span = self._convert_to_otlp(span)

        with self._lock:
            self._pending_spans.append(otlp_span)
            pending_count = len(self._pending_spans)

        if self.debug:
            span_context = span.get_span_context()
            print(
                f"[Foil] Span ended (name={span.name}, "
                f"traceId={format(span_context.trace_id, '032x')}, "
                f"spanId={format(span_context.span_id, '016x')}, "
                f"pendingCount={pending_count})"
            )

        # Export immediately if batch is full
        if pending_count >= self.max_batch_size:
            self._export()

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

        Args:
            timeout_millis: Maximum time to wait (ignored in this implementation)

        Returns:
            True if successful
        """
        if self._pending_spans:
            self._export()
        return True

    def shutdown(self) -> None:
        """Shutdown the processor and flush any remaining spans."""
        self._shutdown = True

        # Stop the timer
        if self._timer:
            self._timer.cancel()
            self._timer = None

        # Flush remaining spans
        self.force_flush()

        if self.debug:
            print("[Foil] SpanProcessor shutdown complete")

    def _start_timer(self) -> None:
        """Start the batch export timer."""
        if self._shutdown:
            return

        self._timer = threading.Timer(
            self.scheduled_delay_ms / 1000.0,
            self._timer_callback,
        )
        self._timer.daemon = True  # Don't prevent process from exiting
        self._timer.start()

    def _timer_callback(self) -> None:
        """Timer callback to export pending spans."""
        if self._shutdown:
            return

        if self._pending_spans:
            self._export()

        # Restart timer
        self._start_timer()

    def _convert_to_otlp(self, span: Any) -> Dict[str, Any]:
        """
        Convert an OpenTelemetry ReadableSpan to OTLP JSON format.

        Args:
            span: The OTEL span

        Returns:
            OTLP span format dictionary
        """
        span_context = span.get_span_context()

        # Convert attributes to OTLP format
        attributes = []
        if hasattr(span, "attributes") and span.attributes:
            for key, value in span.attributes.items():
                attributes.append({
                    "key": key,
                    "value": self._convert_attribute_value(value),
                })

        # Convert events to OTLP format
        events = []
        if hasattr(span, "events") and span.events:
            for event in span.events:
                event_attrs = []
                if hasattr(event, "attributes") and event.attributes:
                    for key, value in event.attributes.items():
                        event_attrs.append({
                            "key": key,
                            "value": self._convert_attribute_value(value),
                        })
                events.append({
                    "timeUnixNano": str(event.timestamp),
                    "name": event.name,
                    "attributes": event_attrs,
                })

        # Convert links to OTLP format
        links = []
        if hasattr(span, "links") and span.links:
            for link in span.links:
                link_attrs = []
                if hasattr(link, "attributes") and link.attributes:
                    for key, value in link.attributes.items():
                        link_attrs.append({
                            "key": key,
                            "value": self._convert_attribute_value(value),
                        })
                link_context = link.context
                links.append({
                    "traceId": format(link_context.trace_id, "032x"),
                    "spanId": format(link_context.span_id, "016x"),
                    "traceState": str(link_context.trace_state) if link_context.trace_state else "",
                    "attributes": link_attrs,
                })

        # Get parent span ID
        parent_span_id = ""
        if hasattr(span, "parent") and span.parent:
            parent_span_id = format(span.parent.span_id, "016x")

        # Get status
        status_code = 0
        status_message = ""
        if hasattr(span, "status") and span.status:
            status_code = span.status.status_code.value if hasattr(span.status.status_code, "value") else int(span.status.status_code)
            status_message = span.status.description or ""

        return {
            "traceId": format(span_context.trace_id, "032x"),
            "spanId": format(span_context.span_id, "016x"),
            "parentSpanId": parent_span_id,
            "traceState": str(span_context.trace_state) if span_context.trace_state else "",
            "name": span.name,
            "kind": span.kind.value if hasattr(span.kind, "value") else int(span.kind),
            "startTimeUnixNano": str(span.start_time),
            "endTimeUnixNano": str(span.end_time),
            "attributes": attributes,
            "droppedAttributesCount": getattr(span, "dropped_attributes", 0),
            "events": events,
            "droppedEventsCount": getattr(span, "dropped_events", 0),
            "links": links,
            "droppedLinksCount": getattr(span, "dropped_links", 0),
            "status": {
                "code": status_code,
                "message": status_message,
            },
        }

    def _convert_attribute_value(self, value: Any) -> Dict[str, Any]:
        """
        Convert an attribute value to OTLP format.

        Args:
            value: The attribute value

        Returns:
            OTLP attribute value dictionary
        """
        if isinstance(value, str):
            return {"stringValue": value}
        if isinstance(value, bool):
            return {"boolValue": value}
        if isinstance(value, int):
            return {"intValue": str(value)}
        if isinstance(value, float):
            return {"doubleValue": value}
        if isinstance(value, (list, tuple)):
            return {
                "arrayValue": {
                    "values": [self._convert_attribute_value(v) for v in value]
                }
            }
        # Fallback: convert to string
        return {"stringValue": str(value)}

    def _export(self) -> None:
        """Export pending spans to Foil."""
        with self._lock:
            if not self._pending_spans:
                return

            # Take the current batch and clear pending
            batch = self._pending_spans[:self.max_batch_size]
            self._pending_spans = self._pending_spans[self.max_batch_size:]

        if self.debug:
            print(f"[Foil] Exporting spans (count={len(batch)})")

        # Build OTLP request
        request_body = {
            "resourceSpans": [
                {
                    "resource": {
                        "attributes": [
                            {
                                "key": "service.name",
                                "value": {"stringValue": os.environ.get("OTEL_SERVICE_NAME", "foil-instrumented-app")},
                            },
                            {
                                "key": "telemetry.sdk.name",
                                "value": {"stringValue": "foil-sdk"},
                            },
                            {
                                "key": "telemetry.sdk.language",
                                "value": {"stringValue": "python"},
                            },
                            {
                                "key": "telemetry.sdk.version",
                                "value": {"stringValue": __version__},
                            },
                        ],
                    },
                    "scopeSpans": [
                        {
                            "scope": {
                                "name": "foil-sdk",
                                "version": __version__,
                            },
                            "spans": batch,
                        },
                    ],
                },
            ],
        }

        try:
            response = requests.post(
                self.endpoint,
                json=request_body,
                headers={
                    "Content-Type": "application/json",
                    "Authorization": self.api_key,
                },
                timeout=self.export_timeout_ms / 1000.0,
            )

            if self.debug:
                print(f"[Foil] Export successful (count={len(batch)}, status={response.status_code})")

            # Check for partial success
            if response.status_code == 200:
                data = response.json()
                if data.get("partialSuccess", {}).get("rejectedSpans", 0) > 0:
                    print(f"[Foil] Warning: Some spans were rejected: {data['partialSuccess']}")

        except Exception as error:
            print(f"[Foil] Export failed: {error} (count={len(batch)})")

            # Re-queue spans on failure (up to a limit to prevent memory issues)
            with self._lock:
                if len(self._pending_spans) < self.max_batch_size * 3:
                    self._pending_spans = batch + self._pending_spans


class Foil:
    """
    Foil convenience class for initializing OpenTelemetry tracing.
    Provides a simple way to set up Foil as the trace exporter.
    """

    _provider: Any = None
    _processor: Optional[FoilSpanProcessor] = None

    @classmethod
    def init(
        cls,
        api_key: str,
        agent_name: Optional[str] = None,
        endpoint: Optional[str] = None,
        debug: bool = False,
    ) -> Dict[str, Any]:
        """
        Initialize Foil OpenTelemetry integration.
        This sets up a TracerProvider with FoilSpanProcessor and registers it globally.

        Note: This requires opentelemetry-sdk to be installed.

        Args:
            api_key: Foil API key (required)
            agent_name: Agent name for grouping traces in Foil
            endpoint: Custom OTLP endpoint URL
            debug: Enable debug logging

        Returns:
            Dict with provider and processor instances

        Example:
            from foil.otel import Foil

            # Initialize with API key
            Foil.init(api_key=os.environ['FOIL_API_KEY'], agent_name='my-agent')

            # Now any OpenTelemetry instrumentation will export to Foil
            # e.g., OpenLLMetry for LLM calls
        """
        if not api_key:
            raise ValueError('Foil API key is required. Pass api_key="sk_live_..."')

        debug = debug or _is_debug_enabled()

        # Load OpenTelemetry SDK
        try:
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider
        except ImportError:
            raise ImportError(
                "OpenTelemetry SDK is required. Install with: pip install foil-sdk[otel]"
            )

        # Set agent name as OTEL service name (used by transformer to create/lookup agent)
        if agent_name:
            os.environ["OTEL_SERVICE_NAME"] = agent_name

        # Create the processor
        cls._processor = FoilSpanProcessor(
            api_key=api_key,
            endpoint=endpoint,
            debug=debug,
        )

        # Create and configure the provider
        cls._provider = TracerProvider()
        cls._provider.add_span_processor(cls._processor)
        trace.set_tracer_provider(cls._provider)

        if debug:
            print("[Foil] Initialized OpenTelemetry provider")

        # Initialize OpenLLMetry (Traceloop) for automatic LLM instrumentation
        # We disable their exporter since we use FoilSpanProcessor instead
        try:
            from traceloop.sdk import Traceloop

            Traceloop.init(
                disable_batch=True,
                exporter=None,  # Disable Traceloop's exporter, we use FoilSpanProcessor
                api_key=None,  # Not sending to Traceloop
            )

            if debug:
                print(
                    "[Foil] Auto-instrumentation enabled for: "
                    "OpenAI, Anthropic, Cohere, Bedrock, etc."
                )
        except ImportError:
            if debug:
                print("[Foil] Note: Install traceloop-sdk for auto-instrumentation of LLM calls")

        if debug:
            print(
                f"[Foil] Ready! All LLM calls will be automatically traced. "
                f"(agentName={agent_name or os.environ.get('OTEL_SERVICE_NAME', 'default')})"
            )

        return {
            "provider": cls._provider,
            "processor": cls._processor,
        }

    @classmethod
    def shutdown(cls) -> None:
        """Shutdown Foil and flush any pending spans."""
        if cls._processor:
            cls._processor.shutdown()
            cls._processor = None
        if cls._provider:
            cls._provider.shutdown()
            cls._provider = None

    @classmethod
    def flush(cls) -> None:
        """Force flush any pending spans."""
        if cls._processor:
            cls._processor.force_flush()


# Attachment counter for unique indexing within a span
# Uses WeakKeyDictionary to track per-span attachment counts
_span_attachment_counts: WeakKeyDictionary = WeakKeyDictionary()


def add_attachment(
    type: str,
    value: str,
    role: str,
    name: Optional[str] = None,
    language: Optional[str] = None,
) -> bool:
    """
    Add an attachment to the current active span.
    Attachments allow you to include additional context like images, code, text,
    or embedded content.

    Args:
        type: Type of attachment: 'image', 'code', 'text', or 'iframe'
        value: The content or URL of the attachment
        role: Whether this is 'input' or 'output'
        name: Optional display name for the attachment
        language: Language for code attachments (e.g., 'python', 'javascript')

    Returns:
        True if attachment was added successfully, False otherwise

    Example:
        # Attach an input image
        add_attachment(type='image', value='https://example.com/image.png', role='input')

        # Attach generated code
        add_attachment(type='code', value='print("hello")', role='output', language='python')

        # Attach a document
        add_attachment(type='text', value='Long document content...', role='input', name='context.txt')
    """
    try:
        from opentelemetry import trace
    except ImportError:
        if _is_debug_enabled():
            print("[Foil] OpenTelemetry not installed, cannot add attachment")
        return False

    span = trace.get_current_span()

    if not span or not span.is_recording():
        if _is_debug_enabled():
            print("[Foil] add_attachment called but no active span found")
        return False

    # Validate attachment type
    valid_types = ["image", "code", "text", "iframe"]
    if type not in valid_types:
        print(f"[Foil] Invalid attachment type: {type}. Must be one of: {', '.join(valid_types)}")
        return False

    if not value:
        print("[Foil] Attachment value is required")
        return False

    if role not in ["input", "output"]:
        print('[Foil] Attachment role must be "input" or "output"')
        return False

    # Get or initialize attachment count for this span
    count = _span_attachment_counts.get(span, 0)

    # Set attributes on the span using Foil's attachment convention
    span.set_attribute(f"foil.attachment.{count}.type", type)
    span.set_attribute(f"foil.attachment.{count}.value", value)
    span.set_attribute(f"foil.attachment.{count}.role", role)

    if name:
        span.set_attribute(f"foil.attachment.{count}.name", name)

    if language and type == "code":
        span.set_attribute(f"foil.attachment.{count}.language", language)

    # Increment counter
    _span_attachment_counts[span] = count + 1

    if _is_debug_enabled():
        print(f"[Foil] Attachment added (index={count}, type={type}, role={role})")

    return True


def attach_image(url: str, role: str, name: Optional[str] = None) -> bool:
    """
    Convenience function to attach an image.

    Args:
        url: Image URL or base64 data URI
        role: 'input' or 'output'
        name: Optional display name

    Returns:
        True if attachment was added successfully
    """
    return add_attachment(type="image", value=url, role=role, name=name)


def attach_code(
    code: str,
    role: str,
    language: Optional[str] = None,
    name: Optional[str] = None,
) -> bool:
    """
    Convenience function to attach code.

    Args:
        code: Code content
        role: 'input' or 'output'
        language: Programming language (e.g., 'python', 'javascript')
        name: Optional display name

    Returns:
        True if attachment was added successfully
    """
    return add_attachment(type="code", value=code, role=role, language=language, name=name)


def attach_text(text: str, role: str, name: Optional[str] = None) -> bool:
    """
    Convenience function to attach text/document.

    Args:
        text: Text content
        role: 'input' or 'output'
        name: Optional display name

    Returns:
        True if attachment was added successfully
    """
    return add_attachment(type="text", value=text, role=role, name=name)
