"""
Foil Tracer - High-level tracing for AI applications.

This module provides the FoilTracer and TraceContext classes for comprehensive
tracing of AI operations including LLM calls, tool executions, and more.
"""

import os
import time
import uuid
from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
from dataclasses import dataclass, field

from .client import Foil, SpanKind


# Debug logging - enable with FOIL_DEBUG=true or pass debug=True to tracer
def _is_debug_enabled() -> bool:
    return os.environ.get("FOIL_DEBUG", "").lower() == "true"


T = TypeVar("T")


@dataclass
class Span:
    """Represents an active span."""
    span_id: str
    span_kind: str
    name: str
    start_time: float
    depth: int
    attributes: Dict[str, Any] = field(default_factory=dict)
    _context: "TraceContext" = field(default=None, repr=False)

    def end(self, result: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        End this span.

        Args:
            result: Span result containing:
                - output: Output from the operation
                - tokens: Token usage {prompt, completion, total}
                - error: Error if operation failed {type, message}
                - ttft: Time to first token (milliseconds)

        Returns:
            Dict with spanId and duration
        """
        if result is None:
            result = {}

        end_time = time.time()
        duration = int((end_time - self.start_time) * 1000)

        end_data = {
            "spanId": self.span_id,
            "traceId": self._context.trace_id,
            "agentName": self._context.agent_name,
            "output": result.get("output"),
            "tokens": result.get("tokens"),
            "timing": {
                "totalDuration": duration,
                "ttft": result.get("ttft"),
            },
            "error": result.get("error"),
        }

        self._context._active_spans.pop(self.span_id, None)

        # Pop this span from the stack
        if self.span_id in self._context._span_stack:
            self._context._span_stack.remove(self.span_id)

        log_data = {"spanId": self.span_id[:8], "duration": f"{duration}ms"}
        if result.get("tokens", {}).get("total"):
            log_data["tokens"] = result["tokens"]["total"]
        if result.get("error"):
            log_data["error"] = result["error"].get("message") or result["error"]
        self._context._log(f"■ END [{self.span_kind}] {self.name}", log_data)

        # Fire and forget
        try:
            self._context._tracer._foil.end_span(end_data)
        except Exception as err:
            print(f"Foil span end failed: {err}")

        return {"spanId": self.span_id, "duration": duration}

    def create_child_context(self) -> "TraceContext":
        """Create a child context for this span."""
        return TraceContext(
            tracer=self._context._tracer,
            trace_id=self._context.trace_id,
            root_span_id=self.span_id,
            agent_name=self._context.agent_name,
            depth=self._context.depth + 1,
        )


class TraceContext:
    """
    TraceContext provides span management within a trace.
    Passed to traced functions to allow creating child spans.
    """

    def __init__(
        self,
        tracer: "FoilTracer",
        trace_id: str,
        root_span_id: str,
        agent_name: str,
        depth: int = 0,
    ):
        self._tracer = tracer
        self.trace_id = trace_id
        self.root_span_id = root_span_id
        self.agent_name = agent_name
        self.depth = depth
        self._active_spans: Dict[str, Span] = {}
        self._debug = tracer._debug or _is_debug_enabled()
        # Stack to track current active span for proper parent-child relationships
        self._span_stack: List[str] = [root_span_id]

    @property
    def current_parent_span_id(self) -> str:
        """Get the current parent span ID (the most recent active span)."""
        return self._span_stack[-1] if self._span_stack else self.root_span_id

    def _log(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
        """Log debug message with indentation based on depth."""
        if not self._debug:
            return
        indent = "  " * self.depth
        data_str = f" {data}" if data else ""
        print(f"[Foil] {indent}{message}{data_str}")

    def start_span(
        self,
        span_kind: Union[SpanKind, str],
        name: str,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> Span:
        """
        Start a new child span.

        Args:
            span_kind: Type of span (llm, tool, chain, etc.)
            name: Name of the span (model name, tool name, etc.)
            attributes: Additional span attributes including:
                - input: Input data
                - properties: Custom properties
                - parent_span_id: Explicit parent span ID (overrides automatic stack-based parent)

        Returns:
            Span object with spanId and methods to end the span
        """
        if attributes is None:
            attributes = {}

        span_id = str(uuid.uuid4())
        start_time = time.time()

        # Convert SpanKind enum to string if needed
        span_kind_str = span_kind.value if isinstance(span_kind, SpanKind) else span_kind

        # Use explicit parent if provided, otherwise use current parent from stack
        parent_span_id = attributes.get("parent_span_id") or self.current_parent_span_id

        span_data = {
            "spanId": span_id,
            "name": name,
            "agentName": self.agent_name,
            "traceId": self.trace_id,
            "parentSpanId": parent_span_id,
            "spanKind": span_kind_str,
            "model": name if span_kind_str == SpanKind.LLM.value else None,
            "input": attributes.get("input"),
            "properties": attributes.get("properties", {}),
        }

        # Push this span onto the stack
        self._span_stack.append(span_id)

        self._log(
            f"▶ START [{span_kind_str}] {name}",
            {"spanId": span_id[:8], "parentId": parent_span_id[:8] if parent_span_id else None},
        )

        # Fire and forget
        try:
            self._tracer._foil.start_span(span_data)
        except Exception as err:
            print(f"Foil span start failed: {err}")

        span = Span(
            span_id=span_id,
            span_kind=span_kind_str,
            name=name,
            start_time=start_time,
            depth=self.depth + 1,
            attributes=attributes,
            _context=self,
        )

        self._active_spans[span_id] = span
        return span

    def wrap_in_span(
        self,
        span_kind: Union[SpanKind, str],
        name: str,
        fn: Callable[[Any], T],
        attributes: Optional[Dict[str, Any]] = None,
    ) -> T:
        """
        Wrap a function in a span.

        Args:
            span_kind: Type of span
            name: Name of the span
            fn: Function to execute (receives this TraceContext as argument)
            attributes: Additional span attributes

        Returns:
            Result of the function
        """
        span = self.start_span(span_kind, name, attributes)

        try:
            result = fn(self)
            span.end({"output": result})
            return result
        except Exception as error:
            span.end({"error": {"type": type(error).__name__, "message": str(error)}})
            raise

    def llm(
        self,
        model: str,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> Span:
        """
        Create an LLM span and return wrapper functions.

        Args:
            model: Model name
            attributes: Additional attributes

        Returns:
            LLM span with helper methods
        """
        return self.start_span(SpanKind.LLM, model, attributes)

    def llm_call(
        self,
        model: str,
        fn: Callable[[Any], T],
        attributes: Optional[Dict[str, Any]] = None,
    ) -> T:
        """
        Wrap a function in an LLM span (auto-creates and closes the span).

        Args:
            model: Model name
            fn: Function to execute
            attributes: Additional attributes

        Returns:
            Result of the function
        """
        return self.wrap_in_span(SpanKind.LLM, model, fn, attributes)

    def tool(
        self,
        tool_name: str,
        fn: Callable[[Any], T],
        attributes: Optional[Dict[str, Any]] = None,
    ) -> T:
        """
        Wrap a tool execution.

        Args:
            tool_name: Name of the tool
            fn: Tool function to execute
            attributes: Additional attributes (arguments, etc.)

        Returns:
            Result of the tool
        """
        return self.wrap_in_span(SpanKind.TOOL, tool_name, fn, attributes)

    def retriever(
        self,
        retriever_name: str,
        fn: Callable[[Any], T],
        attributes: Optional[Dict[str, Any]] = None,
    ) -> T:
        """
        Wrap a retriever operation.

        Args:
            retriever_name: Name of the retriever
            fn: Retriever function to execute
            attributes: Additional attributes

        Returns:
            Retriever results
        """
        return self.wrap_in_span(SpanKind.RETRIEVER, retriever_name, fn, attributes)

    def embedding(
        self,
        model: str,
        fn: Callable[[Any], T],
        attributes: Optional[Dict[str, Any]] = None,
    ) -> T:
        """
        Wrap an embedding operation.

        Args:
            model: Embedding model name
            fn: Embedding function to execute
            attributes: Additional attributes

        Returns:
            Embedding results
        """
        return self.wrap_in_span(SpanKind.EMBEDDING, model, fn, attributes)

    def record_signal(
        self,
        signal_name: str,
        value: Union[bool, int, float, str],
        options: Optional[Dict[str, Any]] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Record a signal for this trace.

        Args:
            signal_name: Name of the signal (e.g., 'user_thumbs', 'sentiment')
            value: Signal value
            options: Additional options including:
                - span_id: Specific span ID to associate with
                - signal_type: Type of signal (feedback, sentiment, quality, etc.)
                - source: Signal source (user, llm, system)
                - metadata: Additional metadata

        Returns:
            Recorded signal or None if failed
        """
        if options is None:
            options = {}

        signal_data = {
            "traceId": self.trace_id,
            "spanId": options.get("span_id"),
            "agentName": self.agent_name,
            "signalName": signal_name,
            "value": value,
            "signalType": options.get("signal_type", "feedback"),
            "source": options.get("source", "user"),
            "metadata": options.get("metadata"),
        }

        self._log(f"📊 SIGNAL: {signal_name}", {"value": value, "traceId": self.trace_id[:8]})

        try:
            return self._tracer._foil.record_signal(signal_data)
        except Exception as err:
            print(f"Foil signal recording failed: {err}")
            return None

    def record_feedback(
        self,
        positive: bool,
        options: Optional[Dict[str, Any]] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Record user feedback (thumbs up/down).

        Args:
            positive: True for thumbs up, False for thumbs down
            options: Additional options

        Returns:
            Recorded signal or None if failed
        """
        if options is None:
            options = {}

        return self.record_signal(
            "user_thumbs",
            positive,
            {
                "signal_type": "feedback",
                "source": "user",
                **options,
            },
        )

    def record_rating(
        self,
        rating: Union[int, float],
        options: Optional[Dict[str, Any]] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Record a user rating.

        Args:
            rating: Rating value (e.g., 1-5)
            options: Additional options

        Returns:
            Recorded signal or None if failed
        """
        if options is None:
            options = {}

        return self.record_signal(
            "user_rating",
            rating,
            {
                "signal_type": "feedback",
                "source": "user",
                **options,
            },
        )


class FoilTracer:
    """
    FoilTracer provides comprehensive tracing for AI applications.
    Captures the full lifecycle: prompt → LLM → tools → response.
    """

    def __init__(
        self,
        foil: Foil,
        agent_name: str = "default-agent",
        default_model: Optional[str] = None,
        debug: bool = False,
    ):
        """
        Initialize the FoilTracer.

        Args:
            foil: Foil client instance
            agent_name: Name of the agent
            default_model: Default model name
            debug: Enable debug logging
        """
        self._foil = foil
        self.agent_name = agent_name
        self.default_model = default_model
        self._active_traces: Dict[str, Dict[str, Any]] = {}
        self._debug = debug or _is_debug_enabled()

    def _log(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
        """Log debug message."""
        if not self._debug:
            return
        data_str = f" {data}" if data else ""
        print(f"[Foil] {message}{data_str}")

    def trace(
        self,
        fn: Callable[[TraceContext], T],
        name: Optional[str] = None,
        trace_id: Optional[str] = None,
        session_id: Optional[str] = None,
        input: Optional[Any] = None,
        properties: Optional[Dict[str, Any]] = None,
        user_id: Optional[str] = None,
        user_properties: Optional[Dict[str, Any]] = None,
        device: Optional[Dict[str, Any]] = None,
    ) -> T:
        """
        Execute a function within a traced context.

        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 for tracking which user triggered this trace
            user_properties: Additional user attributes (plan, company, role, etc.)
            device: Device/platform information containing:
                - platform: Platform type ('web', 'ios', 'android', 'desktop', 'api')
                - os: Operating system (e.g., 'iOS 17.0', 'macOS 14.0')
                - browser: Browser name and version (e.g., 'Chrome 120')
                - device_type: Device type ('phone', 'tablet', 'desktop')
                - app_version: Application version
                - locale: User locale (e.g., 'en-US')
                - timezone: User timezone (e.g., 'America/Los_Angeles')
                - screen_resolution: Screen resolution (e.g., '1920x1080')
                - user_agent: Full user agent string
                - ip: IP address (for geo-location)

        Returns:
            Result of the function
        """
        trace_id = trace_id or str(uuid.uuid4())
        span_id = str(uuid.uuid4())
        start_time = time.time()
        trace_name = name or "trace"

        self._log(
            f"━━━ TRACE START: {trace_name} ━━━",
            {"traceId": trace_id[:8], "spanId": span_id[:8]},
        )

        # Start root span
        root_span_data = {
            "spanId": span_id,
            "name": trace_name,
            "agentName": self.agent_name,
            "traceId": trace_id,
            "parentSpanId": None,
            "spanKind": SpanKind.AGENT.value,
            "input": input,
            "sessionId": session_id,
            "properties": properties or {},
            "userId": user_id,
            "userProperties": user_properties,
            "device": device,
        }

        try:
            self._foil.start_span(root_span_data)
        except Exception as err:
            print(f"Foil trace start failed: {err}")

        context = TraceContext(
            tracer=self,
            trace_id=trace_id,
            root_span_id=span_id,
            agent_name=self.agent_name,
            depth=0,
        )

        # Track active trace
        self._active_traces[trace_id] = {
            "span_id": span_id,
            "start_time": start_time,
            "context": context,
        }

        try:
            result = fn(context)
            end_time = time.time()
            duration = int((end_time - start_time) * 1000)

            self._log(
                f"━━━ TRACE END: {trace_name} ━━━",
                {"traceId": trace_id[:8], "duration": f"{duration}ms"},
            )

            # End root span
            try:
                self._foil.end_span({
                    "spanId": span_id,
                    "traceId": trace_id,
                    "agentName": self.agent_name,
                    "output": result,
                    "timing": {"totalDuration": duration},
                })
            except Exception as err:
                print(f"Foil trace end failed: {err}")

            self._active_traces.pop(trace_id, None)
            return result

        except Exception as error:
            end_time = time.time()
            duration = int((end_time - start_time) * 1000)

            self._log(
                f"━━━ TRACE ERROR: {trace_name} ━━━",
                {"traceId": trace_id[:8], "duration": f"{duration}ms", "error": str(error)},
            )

            try:
                self._foil.end_span({
                    "spanId": span_id,
                    "traceId": trace_id,
                    "agentName": self.agent_name,
                    "error": {"type": type(error).__name__, "message": str(error)},
                    "timing": {"totalDuration": duration},
                })
            except Exception as err:
                print(f"Foil trace end (error) failed: {err}")

            self._active_traces.pop(trace_id, None)
            raise

    def get_trace(self, trace_id: str) -> dict:
        """
        Get the trace data for a completed trace.

        Args:
            trace_id: The trace ID

        Returns:
            Trace data with all spans
        """
        return self._foil.get_trace(trace_id)

    def wrap_openai(self, openai_client: Any, context: Optional[TraceContext] = None) -> Any:
        """
        Create an OpenAI client wrapper with automatic tracing.

        Args:
            openai_client: OpenAI client instance
            context: Optional trace context for creating child spans

        Returns:
            Wrapped OpenAI client with tracing
        """
        tracer = self
        original_create = openai_client.chat.completions.create

        def wrapped_create(*args, **kwargs):
            span = None
            if context:
                span = context.start_span(
                    SpanKind.LLM,
                    kwargs.get("model", "openai"),
                    {
                        "input": kwargs.get("messages"),
                        "properties": {
                            "tools": [t.get("function", {}).get("name") for t in kwargs.get("tools", [])] if kwargs.get("tools") else None
                        },
                    },
                )

            start_time = time.time()
            ttft = None

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

                # Handle streaming
                if kwargs.get("stream"):
                    return tracer._wrap_openai_stream(response, span, start_time)

                # Non-streaming response
                end_time = time.time()
                choice = response.choices[0] if response.choices else None
                output = choice.message.content if choice and choice.message else None
                tool_calls = choice.message.tool_calls if choice and choice.message else None

                if span:
                    span.end({
                        "output": output or tool_calls,
                        "ttft": int((end_time - start_time) * 1000),
                        "tokens": {
                            "prompt": response.usage.prompt_tokens,
                            "completion": response.usage.completion_tokens,
                            "total": response.usage.total_tokens,
                        } if response.usage else None,
                    })

                return response

            except Exception as error:
                if span:
                    span.end({"error": {"type": type(error).__name__, "message": str(error)}})
                raise

        openai_client.chat.completions.create = wrapped_create
        return openai_client

    def _wrap_openai_stream(self, stream, span: Optional[Span], start_time: float):
        """Wrap OpenAI streaming response."""
        first_chunk = True
        ttft = None
        collected_content = ""
        collected_tool_calls = []
        usage = None

        def stream_wrapper():
            nonlocal first_chunk, ttft, collected_content, collected_tool_calls, usage

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

                    delta = chunk.choices[0].delta if chunk.choices else None
                    if delta:
                        if delta.content:
                            collected_content += delta.content

                        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.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

                    if hasattr(chunk, "usage") and chunk.usage:
                        usage = chunk.usage

                    yield chunk

            finally:
                # After stream completes, end the span
                if span:
                    span.end({
                        "output": collected_content or collected_tool_calls or None,
                        "ttft": ttft,
                        "tokens": {
                            "prompt": usage.prompt_tokens,
                            "completion": usage.completion_tokens,
                            "total": usage.total_tokens,
                        } if usage else None,
                    })

        return stream_wrapper()

    def create_tool_executor(
        self,
        tools: Dict[str, Callable],
        context: TraceContext,
    ) -> Callable:
        """
        Create a tool executor that automatically traces tool calls.

        Args:
            tools: Map of tool name to function
            context: Current trace context

        Returns:
            Tool executor function
        """
        import json

        def executor(tool_call) -> Any:
            tool_name = getattr(tool_call, "function", tool_call).name if hasattr(tool_call, "function") else tool_call.get("name")
            tool_fn = tools.get(tool_name)

            if not tool_fn:
                raise ValueError(f"Tool not found: {tool_name}")

            args = getattr(tool_call, "function", tool_call).arguments if hasattr(tool_call, "function") else tool_call.get("arguments")
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except json.JSONDecodeError:
                    pass  # Keep as string if not valid JSON

            return context.tool(
                tool_name,
                lambda _: tool_fn(args) if callable(tool_fn) else tool_fn,
                {"input": args},
            )

        return executor


def create_foil_tracer(
    api_key: str,
    agent_name: str,
    base_url: Optional[str] = None,
    default_model: Optional[str] = None,
    debug: bool = False,
) -> FoilTracer:
    """
    Factory function to create a FoilTracer.

    Args:
        api_key: Foil API key
        agent_name: Name of the agent
        base_url: Foil API base URL
        default_model: Default model name
        debug: Enable debug logging (also enabled via FOIL_DEBUG=true env var)

    Returns:
        Configured FoilTracer instance

    Raises:
        ValueError: If api_key or agent_name is not provided
    """
    if not api_key:
        raise ValueError("api_key is required")
    if not agent_name:
        raise ValueError("agent_name is required")

    foil = Foil(api_key=api_key, base_url=base_url or "https://api.getfoil.ai/api")

    return FoilTracer(
        foil=foil,
        agent_name=agent_name,
        default_model=default_model,
        debug=debug,
    )
