"""Lightweight wrapper around the Grok (xAI) chat-completions API.

All reasoning agents call this module instead of making raw HTTP requests.
The wrapper keeps token usage low by enforcing a system prompt and optional
JSON-mode output.

Configuration
-------------
Set the ``GROK_API_KEY`` environment variable before running SecAgent.
Optionally set ``GROK_MODEL`` (defaults to ``grok-4-1-fast-reasoning``).
"""

from __future__ import annotations

import json
import logging
import os
from typing import Any

import httpx

from secagent.models import TokenUsage

logger = logging.getLogger(__name__)

_DEFAULT_MODEL = "grok-4-1-fast-reasoning"         # for fixer, validator, ai-scanner
_DEFAULT_FAST_MODEL = "grok-4-1-fast-non-reasoning"  # for planner, verifier
_API_URL = "https://api.x.ai/v1/chat/completions"
_TIMEOUT = 60  # seconds


class GrokClient:
    """Thin, reusable client for the Grok chat-completions endpoint."""

    def __init__(
        self,
        api_key: str | None = None,
        model: str | None = None,
        temperature: float = 0.2,
    ) -> None:
        self.api_key = api_key or os.getenv("GROK_API_KEY", "")
        self.model = model or os.getenv("GROK_MODEL", _DEFAULT_MODEL)
        self.temperature = temperature
        self._session_usage = TokenUsage()

        if not self.api_key:
            logger.warning(
                "GROK_API_KEY is not set — LLM agent calls will be skipped. "
                "Set the environment variable to enable AI-powered reasoning."
            )

    # ------------------------------------------------------------------
    # Usage tracking
    # ------------------------------------------------------------------

    @property
    def session_usage(self) -> TokenUsage:
        """Cumulative token usage since last reset (or client creation)."""
        return self._session_usage

    def reset_usage(self) -> TokenUsage:
        """Return a snapshot of current usage and reset the counter to zero."""
        snapshot = TokenUsage(
            prompt_tokens=self._session_usage.prompt_tokens,
            completion_tokens=self._session_usage.completion_tokens,
            reasoning_tokens=self._session_usage.reasoning_tokens,
            calls=self._session_usage.calls,
        )
        self._session_usage = TokenUsage()
        return snapshot

    # ------------------------------------------------------------------
    # Public helpers
    # ------------------------------------------------------------------

    def chat(
        self,
        system_prompt: str,
        user_prompt: str,
        *,
        json_mode: bool = True,
        max_tokens: int = 2048,
    ) -> dict[str, Any] | str:
        """Send a chat-completion request and return the parsed response.

        Parameters
        ----------
        system_prompt:
            Instruction that sets agent behaviour.
        user_prompt:
            The actual task / context payload.
        json_mode:
            When *True* the API is asked to return valid JSON and the
            response is parsed before being returned.
        max_tokens:
            Upper token limit for the response.

        Returns
        -------
        dict | str
            Parsed JSON ``dict`` when *json_mode* is enabled, otherwise
            the raw assistant text.
        """
        if not self.api_key:
            logger.info("Skipping Grok call (no API key configured).")
            return {} if json_mode else ""

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        payload: dict[str, Any] = {
            "model": self.model,
            "temperature": self.temperature,
            "max_tokens": max_tokens,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
        }

        if json_mode:
            payload["response_format"] = {"type": "json_object"}

        try:
            response = httpx.post(
                _API_URL,
                headers=headers,
                json=payload,
                timeout=_TIMEOUT,
            )
            response.raise_for_status()
            data = response.json()

            # ── Accumulate token usage ────────────────────────────────
            usage = data.get("usage", {})
            details = usage.get("completion_tokens_details", {})
            self._session_usage.add(
                usage.get("prompt_tokens", 0),
                usage.get("completion_tokens", 0),
                details.get("reasoning_tokens", 0),
            )

            content: str = data["choices"][0]["message"]["content"]

            if json_mode:
                return json.loads(content)
            return content

        except httpx.HTTPStatusError as exc:
            logger.error("Grok API HTTP error %s: %s", exc.response.status_code, exc.response.text)
            return {} if json_mode else ""
        except (json.JSONDecodeError, KeyError, IndexError) as exc:
            logger.error("Failed to parse Grok response: %s", exc)
            return {} if json_mode else ""
        except httpx.RequestError as exc:
            logger.error("Grok API request failed: %s", exc)
            return {} if json_mode else ""
