"""
AgentBase SDK Client - handles communication with AgentBase API.
"""

import os
import time
import uuid
import queue
import logging
import threading
from dataclasses import dataclass, field

import httpx

logger = logging.getLogger("agentbase")


@dataclass
class SpanData:
    span_id: str
    trace_id: str
    parent_span_id: str = ""
    name: str = ""
    span_type: str = "custom"
    status: str = "running"
    started_at: float = 0.0
    ended_at: float = None
    duration_ms: int = 0
    model: str = ""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cost_usd: float = 0.0
    input: str = ""
    output: str = ""
    error_message: str = ""
    error_type: str = ""
    tool_name: str = ""
    tool_params: str = "{}"
    was_healed: bool = False
    heal_action: str = ""
    heal_details: str = ""
    metadata: dict = field(default_factory=dict)

    def to_dict(self):
        return {
            "span_id": self.span_id, "trace_id": self.trace_id,
            "parent_span_id": self.parent_span_id, "name": self.name,
            "span_type": self.span_type, "status": self.status,
            "started_at": self.started_at, "ended_at": self.ended_at,
            "duration_ms": self.duration_ms, "model": self.model,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_tokens": self.total_tokens, "cost_usd": self.cost_usd,
            "input": self.input[:10000], "output": self.output[:10000],
            "error_message": self.error_message, "error_type": self.error_type,
            "tool_name": self.tool_name, "tool_params": self.tool_params,
            "was_healed": self.was_healed, "heal_action": self.heal_action,
            "heal_details": self.heal_details, "metadata": self.metadata,
        }


@dataclass
class TraceData:
    trace_id: str
    agent_name: str
    session_id: str = ""
    status: str = "running"
    started_at: float = 0.0
    ended_at: float = None
    duration_ms: int = 0
    total_tokens: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_cost_usd: float = 0.0
    model: str = ""
    framework: str = ""
    user_input: str = ""
    final_output: str = ""
    error_message: str = ""
    error_type: str = ""
    metadata: dict = field(default_factory=dict)
    tags: list = field(default_factory=list)

    def to_dict(self):
        return {
            "trace_id": self.trace_id, "agent_name": self.agent_name,
            "session_id": self.session_id, "status": self.status,
            "started_at": self.started_at, "ended_at": self.ended_at,
            "duration_ms": self.duration_ms, "total_tokens": self.total_tokens,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
            "total_cost_usd": self.total_cost_usd, "model": self.model,
            "framework": self.framework,
            "user_input": self.user_input[:10000],
            "final_output": self.final_output[:10000],
            "error_message": self.error_message, "error_type": self.error_type,
            "metadata": self.metadata, "tags": self.tags,
        }


MODEL_PRICING = {
    "claude-sonnet-4-20250514": {"input": 3.0, "output": 15.0},
    "claude-opus-4-20250514": {"input": 15.0, "output": 75.0},
    "gpt-4o": {"input": 2.5, "output": 10.0},
    "gpt-4o-mini": {"input": 0.15, "output": 0.6},
    "_default": {"input": 1.0, "output": 3.0},
}


def calculate_cost(model, prompt_tokens, completion_tokens):
    pricing = MODEL_PRICING.get(model, MODEL_PRICING["_default"])
    input_cost = (prompt_tokens / 1_000_000) * pricing["input"]
    output_cost = (completion_tokens / 1_000_000) * pricing["output"]
    return round(input_cost + output_cost, 6)


class AgentBaseClient:
    def __init__(self, api_key=None, agent_name=None, endpoint="https://api.agentbase.dev",
                 auto_heal=True, environment="production", debug=False,
                 batch_size=50, flush_interval=2.0):
        self.api_key = api_key or os.environ.get("AB_API_KEY", "")
        self.agent_name = agent_name or os.environ.get("AB_AGENT_NAME", "unnamed-agent")
        self.endpoint = endpoint.rstrip("/")
        self.auto_heal = auto_heal
        self.environment = environment
        self.batch_size = batch_size
        self.flush_interval = flush_interval

        if not self.api_key:
            logger.warning("AgentBase: No API key. Set AB_API_KEY or pass api_key to init().")

        self._queue = queue.Queue(maxsize=10000)
        self._flush_thread = None
        self._running = False
        self._http = httpx.Client(
            base_url=self.endpoint,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "X-Agent-Name": self.agent_name,
                "X-Environment": self.environment,
                "X-SDK-Version": "python-0.1.0",
            },
            timeout=10.0,
        )

        if debug:
            logging.basicConfig(level=logging.DEBUG)
            logger.setLevel(logging.DEBUG)

    def start(self):
        self._running = True
        self._flush_thread = threading.Thread(target=self._flush_loop, daemon=True)
        self._flush_thread.start()

    def shutdown(self):
        self._running = False
        self._flush()
        if self._flush_thread:
            self._flush_thread.join(timeout=5.0)
        self._http.close()

    def enqueue_trace(self, trace):
        try:
            self._queue.put_nowait(("trace", trace.to_dict()))
        except queue.Full:
            logger.warning("AgentBase: Queue full, dropping trace")

    def enqueue_span(self, span):
        try:
            self._queue.put_nowait(("span", span.to_dict()))
        except queue.Full:
            logger.warning("AgentBase: Queue full, dropping span")

    def check_heal(self, span):
        if not self.auto_heal or span.status != "error":
            return None
        try:
            resp = self._http.post("/v1/heal/check", json={
                "agent_name": self.agent_name, "span": span.to_dict(),
            })
            if resp.status_code == 200:
                data = resp.json()
                if data.get("should_heal"):
                    return data
        except Exception as e:
            logger.debug("AgentBase: Heal check failed: %s", e)
        return None

    def _flush_loop(self):
        while self._running:
            time.sleep(self.flush_interval)
            self._flush()

    def _flush(self):
        traces = []
        spans = []
        while not self._queue.empty() and (len(traces) + len(spans)) < self.batch_size:
            try:
                item_type, data = self._queue.get_nowait()
                if item_type == "trace":
                    traces.append(data)
                else:
                    spans.append(data)
            except queue.Empty:
                break

        if not traces and not spans:
            return

        payload = {}
        if traces:
            payload["traces"] = traces
        if spans:
            payload["spans"] = spans

        try:
            resp = self._http.post("/v1/ingest", json=payload)
            if resp.status_code != 200:
                logger.warning("AgentBase: Ingest failed (%d)", resp.status_code)
        except Exception as e:
            logger.warning("AgentBase: Flush failed: %s", e)


def generate_id():
    return uuid.uuid4().hex[:16]
