import requests
import time
import uuid
import os
import mimetypes
from functools import wraps
from typing import Optional, Dict, Any, Generator, List, Union, BinaryIO
from dataclasses import dataclass, asdict
from enum import Enum


class SpanKind(str, Enum):
    """Span kinds supported by Foil tracing."""
    AGENT = "agent"
    LLM = "llm"
    TOOL = "tool"
    CHAIN = "chain"
    RETRIEVER = "retriever"
    EMBEDDING = "embedding"
    CUSTOM = "custom"


class MediaCategory(str, Enum):
    """Media categories for multimodal content."""
    IMAGE = "image"
    DOCUMENT = "document"
    SPREADSHEET = "spreadsheet"
    CODE = "code"
    AUDIO = "audio"
    VIDEO = "video"
    ARCHIVE = "archive"
    NOTEBOOK = "notebook"
    OTHER = "other"


@dataclass
class TextBlock:
    """Text content block for multimodal input/output."""
    type: str = "text"
    text: str = ""

    def to_dict(self) -> Dict[str, Any]:
        return {"type": self.type, "text": self.text}


@dataclass
class MediaBlock:
    """Media content block for multimodal input/output."""
    type: str = "media"
    media_id: str = ""
    category: Optional[str] = None
    mime_type: Optional[str] = None
    filename: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        result = {"type": self.type, "mediaId": self.media_id}
        if self.category:
            result["category"] = self.category
        if self.mime_type:
            result["mimeType"] = self.mime_type
        if self.filename:
            result["filename"] = self.filename
        return result


class ContentBlock:
    """Factory for creating content blocks."""

    @staticmethod
    def text(text: str) -> Dict[str, Any]:
        """
        Create a text content block.

        Args:
            text: The text content

        Returns:
            Text content block dictionary
        """
        return {"type": "text", "text": text}

    @staticmethod
    def media(
        media_id: str,
        category: Optional[str] = None,
        mime_type: Optional[str] = None,
        filename: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Create a media content block.

        Args:
            media_id: The media ID from upload_media
            category: Media category
            mime_type: MIME type
            filename: Original filename

        Returns:
            Media content block dictionary
        """
        result = {"type": "media", "mediaId": media_id}
        if category:
            result["category"] = category
        if mime_type:
            result["mimeType"] = mime_type
        if filename:
            result["filename"] = filename
        return result


def content(*parts: Union[str, Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Helper function to create content blocks array from mixed inputs.

    Args:
        *parts: Text strings or content block dictionaries

    Returns:
        Array of content blocks

    Example:
        >>> content("What is in this image?", ContentBlock.media("abc123"))
        [{"type": "text", "text": "What is in this image?"}, {"type": "media", "mediaId": "abc123"}]
    """
    result = []
    for part in parts:
        if isinstance(part, str):
            result.append(ContentBlock.text(part))
        elif isinstance(part, dict):
            result.append(part)
        elif hasattr(part, "to_dict"):
            result.append(part.to_dict())
        else:
            result.append(ContentBlock.text(str(part)))
    return result


class Foil:
    """Foil client for monitoring and logging AI model spans and traces."""

    def __init__(self, api_key: str, base_url: str = "https://api.getfoil.ai/api",
                 agent_name: Optional[str] = None, debug: bool = False):
        """
        Initialize the Foil client.

        Args:
            api_key: Your Foil API key
            base_url: Base URL for the Foil API (default: https://api.getfoil.ai/api)
            agent_name: Name of the agent (required for tracing)
            debug: Enable debug logging
        """
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.agent_name = agent_name
        self._debug = debug
        self._tracer = None

    def trace(self, fn, name=None, trace_id=None, session_id=None, input=None,
              properties=None, user_id=None, user_properties=None, device=None):
        """
        Execute a function within a traced context.
        Requires agent_name to be set on the Foil instance.

        Args:
            fn: Function to execute, receives TraceContext
            name: Name for the root span
            trace_id: Custom trace ID (auto-generated if not provided)
            session_id: Session ID for conversation tracking
            input: Input to record
            properties: Custom properties
            user_id: User identifier
            user_properties: Additional user attributes
            device: Device/platform information

        Returns:
            Result of the function
        """
        if not self.agent_name:
            raise ValueError("agent_name is required for tracing. Pass it to Foil(agent_name=...)")
        if self._tracer is None:
            from .tracer import FoilTracer
            self._tracer = FoilTracer(foil=self, agent_name=self.agent_name, debug=self._debug)
        return self._tracer.trace(
            fn, name=name, trace_id=trace_id, session_id=session_id,
            input=input, properties=properties, user_id=user_id,
            user_properties=user_properties, device=device,
        )

    def _headers(self) -> dict:
        """Get authorization headers."""
        return {"Authorization": self.api_key}

    def create_trace_id(self) -> str:
        """
        Generate a unique trace ID.

        Returns:
            A unique trace ID string
        """
        return str(uuid.uuid4())

    def create_span_id(self) -> str:
        """
        Generate a unique span ID.

        Returns:
            A unique span ID string
        """
        return str(uuid.uuid4())

    # Legacy aliases
    create_trace = create_trace_id
    create_event_id = create_span_id

    def log(self, data: dict) -> None:
        """
        Log a model call to Foil (fire-and-forget).

        Args:
            data: Log data containing:
                - model (required): Model name
                - input (required): Input to the model
                - output (optional): Output from the model
                - latency (optional): Latency in milliseconds
                - tokens (optional): Token usage {prompt, completion, total}
                - ttft (optional): Time to first token in milliseconds
                - status (optional): Status ('success' or 'error')
                - error (optional): Error message if status is 'error'
                - metadata (optional): Additional metadata
        """
        try:
            requests.post(
                f"{self.base_url}/logs",
                json=data,
                headers=self._headers(),
            )
        except Exception as e:
            print(f"Foil Logging Failed: {e}")

    @staticmethod
    def _strip_none(data: dict) -> dict:
        """Recursively remove keys with None values to avoid Zod validation errors (optional != nullable)."""
        result = {}
        for k, v in data.items():
            if v is None:
                continue
            if isinstance(v, dict):
                v = Foil._strip_none(v)
            result[k] = v
        return result

    def start_span(self, data: dict) -> dict:
        """
        Start a span (signals the beginning of an AI model operation).

        DEPRECATED: The SDK should track spans locally and send complete data
        to end_span() instead. This method is kept for backward compatibility
        but the server no longer stores span data in MongoDB on start.

        Args:
            data: Span data containing:
                - spanId (required): Unique identifier for this span
                - name (required): Span name/type
                - agentName (required): Name of the agent
                - input (optional): Input to the model
                - model (optional): Model name/identifier
                - sessionId (optional): Session identifier for grouping related spans
                - properties (optional): Additional custom properties
                - traceId (optional): Trace ID for grouping related spans
                - parentSpanId (optional): Parent span ID for nested spans
                - spanKind (optional): Span kind ('agent', 'llm', 'tool', 'chain', 'retriever', 'embedding', 'custom')
                - experimentId (optional): Experiment identifier
                - variantId (optional): Variant identifier

        Returns:
            The created span object

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.post(
                f"{self.base_url}/traces/spans/start",
                json=self._strip_none(data),
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Start Span Failed: {e}")
            raise

    def end_span(self, data: dict) -> dict:
        """
        End a span (signals the completion of an AI model operation).

        The SDK should now send complete span data here (both start and end data),
        as the server no longer stores start data in MongoDB.

        Args:
            data: Complete span data containing:
                - spanId (required): The spanId of the span
                - agentName (required): Name of the agent

                Start data (previously sent to start_span, now send here):
                - name (optional): Span name/type
                - traceId (optional): Trace ID for grouping related spans
                - parentSpanId (optional): Parent span ID for nested spans
                - spanKind (optional): Span kind
                - input (optional): Input to the model
                - model (optional): Model name/identifier
                - sessionId (optional): Session identifier
                - properties (optional): Additional custom properties
                - llmConfig (optional): LLM configuration for replay
                - device (optional): Device/platform context
                - endUserId (optional): End user identifier
                - endUserProperties (optional): End user attributes
                - experimentId (optional): Experiment identifier
                - variantId (optional): Variant identifier
                - inputAttachments (optional): Input attachments array

                End data:
                - output (optional): Output from the model call
                - tokens (optional): Token usage {prompt, completion, total}
                - timing (optional): Timing metrics {ttft, totalDuration}
                - error (optional): Error info {type, message} if span failed
                - attachments (optional): Output attachments array
                - startTime (optional): ISO string of when span started
                - endTime (optional): ISO string of when span ended

        Returns:
            The span object

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.post(
                f"{self.base_url}/traces/spans/end",
                json=self._strip_none(data),
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil End Span Failed: {e}")
            raise

    # Legacy aliases for backward compatibility
    def start_invocation(self, data: dict) -> dict:
        """
        DEPRECATED: Use start_span() instead.
        Start an invocation (signals the beginning of an AI model call).
        """
        # Convert legacy field names to new format
        span_data = {
            "spanId": data.get("eventId") or data.get("spanId"),
            "name": data.get("event") or data.get("name"),
            "agentName": data.get("agentName", "default"),
            "input": data.get("inputMessage") or data.get("input"),
            "model": data.get("model"),
            "sessionId": data.get("convoId") or data.get("sessionId"),
            "properties": data.get("properties"),
            "traceId": data.get("traceId"),
            "parentSpanId": data.get("parentEventId") or data.get("parentSpanId"),
            "spanKind": data.get("spanKind"),
        }
        # Remove None values
        span_data = {k: v for k, v in span_data.items() if v is not None}
        return self.start_span(span_data)

    def end_invocation(self, data: dict) -> dict:
        """
        DEPRECATED: Use end_span() instead.
        End an invocation (signals the completion of an AI model call).
        """
        # Convert legacy field names to new format
        span_data = {
            "spanId": data.get("eventId") or data.get("spanId"),
            "output": data.get("output"),
            "tokens": data.get("tokens"),
            "timing": data.get("timing"),
            "model": data.get("model"),
            "error": data.get("error"),
        }
        # Remove None values
        span_data = {k: v for k, v in span_data.items() if v is not None}
        return self.end_span(span_data)

    def get_trace(self, trace_id: str) -> dict:
        """
        Get a trace with all its spans.

        Args:
            trace_id: The trace ID

        Returns:
            The trace data with spans and tree structure

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.get(
                f"{self.base_url}/traces/{trace_id}",
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Get Trace Failed: {e}")
            raise

    def list_traces(
        self,
        agent_id: Optional[str] = None,
        agent_name: Optional[str] = None,
        status: Optional[str] = None,
        model: Optional[str] = None,
        search: Optional[str] = None,
        from_date: Optional[str] = None,
        to_date: Optional[str] = None,
        page: int = 1,
        limit: int = 20,
        sort: str = "-startTime",
    ) -> dict:
        """
        List traces with pagination and filters.

        Args:
            agent_id: Filter by agent ID
            agent_name: Filter by agent name
            status: Filter by status ('completed', 'error', 'started')
            model: Filter by model name
            search: Search in span names and outputs
            from_date: Start date filter (ISO format)
            to_date: End date filter (ISO format)
            page: Page number (default: 1)
            limit: Items per page (default: 20, max: 100)
            sort: Sort field and direction (default: '-startTime')

        Returns:
            Paginated list of traces

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        params = {
            "page": page,
            "limit": min(limit, 100),
            "sort": sort,
        }
        if agent_id:
            params["agentId"] = agent_id
        if agent_name:
            params["agentName"] = agent_name
        if status:
            params["status"] = status
        if model:
            params["model"] = model
        if search:
            params["search"] = search
        if from_date:
            params["from"] = from_date
        if to_date:
            params["to"] = to_date

        try:
            response = requests.get(
                f"{self.base_url}/traces",
                params=params,
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil List Traces Failed: {e}")
            raise

    def wrap_openai(self, client, agent_name: str = "openai"):
        """
        Wrap an OpenAI client to automatically log all chat completion calls.
        Supports both streaming and non-streaming responses.

        Args:
            client: An OpenAI client instance
            agent_name: Name of the agent for span tracking (default: "openai")

        Returns:
            The wrapped OpenAI client with automatic logging
        """
        original_create = client.chat.completions.create
        guard = self

        @wraps(original_create)
        def wrapped_create(*args, **kwargs):
            start_time = time.time()
            response = None
            error = None
            status = "success"
            ttft = None
            is_streaming = kwargs.get("stream", False)

            try:
                response = original_create(*args, **kwargs)

                # Handle streaming response
                if is_streaming:
                    return guard._wrap_stream(
                        response, start_time, kwargs, args, agent_name
                    )

                return response
            except Exception as e:
                error = str(e)
                status = "error"
                raise
            finally:
                # Only log for non-streaming responses here
                # Streaming responses are logged after iteration completes
                if not is_streaming:
                    end_time = time.time()
                    latency = int((end_time - start_time) * 1000)

                    input_messages = kwargs.get("messages", args[0] if args else None)
                    output = None
                    tokens = None
                    tool_calls = None

                    if response:
                        choice = response.choices[0] if response.choices else None
                        if choice and choice.message:
                            output = choice.message.content
                            tool_calls = choice.message.tool_calls

                        if response.usage:
                            tokens = {
                                "prompt": response.usage.prompt_tokens,
                                "completion": response.usage.completion_tokens,
                                "total": response.usage.total_tokens,
                            }

                    guard.log({
                        "model": kwargs.get("model", "openai-unknown"),
                        "input": input_messages,
                        "output": output,
                        "latency": latency,
                        "tokens": tokens,
                        "status": status,
                        "error": error,
                        "metadata": {
                            "streaming": False,
                            "tool_calls": [
                                {"id": tc.id, "type": tc.type, "function": {"name": tc.function.name, "arguments": tc.function.arguments}}
                                for tc in tool_calls
                            ] if tool_calls else None,
                        },
                    })

        client.chat.completions.create = wrapped_create
        return client

    def _wrap_stream(self, stream, start_time: float, kwargs: dict, args: tuple, agent_name: str = "openai"):
        """
        Wrap a streaming response to capture TTFT and log after completion.

        Args:
            stream: The streaming response
            start_time: When the request started
            kwargs: Original kwargs passed to create
            args: Original args passed to create
            agent_name: Name of the agent for span tracking

        Returns:
            A generator that yields chunks and logs on completion
        """
        guard = self
        first_chunk = True
        ttft = None
        collected_content = ""
        collected_tool_calls = []
        error = None
        status = "success"

        def stream_wrapper():
            nonlocal first_chunk, ttft, collected_content, collected_tool_calls, error, status

            try:
                for chunk in stream:
                    if first_chunk:
                        ttft = int((time.time() - start_time) * 1000)
                        first_chunk = False

                    # Collect content from chunks
                    if chunk.choices and chunk.choices[0].delta:
                        delta = chunk.choices[0].delta
                        if delta.content:
                            collected_content += delta.content

                        # Collect tool calls from chunks
                        if delta.tool_calls:
                            for tc in delta.tool_calls:
                                while len(collected_tool_calls) <= tc.index:
                                    collected_tool_calls.append({
                                        "id": None,
                                        "type": None,
                                        "function": {"name": "", "arguments": ""},
                                    })
                                if tc.id:
                                    collected_tool_calls[tc.index]["id"] = tc.id
                                if tc.type:
                                    collected_tool_calls[tc.index]["type"] = tc.type
                                if tc.function:
                                    if tc.function.name:
                                        collected_tool_calls[tc.index]["function"]["name"] += tc.function.name
                                    if tc.function.arguments:
                                        collected_tool_calls[tc.index]["function"]["arguments"] += tc.function.arguments

                    yield chunk

            except Exception as e:
                error = str(e)
                status = "error"
                raise
            finally:
                # Log after streaming completes
                end_time = time.time()
                latency = int((end_time - start_time) * 1000)

                input_messages = kwargs.get("messages", args[0] if args else None)

                guard.log({
                    "model": kwargs.get("model", "openai-unknown"),
                    "input": input_messages,
                    "output": collected_content or None,
                    "latency": latency,
                    "ttft": ttft,
                    "status": status,
                    "error": error,
                    "metadata": {
                        "streaming": True,
                        "tool_calls": collected_tool_calls if collected_tool_calls else None,
                    },
                })

        return stream_wrapper()

    def upload_media(
        self,
        file: Union[str, bytes, BinaryIO],
        filename: Optional[str] = None,
        mime_type: Optional[str] = None,
        span_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        direction: Optional[str] = None,
    ) -> dict:
        """
        Upload a media file for multimodal spans.

        Args:
            file: File path string, bytes content, or file-like object
            filename: Override filename (required for bytes)
            mime_type: Override MIME type (auto-detected if not provided)
            span_id: Associate with a span
            trace_id: Associate with a trace
            direction: 'input' or 'output'

        Returns:
            Upload result with mediaId and category

        Raises:
            requests.exceptions.RequestException: If the upload fails
            ValueError: If filename is not provided for bytes input

        Example:
            >>> foil = Foil(api_key="sk_live_...")
            >>> result = foil.upload_media("/path/to/image.png")
            >>> print(result["mediaId"])
        """
        try:
            files = {}
            data = {}

            # Handle different input types
            if isinstance(file, str):
                # File path
                actual_filename = filename or os.path.basename(file)
                actual_mime_type = mime_type or mimetypes.guess_type(file)[0] or "application/octet-stream"
                files["file"] = (actual_filename, open(file, "rb"), actual_mime_type)
            elif isinstance(file, bytes):
                # Bytes - filename is required
                if not filename:
                    raise ValueError("filename is required when uploading bytes")
                actual_mime_type = mime_type or mimetypes.guess_type(filename)[0] or "application/octet-stream"
                files["file"] = (filename, file, actual_mime_type)
            elif hasattr(file, "read"):
                # File-like object
                actual_filename = filename or getattr(file, "name", "upload")
                if isinstance(actual_filename, str) and os.sep in actual_filename:
                    actual_filename = os.path.basename(actual_filename)
                actual_mime_type = mime_type or mimetypes.guess_type(actual_filename)[0] or "application/octet-stream"
                files["file"] = (actual_filename, file, actual_mime_type)
            else:
                raise ValueError("file must be a path string, bytes, or file-like object")

            # Add optional fields
            if span_id:
                data["spanId"] = span_id
            if trace_id:
                data["traceId"] = trace_id
            if direction:
                data["direction"] = direction

            response = requests.post(
                f"{self.base_url}/media/upload",
                files=files,
                data=data if data else None,
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Upload Media Failed: {e}")
            raise

    def get_media(
        self,
        media_id: str,
        content: Optional[str] = None,
    ) -> dict:
        """
        Get media information by ID.

        Args:
            media_id: The media ID
            content: Content type to get URL for ('original', 'extracted', 'preview', 'structured')

        Returns:
            Media information

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            params = {}
            if content:
                params["content"] = content

            response = requests.get(
                f"{self.base_url}/media/{media_id}",
                params=params if params else None,
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Get Media Failed: {e}")
            raise

    def get_media_url(
        self,
        media_id: str,
        content: str = "original",
    ) -> dict:
        """
        Get presigned download URL for media content.

        Args:
            media_id: The media ID
            content: Content type ('original', 'extracted', 'preview', 'structured')

        Returns:
            Presigned URL info with 'url' and 'expiresAt'

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.get(
                f"{self.base_url}/media/{media_id}/url",
                params={"content": content},
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Get Media URL Failed: {e}")
            raise

    def batch_media_info(self, media_ids: List[str]) -> dict:
        """
        Get information for multiple media IDs.

        Args:
            media_ids: Array of media IDs

        Returns:
            Batch media info

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.post(
                f"{self.base_url}/media/batch",
                json={"mediaIds": media_ids},
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Batch Media Info Failed: {e}")
            raise

    def record_signal(self, data: dict) -> dict:
        """
        Record a signal (user feedback, LLM analysis, etc.).

        Args:
            data: Signal data containing:
                - signalName (required): Name of the signal (e.g., 'user_thumbs', 'sentiment')
                - value (required): Signal value (boolean, number, or string)
                - traceId (optional): Associated trace ID
                - spanId (optional): Associated span ID
                - sessionId (optional): Associated session ID
                - agentId (optional): Agent ID
                - agentName (optional): Agent name
                - signalType (optional): Type (feedback, sentiment, quality, completion, safety, effort, engagement, custom)
                - source (optional): Source (user, llm, system)
                - valueType (optional): Explicit value type (boolean, number, category, text)
                - metadata (optional): Additional metadata
                - confidence (optional): Confidence score 0-1 (for LLM signals)
                - reasoning (optional): LLM reasoning (for LLM signals)
                - modelUsed (optional): Model used (for LLM signals)

        Returns:
            The recorded signal

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.post(
                f"{self.base_url}/signals",
                json=self._strip_none(data),
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Record Signal Failed: {e}")
            raise

    def record_signal_batch(self, signals: List[dict]) -> dict:
        """
        Record multiple signals in batch.

        Args:
            signals: Array of signal objects (same format as record_signal)

        Returns:
            The recorded signals

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.post(
                f"{self.base_url}/signals/batch",
                json={"signals": signals},
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Record Signal Batch Failed: {e}")
            raise

    def get_trace_signals(self, trace_id: str) -> dict:
        """
        Get signals for a trace.

        Args:
            trace_id: The trace ID

        Returns:
            Signals associated with the trace

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            response = requests.get(
                f"{self.base_url}/signals/trace/{trace_id}",
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Get Trace Signals Failed: {e}")
            raise

    def get_experiment_variant(self, experiment_id: str, identifier: str) -> dict:
        """
        Get variant assignment for an experiment.
        Uses deterministic hashing so the same identifier always gets the same variant.

        Args:
            experiment_id: The experiment ID
            identifier: Unique identifier (userId, sessionId, etc.) for consistent assignment

        Returns:
            Assignment result with:
                - inExperiment: Whether the identifier is in the experiment
                - experimentId: The experiment ID
                - experimentName: The experiment name
                - variantId: Assigned variant ID
                - variantName: Assigned variant name
                - config: Variant configuration (model, systemPrompt, temperature, etc.)
                - message: Message if excluded from experiment

        Raises:
            ValueError: If experimentId or identifier is not provided
            requests.exceptions.RequestException: If the API call fails
        """
        if not experiment_id:
            raise ValueError("experiment_id is required")
        if not identifier:
            raise ValueError("identifier is required")

        try:
            response = requests.get(
                f"{self.base_url}/experiments/{experiment_id}/assign",
                params={"identifier": identifier},
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Get Experiment Variant Failed: {e}")
            raise

    # ============================================
    # Semantic Search Methods
    # ============================================

    def semantic_search(
        self,
        query: str,
        agent_id: Optional[str] = None,
        agent_name: Optional[str] = None,
        from_date: Optional[str] = None,
        to_date: Optional[str] = None,
        limit: int = 20,
        offset: int = 0,
        threshold: float = 0.3,
    ) -> dict:
        """
        Search spans using natural language queries.
        Uses AI embeddings to find semantically similar content.

        Args:
            query: Natural language search query (e.g., "conversations about refunds")
            agent_id: Filter by specific agent ID
            agent_name: Filter by agent name
            from_date: Start date filter (ISO format)
            to_date: End date filter (ISO format)
            limit: Maximum results to return (default: 20)
            offset: Offset for pagination (default: 0)
            threshold: Minimum similarity threshold 0-1 (default: 0.3)

        Returns:
            Search results with:
                - results: Array of matching spans with similarity scores
                - pagination: Pagination info (total, limit, offset, hasMore)
                - query: Query info (text, threshold, embedding status)

        Raises:
            ValueError: If query is not provided
            requests.exceptions.RequestException: If the API call fails

        Example:
            >>> results = foil.semantic_search(
            ...     "conversations about refunds",
            ...     agent_id="agent-123",
            ...     limit=10,
            ...     threshold=0.4
            ... )
        """
        if not query:
            raise ValueError("query is required")

        try:
            filters = {}
            if agent_id:
                filters["agentId"] = agent_id
            if agent_name:
                filters["agentName"] = agent_name
            if from_date:
                filters["from"] = from_date
            if to_date:
                filters["to"] = to_date

            response = requests.post(
                f"{self.base_url}/semantic-search/query",
                json={
                    "query": query,
                    "filters": filters if filters else None,
                    "limit": limit,
                    "offset": offset,
                    "threshold": threshold,
                },
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Semantic Search Failed: {e}")
            raise

    def find_similar_traces(
        self,
        trace_id: str,
        limit: int = 10,
        threshold: float = 0.3,
    ) -> dict:
        """
        Find traces similar to a given trace.
        Uses the trace's content embeddings to find related conversations.

        Args:
            trace_id: The trace ID to find similar traces for
            limit: Maximum results to return (default: 10)
            threshold: Minimum similarity threshold 0-1 (default: 0.3)

        Returns:
            Similar traces with:
                - results: Array of similar traces with similarity scores
                - sourceTrace: Info about the source trace

        Raises:
            ValueError: If trace_id is not provided
            requests.exceptions.RequestException: If the API call fails

        Example:
            >>> similar = foil.find_similar_traces(
            ...     "trace-abc-123",
            ...     limit=5,
            ...     threshold=0.5
            ... )
        """
        if not trace_id:
            raise ValueError("trace_id is required")

        try:
            response = requests.get(
                f"{self.base_url}/semantic-search/similar/{trace_id}",
                params={"limit": limit, "threshold": threshold},
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Find Similar Traces Failed: {e}")
            raise

    def get_semantic_search_status(self, agent_id: Optional[str] = None) -> dict:
        """
        Get semantic search status and embedding statistics.
        Shows how many spans have been embedded and are searchable.

        Args:
            agent_id: Optional agent ID to get stats for specific agent

        Returns:
            Embedding status with:
                - embeddedSpans: Number of spans with embeddings
                - totalSpans: Total eligible spans
                - coveragePercent: Percentage of spans embedded
                - ready: Whether semantic search is ready

        Raises:
            requests.exceptions.RequestException: If the API call fails
        """
        try:
            params = {}
            if agent_id:
                params["agentId"] = agent_id

            response = requests.get(
                f"{self.base_url}/semantic-search/status",
                params=params if params else None,
                headers=self._headers(),
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"Foil Get Semantic Search Status Failed: {e}")
            raise

