"""VelenAI client — handles auth, batching, and delivery."""

from __future__ import annotations

import asyncio
import functools
import hashlib
import hmac
import inspect
import json
import logging
import os
import threading
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Callable

import httpx

logger = logging.getLogger("velenai")


@dataclass
class TraceContext:
    """Mutable context for a tracked agent run."""

    agent_id: str
    run_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    success: bool = True
    tokens: int | None = None
    cost_usd: float | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    _start: float = field(default_factory=time.monotonic)

    @property
    def elapsed_ms(self) -> int:
        return int((time.monotonic() - self._start) * 1000)


class VelenAI:
    """VelenAI SDK client.

    Args:
        api_key: HMAC API key (or VELENAI_API_KEY env var)
        api_secret: HMAC API secret (or VELENAI_API_SECRET env var)
        host: VelenAI server URL (or VELENAI_HOST env var)
        flush_interval: Seconds between batch flushes (default: 5)
        max_batch_size: Max events per flush (default: 50)
        enabled: Set False to disable all tracking (default: True)
    """

    def __init__(
        self,
        api_key: str | None = None,
        api_secret: str | None = None,
        host: str | None = None,
        flush_interval: float = 5.0,
        max_batch_size: int = 50,
        enabled: bool = True,
    ):
        self._api_key = api_key or os.environ.get("VELENAI_API_KEY", "")
        self._api_secret = api_secret or os.environ.get("VELENAI_API_SECRET", "")
        self._host = (
            host or os.environ.get("VELENAI_HOST", "http://localhost:8000")
        ).rstrip("/")
        self._enabled = enabled and bool(self._api_key and self._api_secret)
        self._flush_interval = flush_interval
        self._max_batch_size = max_batch_size

        self._queue: list[dict] = []
        self._lock = threading.Lock()
        self._client: httpx.Client | None = None
        self._async_client: httpx.AsyncClient | None = None

        # Start background flusher
        if self._enabled:
            self._flusher = threading.Thread(target=self._flush_loop, daemon=True)
            self._flusher.start()

    def _get_client(self) -> httpx.Client:
        if self._client is None or self._client.is_closed:
            self._client = httpx.Client(timeout=httpx.Timeout(10.0))
        return self._client

    def _get_async_client(self) -> httpx.AsyncClient:
        if self._async_client is None or self._async_client.is_closed:
            self._async_client = httpx.AsyncClient(timeout=httpx.Timeout(10.0))
        return self._async_client

    def _sign(self, body: bytes, timestamp: int) -> str:
        message = body + str(timestamp).encode()
        return hmac.new(self._api_secret.encode(), message, hashlib.sha256).hexdigest()

    def _enqueue(self, event: dict):
        if not self._enabled:
            return
        with self._lock:
            self._queue.append(event)
            if len(self._queue) >= self._max_batch_size:
                self._flush_sync()

    def _build_event(
        self,
        agent_id: str,
        success: bool = True,
        latency_ms: int = 0,
        tokens: int | None = None,
        cost_usd: float | None = None,
        metadata: dict | None = None,
    ) -> dict:
        return {
            "agentId": agent_id,
            "runId": str(uuid.uuid4()),
            "status": "success" if success else "failure",
            "latencyMs": latency_ms,
            "totalTokens": tokens,
            "costUsd": cost_usd,
            "metadata": metadata or {},
        }

    def _send_event(self, event: dict) -> bool:
        """Send a single event synchronously. Returns True on success."""
        body = json.dumps(event).encode()
        timestamp = int(time.time())
        headers = {
            "Content-Type": "application/json",
            "X-VelenAI-Key": self._api_key,
            "X-VelenAI-Timestamp": str(timestamp),
            "X-VelenAI-Signature": self._sign(body, timestamp),
        }
        try:
            resp = self._get_client().post(
                f"{self._host}/telemetry", content=body, headers=headers
            )
            return resp.status_code == 200
        except Exception as e:
            logger.debug(f"VelenAI send failed: {e}")
            return False

    async def _send_event_async(self, event: dict) -> bool:
        """Send a single event asynchronously."""
        body = json.dumps(event).encode()
        timestamp = int(time.time())
        headers = {
            "Content-Type": "application/json",
            "X-VelenAI-Key": self._api_key,
            "X-VelenAI-Timestamp": str(timestamp),
            "X-VelenAI-Signature": self._sign(body, timestamp),
        }
        try:
            resp = await self._get_async_client().post(
                f"{self._host}/telemetry", content=body, headers=headers
            )
            return resp.status_code == 200
        except Exception as e:
            logger.debug(f"VelenAI async send failed: {e}")
            return False

    def _flush_sync(self):
        """Flush queued events. Called with lock held or from flusher thread."""
        events = self._queue[:]
        self._queue.clear()
        for event in events:
            self._send_event(event)

    def _flush_loop(self):
        """Background thread that flushes the queue periodically."""
        while True:
            time.sleep(self._flush_interval)
            with self._lock:
                if self._queue:
                    self._flush_sync()

    # ── Public API ────────────────────────────────────────────────────

    def emit(
        self,
        agent_id: str,
        success: bool = True,
        latency_ms: int = 0,
        tokens: int | None = None,
        cost_usd: float | None = None,
        metadata: dict | None = None,
    ):
        """Emit a single agent run event (queued for batch delivery)."""
        event = self._build_event(
            agent_id, success, latency_ms, tokens, cost_usd, metadata
        )
        self._enqueue(event)

    async def emit_async(
        self,
        agent_id: str,
        success: bool = True,
        latency_ms: int = 0,
        tokens: int | None = None,
        cost_usd: float | None = None,
        metadata: dict | None = None,
    ):
        """Emit a single agent run event asynchronously (sent immediately)."""
        if not self._enabled:
            return
        event = self._build_event(
            agent_id, success, latency_ms, tokens, cost_usd, metadata
        )
        await self._send_event_async(event)

    @contextmanager
    def trace(self, agent_id: str, **initial_metadata):
        """Context manager for tracking an agent run.

        Usage:
            with vai.trace("my-agent") as t:
                result = do_work()
                t.success = True
                t.tokens = 150
        """
        ctx = TraceContext(agent_id=agent_id, metadata=initial_metadata)
        try:
            yield ctx
        except Exception:
            ctx.success = False
            raise
        finally:
            self.emit(
                agent_id=ctx.agent_id,
                success=ctx.success,
                latency_ms=ctx.elapsed_ms,
                tokens=ctx.tokens,
                cost_usd=ctx.cost_usd,
                metadata=ctx.metadata,
            )

    def track(self, agent_id: str, **default_metadata) -> Callable:
        """Decorator for tracking agent function calls.

        Works with both sync and async functions.

        Usage:
            @vai.track("my-agent")
            def run_agent(query: str):
                return process(query)

            @vai.track("my-async-agent")
            async def run_agent(query: str):
                return await process(query)
        """

        def decorator(fn: Callable) -> Callable:
            if inspect.iscoroutinefunction(fn):

                @functools.wraps(fn)
                async def async_wrapper(*args, **kwargs):
                    start = time.monotonic()
                    success = True
                    try:
                        result = await fn(*args, **kwargs)
                        return result
                    except Exception:
                        success = False
                        raise
                    finally:
                        await self.emit_async(
                            agent_id=agent_id,
                            success=success,
                            latency_ms=int((time.monotonic() - start) * 1000),
                            metadata=default_metadata,
                        )

                return async_wrapper
            else:

                @functools.wraps(fn)
                def sync_wrapper(*args, **kwargs):
                    start = time.monotonic()
                    success = True
                    try:
                        result = fn(*args, **kwargs)
                        return result
                    except Exception:
                        success = False
                        raise
                    finally:
                        self.emit(
                            agent_id=agent_id,
                            success=success,
                            latency_ms=int((time.monotonic() - start) * 1000),
                            metadata=default_metadata,
                        )

                return sync_wrapper
            return decorator

        return decorator

    def flush(self):
        """Manually flush all queued events. Call before shutdown."""
        with self._lock:
            if self._queue:
                self._flush_sync()

    def shutdown(self):
        """Flush and close connections."""
        self.flush()
        if self._client:
            self._client.close()
        if self._async_client:
            try:
                asyncio.get_event_loop().run_until_complete(self._async_client.aclose())
            except Exception:
                pass
