"""
A Span is an immutable record/snapshot of something that happend
"""

from __future__ import annotations

import enum
from dataclasses import dataclass, field
from typing import Any, Union

from agenx._internal._clock import (
    Timestamp,
    duration_ms, 
    wall_ns_to_iso
)


class SpanType(str, enum.Enum):
    """Type of unit of work this span represents"""
    ROOT      = "root"
    STEP      = "step"
    LLM       = "llm"
    RETRIEVER = "retriever"
    TOOL      = "tool"
    GUARD     = "guard"
    AGENT     = "agent"


class SpanStatus(str, enum.Enum):
    """Lifecycle state of a span"""
    RUNNING = "running"  
    OK      = "ok"       
    ERROR   = "error"  


@dataclass(slots=True)
class StepAttributes:
    """
    Attributes for SpanType.STEP and SpanType.ROOT

    The generic container for any discrete unit of business logic
    """
    input:    Any                  = None
    output:   Any                  = None
    metadata: dict[str, Any]       = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "input":    self.input,
            "output":   self.output,
            "metadata": self.metadata,
        }
    

@dataclass(slots=True)
class LLMAttributes:
    """
    Attributes for SpanType.LLM
    """
    model:         str            = ""
    provider:      str            = ""
    prompt:        str            = ""
    temperature:   float          = 0.0
    output:        str            = ""
    tokens_input:  int            = 0
    tokens_output: int            = 0
    tokens_total:  int            = 0
    metadata:      dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "model":         self.model,
            "provider":      self.provider,
            "prompt":        self.prompt,
            "temperature":   self.temperature,
            "output":        self.output,
            "tokens_input":  self.tokens_input,
            "tokens_output": self.tokens_output,
            "tokens_total":  self.tokens_total,
            "metadata":      self.metadata,
        }
    

@dataclass(slots=True)
class RetrieverAttributes:
    """
    Attributes for SpanType.RETRIEVER
    """
    query:      str            = ""
    documents:  list[Any]      = field(default_factory=list)
    top_k:      int            = 0
    returned:   int            = 0   
    metadata:   dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "query":     self.query,
            "documents": self.documents,
            "top_k":     self.top_k,
            "returned":  self.returned,
            "metadata":  self.metadata,
        }


@dataclass(slots=True)
class ToolAttributes:
    """
    Attributes for SpanType.TOOL
    """
    tool_name:  str            = ""
    input:      Any            = None   
    output:     Any            = None   
    metadata:   dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "tool_name": self.tool_name,
            "input":     self.input,
            "output":    self.output,
            "metadata":  self.metadata,
        }
    

@dataclass(slots=True)
class GuardAttributes:
    """
    Attributes for SpanType.GUARD
    """
    guard_name: str            = ""
    input:      Any            = None   
    passed:     bool           = True  
    reason:     str            = ""    
    metadata:   dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "guard_name": self.guard_name,
            "input":      self.input,
            "passed":     self.passed,
            "reason":     self.reason,
            "metadata":   self.metadata,
        }


@dataclass(slots=True)
class AgentAttributes:
    """
    Attributes for SpanType.AGENT.

    An agent is an autonomous loop that plans and calls tools. Captures
    which tools were available and how many iterations it ran.
    """
    input:      Any            = None
    output:     Any            = None
    tools:      list[str]      = field(default_factory=list)  
    iterations: int            = 0
    metadata:   dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "input":      self.input,
            "output":     self.output,
            "tools":      self.tools,
            "iterations": self.iterations,
            "metadata":   self.metadata,
        }
    

@dataclass(slots=True)
class SpanEvent:
    """
    A timestamped point-in-time event attached to a span.

    To be used for things that happen *inside* a span but aren't worth their own span
    """
    name: str
    timestamp: Timestamp
    attributes: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "name":         self.name,
            "timestamp":    wall_ns_to_iso(self.timestamp.wall_ns),
            "attribute":    self.attributes
        } 
    

SpanAttributes = Union[
    StepAttributes,
    LLMAttributes,
    RetrieverAttributes,
    ToolAttributes,
    GuardAttributes,
    AgentAttributes,
]


SPAN_TYPE_TO_ATTRIBUTES: dict[SpanType, type] = {
    SpanType.ROOT:      StepAttributes,    
    SpanType.STEP:      StepAttributes,
    SpanType.LLM:       LLMAttributes,
    SpanType.RETRIEVER: RetrieverAttributes,
    SpanType.TOOL:      ToolAttributes,
    SpanType.GUARD:     GuardAttributes,
    SpanType.AGENT:     AgentAttributes,
}


@dataclass(slots=True)
class Span:
    """
    The core data record for a single unit of traced work.

    This is a mutable dataclass — LiveSpan writes into it as the work progresses.
    After the span ends, treat it as read-only.
    """
    span_id:   str
    trace_id:  str
    parent_id: str | None

    name:      str
    type:      SpanType

    start_time: Timestamp
    end_time:   Timestamp | None = None

    status: SpanStatus = SpanStatus.RUNNING
    error:  BaseException | None = None

    attributes: SpanAttributes = field(default_factory=StepAttributes)
    events:     list[SpanEvent] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a plain dict. Safe to JSON-encode"""
        return {
            "span_id":    self.span_id,
            "trace_id":   self.trace_id,
            "parent_id":  self.parent_id,
            "name":       self.name,
            "type":       self.type.value,
            "status":     self.status.value,
            "start_time": wall_ns_to_iso(self.start_time.wall_ns),
            "end_time":   wall_ns_to_iso(self.end_time.wall_ns) if self.end_time else None,
            "duration_ms": (
                duration_ms(self.start_time, self.end_time)
                if self.end_time else None
            ),
            "error":      repr(self.error) if self.error else None,
            "attributes": self.attributes,
            "events": [
                {
                    "name":       e.name,
                    "timestamp":  wall_ns_to_iso(e.timestamp.wall_ns),
                    "attributes": e.attributes,
                }
                for e in self.events
            ],
        }

