"""
vela - Lightweight AI agent tracing library.

Usage:
    import vela

    @vela.trace
    def my_agent(prompt):
        ...

    # Optional: custom agent name and per-call cost
    @vela.trace(name="summarizer", cost_usd=0.0005)
    def summarize(text):
        ...

    # Optional: connect to Vela (call before running pipeline)
    vela.init(api_key="vela_abc123", cost_limit_usd=0.10, loop_limit=5)

    # Optional: start a new session (new pipeline run)
    vela.new_session(label="my run")

    # Optional: mark session complete when done
    vela.end_session()
"""

import json
import time
import uuid
import inspect
import functools
import datetime

from . import transport


# ── Model pricing (per 1 000 tokens: input, output) ──────────────────────────

MODEL_COSTS: dict[str, tuple[float, float]] = {
    "gpt-4o":             (0.0025,   0.01),
    "gpt-4o-mini":        (0.000150, 0.000600),
    "gpt-4-turbo":        (0.01,     0.03),
    "gpt-3.5-turbo":      (0.0005,   0.0015),
    "claude-opus-4-6":    (0.015,    0.075),
    "claude-sonnet-4-6":  (0.003,    0.015),
    "claude-haiku-4-5":   (0.00025,  0.00125),
    "gemini-1.5-pro":     (0.00125,  0.005),
    "gemini-1.5-flash":   (0.000075, 0.0003),
}


# ── Exceptions ───────────────────────────────────────────────────────────────

class VelaHaltException(Exception):
    """Raised when a session is halted due to a governance limit being exceeded."""


# ── Credentials / limits ─────────────────────────────────────────────────────

_initialized: bool = False
_cost_limit_usd: float | None = None
_loop_limit: int | None = None


def init(
    api_key: str,
    ingest_url: str | None = None,
    cost_limit_usd: float | None = None,
    loop_limit: int | None = None,
) -> None:
    """Connect Vela. Call once before running your pipeline."""
    global _initialized, _cost_limit_usd, _loop_limit
    transport.init(api_key, ingest_url)
    _cost_limit_usd = cost_limit_usd
    _loop_limit = loop_limit
    _initialized = True


# ── Session management ──────────────────────────────────────────────────────

_session_id: str = str(uuid.uuid4())
_trace_file: str = "vela_traces.json"
_sequence: int = 0
_session_total_cost: float = 0.0
_session_agent_counts: dict[str, int] = {}


def new_session(label: str | None = None) -> str:
    """Start a new session (new pipeline run) and return its ID."""
    global _session_id, _sequence, _session_total_cost, _session_agent_counts
    _session_id = str(uuid.uuid4())
    _sequence = 0
    _session_total_cost = 0.0
    _session_agent_counts = {}
    if _initialized:
        transport.create_session(
            _session_id,
            label=label,
            cost_limit_usd=_cost_limit_usd,
            loop_limit=_loop_limit,
        )
    return _session_id


def end_session() -> None:
    """Mark the current session as complete."""
    if _initialized:
        transport.complete_session(_session_id)


def get_session_id() -> str:
    """Return the current session ID."""
    return _session_id


def set_trace_file(path: str) -> None:
    """Override the default trace output file path."""
    global _trace_file
    _trace_file = path


# ── Internal helpers ────────────────────────────────────────────────────────

def _serializable(obj):
    """Best-effort conversion to a JSON-serializable value."""
    try:
        json.dumps(obj)
        return obj
    except (TypeError, ValueError):
        return repr(obj)


def _build_event(agent_name: str, args: tuple, kwargs: dict) -> dict:
    global _sequence
    _sequence += 1
    return {
        "session_id": _session_id,
        "agent": agent_name,
        "created_at": datetime.datetime.utcnow().isoformat() + "Z",
        "sequence": _sequence,
        "input": {
            "args": [_serializable(a) for a in args],
            "kwargs": {k: _serializable(v) for k, v in kwargs.items()},
        },
        "output": None,
        "error": None,
        "latency_ms": None,
    }


def _elapsed_ms(t0: float) -> float:
    return round((time.perf_counter() - t0) * 1000, 3)


def _apply_cost_metadata(
    event: dict,
    model: str | None,
    token_input: int | None,
    token_output: int | None,
    cost_usd: float | None,
) -> None:
    """Populate model/token/cost fields on a trace event in-place."""
    if model is not None:
        event["model"] = model
    if token_input is not None:
        event["token_input"] = token_input
    if token_output is not None:
        event["token_output"] = token_output

    if cost_usd is not None:
        event["cost_usd"] = cost_usd
    elif model is not None and token_input is not None and token_output is not None:
        prices = MODEL_COSTS.get(model)
        if prices:
            event["cost_usd"] = (
                token_input  / 1000 * prices[0]
                + token_output / 1000 * prices[1]
            )


def _write_trace(event: dict) -> None:
    """Print the trace event to console, append to file, and send to Vela."""
    line = json.dumps(event)
    print(f"[vela] {line}")
    with open(_trace_file, "a") as f:
        f.write(line + "\n")
    if _initialized:
        transport.send_trace(event)


def _do_halt(reason: str, message: str) -> None:
    """Halt the current session and raise VelaHaltException."""
    if _initialized:
        transport.halt_session(_session_id, reason)
    raise VelaHaltException(message)


def _check_limits(agent_name: str, event: dict) -> None:
    """Update per-session tracking and check governance limits."""
    global _session_total_cost, _session_agent_counts

    cost = event.get("cost_usd") or 0.0
    _session_total_cost += cost
    _session_agent_counts[agent_name] = _session_agent_counts.get(agent_name, 0) + 1

    if _cost_limit_usd is not None and _session_total_cost > _cost_limit_usd:
        _do_halt(
            "cost_limit_exceeded",
            f"Session halted: cost limit of ${_cost_limit_usd} exceeded "
            f"(session total: ${_session_total_cost:.6f})",
        )

    if _loop_limit is not None and _session_agent_counts[agent_name] > _loop_limit:
        count = _session_agent_counts[agent_name]
        _do_halt(
            "loop_detected",
            f"Session halted: agent '{agent_name}' fired {count} times, "
            f"exceeding loop limit of {_loop_limit}",
        )


# ── Decorator ───────────────────────────────────────────────────────────────

def trace(
    fn=None,
    *,
    name: str = None,
    model: str | None = None,
    token_input: int | None = None,
    token_output: int | None = None,
    cost_usd: float | None = None,
):
    """
    Decorator that traces an AI agent function call.

    Captures: agent name, input args/kwargs, output, latency (ms),
    UTC timestamp, session ID, and any raised exception.

    Supports both sync and async functions.

    Args:
        name:         Override the agent name shown in the dashboard.
        model:        Model name (e.g. "gpt-4o", "claude-sonnet-4-6").
        token_input:  Input tokens consumed by this agent call.
        token_output: Output tokens produced by this agent call.
        cost_usd:     Explicit cost in USD. If omitted and model+tokens are
                      provided, cost is inferred from MODEL_COSTS automatically.
    """
    if fn is None:
        return functools.partial(
            trace,
            name=name,
            model=model,
            token_input=token_input,
            token_output=token_output,
            cost_usd=cost_usd,
        )

    agent_name = name or fn.__qualname__

    if inspect.iscoroutinefunction(fn):
        @functools.wraps(fn)
        async def async_wrapper(*args, **kwargs):
            event = _build_event(agent_name, args, kwargs)
            _apply_cost_metadata(event, model, token_input, token_output, cost_usd)
            t0 = time.perf_counter()
            try:
                result = await fn(*args, **kwargs)
            except VelaHaltException:
                raise
            except Exception as exc:
                event["error"] = repr(exc)
                event["latency_ms"] = _elapsed_ms(t0)
                _write_trace(event)
                if _initialized:
                    transport.error_session(_session_id)
                raise
            event["output"] = _serializable(result)
            event["latency_ms"] = _elapsed_ms(t0)
            _write_trace(event)
            _check_limits(agent_name, event)
            return result

        return async_wrapper

    @functools.wraps(fn)
    def sync_wrapper(*args, **kwargs):
        event = _build_event(agent_name, args, kwargs)
        _apply_cost_metadata(event, model, token_input, token_output, cost_usd)
        t0 = time.perf_counter()
        try:
            result = fn(*args, **kwargs)
        except VelaHaltException:
            raise
        except Exception as exc:
            event["error"] = repr(exc)
            event["latency_ms"] = _elapsed_ms(t0)
            _write_trace(event)
            if _initialized:
                transport.error_session(_session_id)
            raise
        event["output"] = _serializable(result)
        event["latency_ms"] = _elapsed_ms(t0)
        _write_trace(event)
        _check_limits(agent_name, event)
        return result

    return sync_wrapper
