"""
LiveSpan: a running span you can interact with

LiveSpan wraps a Span (the data record) and provides:
  - Context manager protocol (__enter__ / __exit__)
  - Mutation helpers: set_attribute(), add_event()
  - Guaranteed on_end callback firing, even on exception or GC
"""

from __future__ import annotations

from typing import Any, Callable

from agenx._internal import now
from agenx.trace.spans.models import (
    Span,
    SpanEvent,
    SpanStatus,
    SpanType,
    StepAttributes,
    LLMAttributes,
    RetrieverAttributes,
    ToolAttributes,
    GuardAttributes,
    AgentAttributes
)

class LiveSpan:
    """
    A span that is currently in flight
    """

    __slots__ = ("_span", "_callbacks", "_ended")

    def __init__(
        self,
        span: Span,
        on_end_callbacks: list[Callable[[Span], None]]
    ) -> None:
        self._span = span
        self._callbacks = on_end_callbacks
        self._ended = False

    @property
    def step(self) -> StepAttributes:
        """
        Access step/root attributes with full type safety
        """
        if not isinstance(self._span.attributes, StepAttributes):
            raise TypeError(
                f"span '{self._span.name}' is type {self._span.type.value!r}, "
                f"not 'step' or 'root'. Use s.{self._span.type.value} instead"
            )
        return self._span.attributes
    
    @property
    def llm(self) -> LLMAttributes:
        """
        Access LLM attributes with full type safety
        """
        if not isinstance(self._span.attributes, LLMAttributes):
            raise TypeError(
                f"span '{self._span.name}' is type {self._span.type.value!r}, "
                f"not 'llm'. Use s.{self._span.type.value} instead"
            )
        return self._span.attributes
    
    @property
    def retriever(self) -> RetrieverAttributes:
        """
        Access retriever attributes with full type safety
        """
        if not isinstance(self._span.attributes, RetrieverAttributes):
            raise TypeError(
                f"span '{self._span.name}' is type {self._span.type.value!r}, "
                f"not 'retriever'. Use s.{self._span.type.value} instead"
            )
        return self._span.attributes
    
    @property
    def tool(self) -> ToolAttributes:
        """
        Access tool attributes with full type safety
        """
        if not isinstance(self._span.attributes, ToolAttributes):
            raise TypeError(
                f"span '{self._span.name}' is type {self._span.type.value!r}, "
                f"not 'tool'. Use s.{self._span.type.value} instead"
            )
        return self._span.attributes
    
    @property
    def guard(self) -> GuardAttributes:
        """
        Access guard attributes with full type safety
        """
        if not isinstance(self._span.attributes, GuardAttributes):
            raise TypeError(
                f"span '{self._span.name}' is type {self._span.type.value!r}, "
                f"not 'guard'. Use s.{self._span.type.value} instead"
            )
        return self._span.attributes

    @property
    def agent(self) -> AgentAttributes:
        """
        Access agent attributes with full type safety
        """
        if not isinstance(self._span.attributes, AgentAttributes):
            raise TypeError(
                f"span '{self._span.name}' is type {self._span.type.value!r}, "
                f"not 'agent'. Use s.{self._span.type.value} instead"
            )
        return self._span.attributes

    def set_attribute(self, key: str, value: Any) -> "LiveSpan":
        """
        Directly set a key on the underlying attributes dataclass, or fall
        back to attributes.metadata for unknown keys
        """
        try:
            if self._ended:
                return self
            attrs = self._span.attributes
            if hasattr(attrs, key):
                setattr(attrs, key, value)
            elif hasattr(attrs, "metadata"):
                attrs.metadata[key] = value    # adds unknown keys to metadata object
        except Exception:
            pass
        return self
    
    def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> "LiveSpan":
        """
        Record a named point-in-time event on this span
        """
        try:
            if not self._ended:
                self._span.events.append(
                    SpanEvent(
                        name=name,
                        timestamp=now(),
                        attributes=attributes or {},
                    )
                )
        except Exception:
            pass
        return self
        
    @property
    def span_id(self) -> str:
        return self._span.span_id

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

    @property
    def parent_id(self) -> str | None:
        return self._span.parent_id

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

    @property
    def type(self) -> SpanType:
        return self._span.type

    @property
    def status(self) -> SpanStatus:
        return self._span.status

    @property
    def ended(self) -> bool:
        return self._ended
    
    def __enter__(self) -> "LiveSpan":
        return self
    
    def __exit__(
        self, 
        exc_type: type | None, 
        exc_val: BaseException | None, 
        tb: Any
    ) -> bool:
        self._finish(exc_val)
        return False
    
    def _finish(self, exc: BaseException | None = None) -> None:
        """
        Finalise the span. Idempotent — only the first call does anything

        Called by:
        - __exit__ on normal or exceptional context manager unwind
        - __del__ as a last-resort for abandoned spans
        """
        if self._ended:
            return
        self._ended = True

        try:
            self._span.end_time = now()
            if exc is not None:
                self._span.status = SpanStatus.ERROR
                self._span.error = exc
            else:
                self._span.status = SpanStatus.OK
        except Exception:
            pass

        for callback in self._callbacks:
            try:
                callback(self._span)
            except Exception:
                pass

    def __del__(self) -> None:
        """
        Safety net for spans abandoned without calling __exit__
        """
        if not self._ended:
            try:
                self._finish()
            except Exception:
                pass

    def __repr__(self) -> str:
        return (
            f"LiveSpan("
            f"name={self._span.name!r}, "
            f"type={self._span.type.value}, "
            f"status={self._span.status.value}, "
            f"trace_id={self._span.trace_id[:8]}...)"
        )
