"""
Core Syntropy client — sends traces to the ingestion service.
"""

from __future__ import annotations

import asyncio
import atexit
import logging
import threading
import time
from dataclasses import dataclass, field
from queue import Empty, Queue
from typing import Any, Optional

import httpx

logger = logging.getLogger("syntropy")


@dataclass
class SyntropyConfig:
    """Configuration for the Syntropy client."""

    api_key: str
    """API key for authentication (starts with syn_)."""

    base_url: str = "http://localhost:4000"
    """Base URL of the Syntropy ingestion service."""

    agent_id: str = ""
    """Default agent ID attached to all traces."""

    flush_interval: float = 2.0
    """How often to flush the trace buffer (seconds)."""

    batch_size: int = 50
    """Maximum traces per batch request."""

    max_queue_size: int = 10_000
    """Maximum number of buffered traces before dropping."""

    timeout: float = 10.0
    """HTTP request timeout (seconds)."""

    enabled: bool = True
    """Set to False to disable all tracing (dry-run mode)."""

    metadata: dict[str, Any] = field(default_factory=dict)
    """Default metadata attached to all traces."""


@dataclass
class TraceResult:
    """Result of a trace submission."""

    success: bool
    status: str = ""
    error: Optional[str] = None
    trace_id: Optional[str] = None


class Syntropy:
    """
    Syntropy AI Observability client.

    Buffers traces and sends them in batches to the ingestion service.
    Thread-safe. Flushes automatically on exit.
    """

    _instance: Optional["Syntropy"] = None

    @classmethod
    def init(
        cls,
        api_key: str = "",
        base_url: str = "http://localhost:4000",
        agent_id: str = "",
        *,
        auto_instrument: bool = True,
        **kwargs: Any,
    ) -> "Syntropy":
        """
        Initialize Syntropy and auto-instrument detected AI providers.

        Creates a singleton instance. Subsequent calls return the same instance.

        Usage:
            from syntropy import Syntropy
            syn = Syntropy.init(api_key="syn_...", agent_id="my-agent")
            # OpenAI, Anthropic, Mistral, etc. are now auto-traced.

        Args:
            api_key: Your Syntropy API key.
            base_url: Ingestion service URL.
            agent_id: Default agent identifier.
            auto_instrument: If True, auto-patch detected providers (default True).

        Returns:
            The Syntropy singleton instance.
        """
        if cls._instance is not None:
            return cls._instance

        instance = cls(
            api_key=api_key,
            base_url=base_url,
            agent_id=agent_id,
            **kwargs,
        )
        cls._instance = instance

        if auto_instrument:
            from syntropy.instrument import instrument
            instrument(instance)

        return instance

    @classmethod
    def get_instance(cls) -> "Syntropy":
        """Return the singleton created by init(). Raises if not initialized."""
        if cls._instance is None:
            raise RuntimeError(
                "Syntropy.init() has not been called. "
                "Call Syntropy.init(api_key=...) first."
            )
        return cls._instance

    def __init__(
        self,
        api_key: str = "",
        base_url: str = "http://localhost:4000",
        agent_id: str = "",
        *,
        config: Optional[SyntropyConfig] = None,
        **kwargs: Any,
    ):
        if config:
            self._config = config
        else:
            self._config = SyntropyConfig(
                api_key=api_key or kwargs.get("api_key", ""),
                base_url=base_url,
                agent_id=agent_id,
                **{
                    k: v
                    for k, v in kwargs.items()
                    if k in SyntropyConfig.__dataclass_fields__
                },
            )

        if not self._config.api_key:
            raise ValueError("Syntropy API key is required")

        self._queue: Queue[dict[str, Any]] = Queue(
            maxsize=self._config.max_queue_size
        )
        self._http = httpx.Client(
            base_url=self._config.base_url,
            headers={
                "X-API-Key": self._config.api_key,
                "Content-Type": "application/json",
            },
            timeout=self._config.timeout,
        )
        self._stopped = False
        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True
        )
        self._flush_thread.start()
        atexit.register(self.shutdown)

    # ---- Public API ----

    def trace(
        self,
        prompt: str,
        response: str,
        model: str,
        *,
        agent_id: str = "",
        tokens_in: int = 0,
        tokens_out: int = 0,
        latency_ms: int = 0,
        status: str = "success",
        metadata: Optional[dict[str, Any]] = None,
    ) -> TraceResult:
        """
        Record a single LLM interaction trace.

        Returns immediately — the trace is buffered and sent asynchronously.
        """
        if not self._config.enabled:
            return TraceResult(success=True, status="disabled")

        payload = {
            "agent_id": agent_id or self._config.agent_id,
            "prompt": prompt,
            "response": response,
            "model": model,
            "tokens_in": tokens_in,
            "tokens_out": tokens_out,
            "latency_ms": latency_ms,
            "status": status,
            "metadata": {
                **self._config.metadata,
                **(metadata or {}),
            },
        }

        if not payload["agent_id"]:
            return TraceResult(
                success=False, error="agent_id is required"
            )

        try:
            self._queue.put_nowait(payload)
            return TraceResult(success=True, status="queued")
        except Exception:
            logger.warning("Syntropy trace queue full — dropping trace")
            return TraceResult(success=False, error="Queue full")

    def trace_sync(
        self,
        prompt: str,
        response: str,
        model: str,
        **kwargs: Any,
    ) -> TraceResult:
        """
        Send a single trace synchronously (blocking).

        Use this when you need confirmation the trace was received.
        """
        if not self._config.enabled:
            return TraceResult(success=True, status="disabled")

        payload = {
            "agent_id": kwargs.get("agent_id") or self._config.agent_id,
            "prompt": prompt,
            "response": response,
            "model": model,
            "tokens_in": kwargs.get("tokens_in", 0),
            "tokens_out": kwargs.get("tokens_out", 0),
            "latency_ms": kwargs.get("latency_ms", 0),
            "status": kwargs.get("status", "success"),
            "metadata": {
                **self._config.metadata,
                **(kwargs.get("metadata") or {}),
            },
        }

        try:
            resp = self._http.post("/v1/ingest", json=payload)
            if resp.status_code < 300:
                data = resp.json()
                return TraceResult(
                    success=True,
                    status=data.get("status", "accepted"),
                    trace_id=data.get("id"),
                )
            return TraceResult(
                success=False,
                error=f"HTTP {resp.status_code}: {resp.text}",
            )
        except httpx.HTTPError as e:
            return TraceResult(success=False, error=str(e))

    def flush(self) -> int:
        """Flush all buffered traces immediately. Returns number sent."""
        return self._flush_batch()

    def shutdown(self) -> None:
        """Flush remaining traces and close the client."""
        if self._stopped:
            return
        self._stopped = True
        self._flush_batch()
        self._http.close()

    # ---- Batch endpoint support ----

    def trace_batch(
        self, traces: list[dict[str, Any]]
    ) -> TraceResult:
        """
        Send multiple traces in one request via /v1/ingest/batch.
        """
        if not self._config.enabled:
            return TraceResult(success=True, status="disabled")

        for t in traces:
            if not t.get("agent_id"):
                t["agent_id"] = self._config.agent_id
            t.setdefault("metadata", {}).update(self._config.metadata)

        try:
            resp = self._http.post(
                "/v1/ingest/batch", json={"traces": traces}
            )
            if resp.status_code < 300:
                return TraceResult(success=True, status="accepted")
            return TraceResult(
                success=False,
                error=f"HTTP {resp.status_code}: {resp.text}",
            )
        except httpx.HTTPError as e:
            return TraceResult(success=False, error=str(e))

    # ---- Internal ----

    def _flush_loop(self) -> None:
        """Background thread that flushes the queue periodically."""
        while not self._stopped:
            time.sleep(self._config.flush_interval)
            try:
                self._flush_batch()
            except Exception as e:
                logger.warning("Syntropy flush error: %s", e)

    def _flush_batch(self) -> int:
        """Drain queue and send as batch. Returns count sent."""
        traces: list[dict[str, Any]] = []
        while len(traces) < self._config.batch_size:
            try:
                traces.append(self._queue.get_nowait())
            except Empty:
                break

        if not traces:
            return 0

        try:
            resp = self._http.post(
                "/v1/ingest/batch", json={"traces": traces}
            )
            if resp.status_code >= 300:
                logger.warning(
                    "Syntropy batch send failed: HTTP %d", resp.status_code
                )
        except httpx.HTTPError as e:
            logger.warning("Syntropy batch send error: %s", e)

        return len(traces)


# ================================================================
# Async variant
# ================================================================


class AsyncSyntropy:
    """
    Async Syntropy client for use with asyncio / FastAPI / etc.
    """

    def __init__(
        self,
        api_key: str = "",
        base_url: str = "http://localhost:4000",
        agent_id: str = "",
        *,
        config: Optional[SyntropyConfig] = None,
    ):
        if config:
            self._config = config
        else:
            self._config = SyntropyConfig(
                api_key=api_key, base_url=base_url, agent_id=agent_id
            )

        if not self._config.api_key:
            raise ValueError("Syntropy API key is required")

        self._http = httpx.AsyncClient(
            base_url=self._config.base_url,
            headers={
                "X-API-Key": self._config.api_key,
                "Content-Type": "application/json",
            },
            timeout=self._config.timeout,
        )
        self._buffer: list[dict[str, Any]] = []
        self._lock = asyncio.Lock()

    async def trace(
        self,
        prompt: str,
        response: str,
        model: str,
        *,
        agent_id: str = "",
        tokens_in: int = 0,
        tokens_out: int = 0,
        latency_ms: int = 0,
        status: str = "success",
        metadata: Optional[dict[str, Any]] = None,
    ) -> TraceResult:
        """Record a trace asynchronously."""
        if not self._config.enabled:
            return TraceResult(success=True, status="disabled")

        payload = {
            "agent_id": agent_id or self._config.agent_id,
            "prompt": prompt,
            "response": response,
            "model": model,
            "tokens_in": tokens_in,
            "tokens_out": tokens_out,
            "latency_ms": latency_ms,
            "status": status,
            "metadata": {**self._config.metadata, **(metadata or {})},
        }

        try:
            resp = await self._http.post("/v1/ingest", json=payload)
            if resp.status_code < 300:
                data = resp.json()
                return TraceResult(
                    success=True,
                    status=data.get("status", "accepted"),
                    trace_id=data.get("id"),
                )
            return TraceResult(
                success=False,
                error=f"HTTP {resp.status_code}: {resp.text}",
            )
        except httpx.HTTPError as e:
            return TraceResult(success=False, error=str(e))

    async def trace_batch(
        self, traces: list[dict[str, Any]]
    ) -> TraceResult:
        """Send multiple traces in one request."""
        if not self._config.enabled:
            return TraceResult(success=True, status="disabled")

        for t in traces:
            if not t.get("agent_id"):
                t["agent_id"] = self._config.agent_id

        try:
            resp = await self._http.post(
                "/v1/ingest/batch", json={"traces": traces}
            )
            if resp.status_code < 300:
                return TraceResult(success=True, status="accepted")
            return TraceResult(
                success=False,
                error=f"HTTP {resp.status_code}: {resp.text}",
            )
        except httpx.HTTPError as e:
            return TraceResult(success=False, error=str(e))

    async def flush(self) -> int:
        """
        Flush all buffered traces immediately.

        This is critical for serverless environments (AWS Lambda, etc.)
        where the container freezes immediately after the response is sent.
        Call ``await client.flush()`` at the end of your handler.

        Returns:
            Number of traces sent.
        """
        async with self._lock:
            if not self._buffer:
                return 0
            batch = list(self._buffer)
            self._buffer.clear()

        try:
            resp = await self._http.post(
                "/v1/ingest/batch", json={"traces": batch}
            )
            if resp.status_code >= 300:
                logger.warning(
                    "Syntropy async flush failed: HTTP %d", resp.status_code
                )
        except httpx.HTTPError as e:
            logger.warning("Syntropy async flush error: %s", e)

        return len(batch)

    async def shutdown(self) -> None:
        """Flush remaining traces and close the async HTTP client."""
        await self.flush()
        await self._http.aclose()

    async def close(self) -> None:
        """Close the async HTTP client."""
        await self._http.aclose()
