"""Core Tokengrip client with batching and auto-flush."""

from __future__ import annotations

import atexit
import json
import logging
import threading
import urllib.error
import urllib.request
from typing import Any, Dict, List, Optional

from tokengrip.types import TrackUsagePayload

logger = logging.getLogger("tokengrip")

_SDK_VERSION = "0.1.0"


class Tokengrip:
    """Buffer usage records and send them in batches to the Tokengrip API.

    Example::

        tg = Tokengrip(api_key="tq_live_abc123", project_id="my-project")
        tg.track(provider="openai", model="gpt-4.1",
                 inputTokens=500, outputTokens=150)
        tg.shutdown()
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://tokengrip.com",
        project_id: str = "",
        agent_name: str = "",
        debug: bool = False,
        batch_size: int = 10,
        flush_interval_ms: int = 5000,
    ) -> None:
        if not api_key:
            raise ValueError(
                "Tokengrip: api_key is required. "
                "Get yours at https://tokengrip.com/settings/api-keys"
            )

        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._project_id = project_id
        self._agent_name = agent_name
        self._debug = debug
        self._batch_size = batch_size
        self._flush_interval_s = flush_interval_ms / 1000.0

        self._buffer: List[Dict[str, Any]] = []
        self._lock = threading.Lock()
        self._is_flushing = False
        self._is_shutdown = False

        self._timer: Optional[threading.Timer] = None
        if self._flush_interval_s > 0:
            self._schedule_flush()

        atexit.register(self.shutdown)

    def _schedule_flush(self) -> None:
        if self._is_shutdown:
            return
        self._timer = threading.Timer(self._flush_interval_s, self._periodic_flush)
        self._timer.daemon = True
        self._timer.start()

    def _periodic_flush(self) -> None:
        self.flush()
        self._schedule_flush()

    def track(
        self,
        payload: Optional[TrackUsagePayload] = None,
        **kwargs: Any,
    ) -> None:
        """Buffer a single usage record.

        Accepts a ``TrackUsagePayload`` dict, keyword arguments, or both.
        Never raises — errors are logged at debug level when ``debug=True``.
        """
        if self._is_shutdown:
            self._log("warn", "track() called after shutdown — record dropped")
            return

        try:
            record: Dict[str, Any] = dict(payload) if payload else {}
            record.update(kwargs)

            # Compute totalTokens if absent
            input_t = record.get("inputTokens", 0)
            output_t = record.get("outputTokens", 0)
            if "totalTokens" not in record:
                record["totalTokens"] = input_t + output_t

            # Apply defaults from config
            if not record.get("projectId") and self._project_id:
                record["projectId"] = self._project_id
            if not record.get("agentName") and self._agent_name:
                record["agentName"] = self._agent_name

            # Strip falsy optional fields for a clean payload
            for key in ("projectId", "agentName", "taskType", "metadata"):
                if key in record and not record[key]:
                    del record[key]
            if "latencyMs" in record and record["latencyMs"] is None:
                del record["latencyMs"]

            with self._lock:
                self._buffer.append(record)
                buffer_len = len(self._buffer)

            self._log(
                "debug",
                f"Buffered: {record.get('provider')}/{record.get('model')} "
                f"({input_t}+{output_t} tokens)",
            )

            if buffer_len >= self._batch_size:
                self.flush()

        except Exception as e:
            self._log("warn", f"Failed to buffer record: {e}")

    def flush(self) -> None:
        """Force-flush all buffered records to the API. Thread-safe."""
        with self._lock:
            if not self._buffer or self._is_flushing:
                return
            self._is_flushing = True
            records = self._buffer[:]
            self._buffer.clear()

        try:
            self._send_batch(records)
        except Exception as e:
            self._log("warn", f"Flush failed: {e}")
        finally:
            with self._lock:
                self._is_flushing = False

    def shutdown(self) -> None:
        """Flush remaining records and stop the auto-flush timer."""
        if self._is_shutdown:
            return
        self._is_shutdown = True

        if self._timer is not None:
            self._timer.cancel()
            self._timer = None

        with self._lock:
            self._is_flushing = False

        self.flush()
        self._log("debug", "Tokengrip client shut down")

    def _send_batch(self, records: List[Dict[str, Any]]) -> None:
        if not records:
            return

        url = f"{self._base_url}/api/v1/track"
        self._log("debug", f"Sending {len(records)} record(s) to {url}")

        data = json.dumps(records).encode("utf-8")
        req = urllib.request.Request(
            url,
            data=data,
            method="POST",
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self._api_key}",
                "User-Agent": f"tokengrip-python/{_SDK_VERSION}",
            },
        )

        try:
            with urllib.request.urlopen(req, timeout=10) as resp:
                self._log(
                    "debug", f"Sent {len(records)} record(s), status {resp.status}"
                )
        except urllib.error.HTTPError as e:
            body = ""
            try:
                body = e.read().decode("utf-8", errors="replace")
            except Exception:
                pass
            self._log("warn", f"API returned {e.code}: {body}")
        except Exception as e:
            self._log("warn", f"Failed to send batch: {e}")

    def _log(self, level: str, message: str) -> None:
        if not self._debug:
            return
        prefix = "[tokengrip]"
        if level == "warn":
            logger.warning("%s %s", prefix, message)
        else:
            logger.debug("%s %s", prefix, message)
