from __future__ import annotations

import abc
import contextvars
import hashlib
import threading
from collections import OrderedDict
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Any

from ..logger import logger
from . import util
from .processor_interface import TracingProcessor
from .scope import Scope


class Trace(abc.ABC):
    """A complete end-to-end workflow containing related spans and metadata.

    A trace represents a logical workflow or operation (e.g., "Customer Service Query"
    or "Code Generation") and contains all the spans (individual operations) that occur
    during that workflow.

    Example:
        ```python
        # Basic trace usage
        with trace("Order Processing") as t:
            validation_result = await Runner.run(validator, order_data)
            if validation_result.approved:
                await Runner.run(processor, order_data)

        # Trace with metadata and grouping
        with trace(
            "Customer Service",
            group_id="chat_123",
            metadata={"customer": "user_456"}
        ) as t:
            result = await Runner.run(support_agent, query)
        ```

    Notes:
        - Use descriptive workflow names
        - Group related traces with consistent group_ids
        - Add relevant metadata for filtering/analysis
        - Use context managers for reliable cleanup
        - Consider privacy when adding trace data
    """

    @abc.abstractmethod
    def __enter__(self) -> Trace:
        pass

    @abc.abstractmethod
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    @abc.abstractmethod
    def start(self, mark_as_current: bool = False):
        """Start the trace and optionally mark it as the current trace.

        Args:
            mark_as_current: If true, marks this trace as the current trace
                in the execution context.

        Notes:
            - Must be called before any spans can be added
            - Only one trace can be current at a time
            - Thread-safe when using mark_as_current
        """
        pass

    @abc.abstractmethod
    def finish(self, reset_current: bool = False):
        """Finish the trace and optionally reset the current trace.

        Args:
            reset_current: If true, resets the current trace to the previous
                trace in the execution context.

        Notes:
            - Must be called to complete the trace
            - Finalizes all open spans
            - Thread-safe when using reset_current
        """
        pass

    @property
    @abc.abstractmethod
    def trace_id(self) -> str:
        """Get the unique identifier for this trace.

        Returns:
            str: The trace's unique ID in the format 'trace_<32_alphanumeric>'

        Notes:
            - IDs are globally unique
            - Used to link spans to their parent trace
            - Can be used to look up traces in the dashboard
        """
        pass

    @property
    @abc.abstractmethod
    def name(self) -> str:
        """Get the human-readable name of this workflow trace.

        Returns:
            str: The workflow name (e.g., "Customer Service", "Data Processing")

        Notes:
            - Should be descriptive and meaningful
            - Used for grouping and filtering in the dashboard
            - Helps identify the purpose of the trace
        """
        pass

    @abc.abstractmethod
    def export(self) -> dict[str, Any] | None:
        """Export the trace data as a serializable dictionary.

        Returns:
            dict | None: Dictionary containing trace data, or None if tracing is disabled.

        Notes:
            - Includes all spans and their data
            - Used for sending traces to backends
            - May include metadata and group ID
        """
        pass

    @property
    @abc.abstractmethod
    def tracing_api_key(self) -> str | None:
        """The API key to use when exporting this trace and its spans."""
        pass

    def to_json(self, *, include_tracing_api_key: bool = False) -> dict[str, Any] | None:
        """Serialize trace metadata for persistence or transport.

        Args:
            include_tracing_api_key: When True, include the tracing API key. Defaults to False
                to avoid persisting secrets unintentionally.
        """
        exported = self.export()
        if exported is None:
            return None
        payload = dict(exported)
        if include_tracing_api_key and self.tracing_api_key:
            payload["tracing_api_key"] = self.tracing_api_key
        return payload


def _hash_tracing_api_key(tracing_api_key: str | None) -> str | None:
    # Persist only a fingerprint so resumed runs can verify the same explicit
    # tracing key without storing the secret.
    if tracing_api_key is None:
        return None
    return hashlib.sha256(tracing_api_key.encode("utf-8")).hexdigest()


@dataclass
class TraceState:
    """Serializable trace metadata for run state persistence."""

    trace_id: str | None = None
    workflow_name: str | None = None
    group_id: str | None = None
    metadata: dict[str, Any] | None = None
    tracing_api_key: str | None = None
    tracing_api_key_hash: str | None = None
    object_type: str | None = None
    extra: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_trace(cls, trace: Trace | None) -> TraceState | None:
        if trace is None:
            return None
        payload = trace.to_json(include_tracing_api_key=True)
        return cls.from_json(payload)

    @classmethod
    def from_json(cls, payload: Mapping[str, Any] | None) -> TraceState | None:
        if not payload:
            return None
        data = dict(payload)
        object_type = data.pop("object", None)
        trace_id = data.pop("id", None) or data.pop("trace_id", None)
        workflow_name = data.pop("workflow_name", None)
        group_id = data.pop("group_id", None)
        metadata_value = data.pop("metadata", None)
        metadata = metadata_value if isinstance(metadata_value, dict) else None
        tracing_api_key = data.pop("tracing_api_key", None)
        tracing_api_key_hash = data.pop("tracing_api_key_hash", None)
        resolved_tracing_api_key = tracing_api_key if isinstance(tracing_api_key, str) else None
        resolved_tracing_api_key_hash = _hash_tracing_api_key(resolved_tracing_api_key)
        # Secure snapshots may strip the raw key, so keep the stored
        # fingerprint for resume-time matching.
        if resolved_tracing_api_key_hash is None and isinstance(tracing_api_key_hash, str):
            resolved_tracing_api_key_hash = tracing_api_key_hash
        return cls(
            trace_id=trace_id if isinstance(trace_id, str) else None,
            workflow_name=workflow_name if isinstance(workflow_name, str) else None,
            group_id=group_id if isinstance(group_id, str) else None,
            metadata=metadata,
            tracing_api_key=resolved_tracing_api_key,
            tracing_api_key_hash=resolved_tracing_api_key_hash,
            object_type=object_type if isinstance(object_type, str) else None,
            extra=data,
        )

    def to_json(self, *, include_tracing_api_key: bool = False) -> dict[str, Any] | None:
        if (
            self.trace_id is None
            and self.workflow_name is None
            and self.group_id is None
            and self.metadata is None
            and self.tracing_api_key is None
            and self.tracing_api_key_hash is None
            and self.object_type is None
            and not self.extra
        ):
            return None
        payload: dict[str, Any] = {}
        if self.object_type:
            payload["object"] = self.object_type
        if self.trace_id:
            payload["id"] = self.trace_id
        if self.workflow_name is not None:
            payload["workflow_name"] = self.workflow_name
        if self.group_id is not None:
            payload["group_id"] = self.group_id
        if self.metadata is not None:
            payload["metadata"] = dict(self.metadata)
        if include_tracing_api_key and self.tracing_api_key:
            payload["tracing_api_key"] = self.tracing_api_key
        if self.tracing_api_key_hash:
            # Always persist the fingerprint so default RunState snapshots
            # can still validate explicit resume keys.
            payload["tracing_api_key_hash"] = self.tracing_api_key_hash
        for key, value in self.extra.items():
            if key not in payload:
                payload[key] = value
        return payload


_MAX_STARTED_TRACE_IDS = 4096
_started_trace_ids: OrderedDict[str, None] = OrderedDict()
_started_trace_ids_lock = threading.Lock()


def _mark_trace_id_started(trace_id: str | None) -> None:
    if not trace_id or trace_id == "no-op":
        return
    with _started_trace_ids_lock:
        if trace_id in _started_trace_ids:
            _started_trace_ids.move_to_end(trace_id)
        else:
            _started_trace_ids[trace_id] = None

        while len(_started_trace_ids) > _MAX_STARTED_TRACE_IDS:
            _started_trace_ids.popitem(last=False)


def _trace_id_was_started(trace_id: str | None) -> bool:
    if not trace_id or trace_id == "no-op":
        return False
    with _started_trace_ids_lock:
        return trace_id in _started_trace_ids


class ReattachedTrace(Trace):
    """A trace context rebuilt from persisted state without re-emitting trace start events."""

    __slots__ = (
        "_name",
        "_trace_id",
        "_tracing_api_key",
        "group_id",
        "metadata",
        "_prev_context_token",
        "_started",
    )

    def __init__(
        self,
        *,
        name: str,
        trace_id: str,
        group_id: str | None,
        metadata: dict[str, Any] | None,
        tracing_api_key: str | None,
    ) -> None:
        self._name = name
        self._trace_id = trace_id
        self._tracing_api_key = tracing_api_key
        self.group_id = group_id
        self.metadata = metadata
        self._prev_context_token: contextvars.Token[Trace | None] | None = None
        self._started = False

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

    @property
    def name(self) -> str:
        return self._name

    @property
    def tracing_api_key(self) -> str | None:
        return self._tracing_api_key

    def start(self, mark_as_current: bool = False):
        if self._started:
            return

        self._started = True
        _mark_trace_id_started(self.trace_id)

        if mark_as_current:
            self._prev_context_token = Scope.set_current_trace(self)

    def finish(self, reset_current: bool = False):
        if not self._started:
            return

        if reset_current and self._prev_context_token is not None:
            Scope.reset_current_trace(self._prev_context_token)
            self._prev_context_token = None

    def __enter__(self) -> Trace:
        if self._started:
            if not self._prev_context_token:
                logger.error("Trace already started but no context token set")
            return self

        self.start(mark_as_current=True)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.finish(reset_current=exc_type is not GeneratorExit)

    def export(self) -> dict[str, Any] | None:
        return {
            "object": "trace",
            "id": self.trace_id,
            "workflow_name": self.name,
            "group_id": self.group_id,
            "metadata": self.metadata,
        }


def reattach_trace(trace_state: TraceState, *, tracing_api_key: str | None = None) -> Trace | None:
    """Build a live trace context from persisted state without notifying processors."""
    if trace_state.trace_id is None:
        return None
    return ReattachedTrace(
        name=trace_state.workflow_name or "Agent workflow",
        trace_id=trace_state.trace_id,
        group_id=trace_state.group_id,
        metadata=dict(trace_state.metadata) if trace_state.metadata is not None else None,
        tracing_api_key=(
            trace_state.tracing_api_key
            if trace_state.tracing_api_key is not None
            else tracing_api_key
        ),
    )


class NoOpTrace(Trace):
    """A no-op implementation of Trace that doesn't record any data.

    Used when tracing is disabled but trace operations still need to work.
    Maintains proper context management but doesn't store or export any data.

    Example:
        ```python
        # When tracing is disabled, traces become NoOpTrace
        with trace("Disabled Workflow") as t:
            # Operations still work but nothing is recorded
            await Runner.run(agent, "query")
        ```
    """

    def __init__(self):
        self._started = False
        self._prev_context_token: contextvars.Token[Trace | None] | None = None

    def __enter__(self) -> Trace:
        if self._started:
            if not self._prev_context_token:
                logger.error("Trace already started but no context token set")
            return self

        self._started = True
        self.start(mark_as_current=True)

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.finish(reset_current=True)

    def start(self, mark_as_current: bool = False):
        if mark_as_current:
            self._prev_context_token = Scope.set_current_trace(self)

    def finish(self, reset_current: bool = False):
        if reset_current and self._prev_context_token is not None:
            Scope.reset_current_trace(self._prev_context_token)
            self._prev_context_token = None

    @property
    def trace_id(self) -> str:
        """The trace's unique identifier.

        Returns:
            str: A unique ID for this trace.
        """
        return "no-op"

    @property
    def name(self) -> str:
        """The workflow name for this trace.

        Returns:
            str: Human-readable name describing this workflow.
        """
        return "no-op"

    def export(self) -> dict[str, Any] | None:
        """Export the trace data as a dictionary.

        Returns:
            dict | None: Trace data in exportable format, or None if no data.
        """
        return None

    @property
    def tracing_api_key(self) -> str | None:
        return None


NO_OP_TRACE = NoOpTrace()


class TraceImpl(Trace):
    """
    A trace that will be recorded by the tracing library.
    """

    __slots__ = (
        "_name",
        "_trace_id",
        "_tracing_api_key",
        "group_id",
        "metadata",
        "_prev_context_token",
        "_processor",
        "_started",
    )

    def __init__(
        self,
        name: str,
        trace_id: str | None,
        group_id: str | None,
        metadata: dict[str, Any] | None,
        processor: TracingProcessor,
        tracing_api_key: str | None = None,
    ):
        self._name = name
        self._trace_id = trace_id or util.gen_trace_id()
        self._tracing_api_key = tracing_api_key
        self.group_id = group_id
        self.metadata = metadata
        self._prev_context_token: contextvars.Token[Trace | None] | None = None
        self._processor = processor
        self._started = False

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

    @property
    def name(self) -> str:
        return self._name

    @property
    def tracing_api_key(self) -> str | None:
        return self._tracing_api_key

    def start(self, mark_as_current: bool = False):
        if self._started:
            return

        self._started = True
        self._processor.on_trace_start(self)
        _mark_trace_id_started(self.trace_id)

        if mark_as_current:
            self._prev_context_token = Scope.set_current_trace(self)

    def finish(self, reset_current: bool = False):
        if not self._started:
            return

        self._processor.on_trace_end(self)

        if reset_current and self._prev_context_token is not None:
            Scope.reset_current_trace(self._prev_context_token)
            self._prev_context_token = None

    def __enter__(self) -> Trace:
        if self._started:
            if not self._prev_context_token:
                logger.error("Trace already started but no context token set")
            return self

        self.start(mark_as_current=True)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.finish(reset_current=exc_type is not GeneratorExit)

    def export(self) -> dict[str, Any] | None:
        return {
            "object": "trace",
            "id": self.trace_id,
            "workflow_name": self.name,
            "group_id": self.group_id,
            "metadata": self.metadata,
        }