"""
AgentBase Tracer - Context managers for creating traces and spans.
"""

import time
import contextvars
from agentbase_sdk.client import AgentBaseClient, TraceData, SpanData, generate_id, calculate_cost

_current_trace = contextvars.ContextVar("_current_trace", default=None)
_current_span = contextvars.ContextVar("_current_span", default=None)


class Tracer:
    def __init__(self, name, session_id="", tags=None, metadata=None):
        self._trace = TraceData(
            trace_id=generate_id(), agent_name=name,
            session_id=session_id, tags=tags or [], metadata=metadata or {},
        )
        self._token = None

    def __enter__(self):
        self._trace.started_at = time.time()
        self._trace.status = "running"
        self._token = _current_trace.set(self._trace)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._trace.ended_at = time.time()
        self._trace.duration_ms = int((self._trace.ended_at - self._trace.started_at) * 1000)
        if exc_type is not None:
            self._trace.status = "error"
            self._trace.error_message = str(exc_val)
            self._trace.error_type = exc_type.__name__
        elif self._trace.status == "running":
            self._trace.status = "success"
        import agentbase_sdk
        try:
            client = agentbase_sdk.get_client()
            client.enqueue_trace(self._trace)
        except RuntimeError:
            pass
        if self._token:
            _current_trace.reset(self._token)
        return False

    @property
    def trace_id(self):
        return self._trace.trace_id

    def set_input(self, text):
        self._trace.user_input = str(text)

    def set_output(self, text):
        self._trace.final_output = str(text)

    def set_model(self, model):
        self._trace.model = model

    def set_framework(self, framework):
        self._trace.framework = framework

    def set_error(self, error_type, message):
        self._trace.status = "error"
        self._trace.error_type = error_type
        self._trace.error_message = message

    def set_metadata(self, key, value):
        self._trace.metadata[key] = value

    def add_tag(self, tag):
        self._trace.tags.append(tag)

    def add_tokens(self, prompt=0, completion=0):
        self._trace.prompt_tokens += prompt
        self._trace.completion_tokens += completion
        self._trace.total_tokens += prompt + completion

    def add_cost(self, cost_usd):
        self._trace.total_cost_usd += cost_usd


class Span:
    def __init__(self, name, span_type="custom", metadata=None):
        current_trace = _current_trace.get()
        parent_span = _current_span.get()
        self._span = SpanData(
            span_id=generate_id(),
            trace_id=current_trace.trace_id if current_trace else generate_id(),
            parent_span_id=parent_span.span_id if parent_span else "",
            name=name, span_type=span_type, metadata=metadata or {},
        )
        self._token = None
        self._trace = current_trace

    def __enter__(self):
        self._span.started_at = time.time()
        self._span.status = "running"
        self._token = _current_span.set(self._span)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._span.ended_at = time.time()
        self._span.duration_ms = int((self._span.ended_at - self._span.started_at) * 1000)
        if exc_type is not None:
            self._span.status = "error"
            self._span.error_message = str(exc_val)
            self._span.error_type = exc_type.__name__
            self._try_heal()
        elif self._span.status == "running":
            self._span.status = "success"
        if self._span.total_tokens > 0 and self._span.model:
            self._span.cost_usd = calculate_cost(
                self._span.model, self._span.prompt_tokens, self._span.completion_tokens)
        if self._trace:
            self._trace.total_tokens += self._span.total_tokens
            self._trace.prompt_tokens += self._span.prompt_tokens
            self._trace.completion_tokens += self._span.completion_tokens
            self._trace.total_cost_usd += self._span.cost_usd
        import agentbase_sdk
        try:
            client = agentbase_sdk.get_client()
            client.enqueue_span(self._span)
        except RuntimeError:
            pass
        if self._token:
            _current_span.reset(self._token)
        return False

    def _try_heal(self):
        import agentbase_sdk
        try:
            client = agentbase_sdk.get_client()
            heal_result = client.check_heal(self._span)
            if heal_result:
                self._span.was_healed = True
                self._span.heal_action = heal_result.get("action", "")
                self._span.heal_details = str(heal_result.get("details", ""))
                self._span.status = "healed"
        except Exception:
            pass

    @property
    def span_id(self):
        return self._span.span_id

    def set_input(self, text):
        self._span.input = str(text)

    def set_output(self, text):
        self._span.output = str(text)

    def set_model(self, model):
        self._span.model = model

    def set_tokens(self, prompt_tokens=0, completion_tokens=0, total_tokens=None):
        self._span.prompt_tokens = prompt_tokens
        self._span.completion_tokens = completion_tokens
        self._span.total_tokens = total_tokens or (prompt_tokens + completion_tokens)

    def set_tool(self, tool_name, params=None):
        self._span.tool_name = tool_name
        self._span.span_type = "tool"
        if params:
            import json
            self._span.tool_params = json.dumps(params)

    def set_error(self, error_type, message):
        self._span.status = "error"
        self._span.error_type = error_type
        self._span.error_message = message

    def set_metadata(self, key, value):
        self._span.metadata[key] = value


def trace(name, session_id="", tags=None, metadata=None):
    return Tracer(name, session_id=session_id, tags=tags, metadata=metadata)


def span(name, span_type="custom", metadata=None):
    return Span(name, span_type=span_type, metadata=metadata)
