"""
Groq Client - Ultra-fast LLM inference

Provides LLM completion via Groq API with industry-leading inference speeds.
Groq specializes in fast inference using custom hardware (LPU - Language Processing Unit).
"""

import os
from collections import defaultdict
from typing import Any, Optional

from groknroll.clients.base_lm import BaseLM
from groknroll.core.types import ModelUsageSummary, UsageSummary


class GroqClient(BaseLM):
    """
    Groq Language Model Client

    Groq provides ultra-fast inference for open-source models like:
    - Llama 3.3 70B
    - Mixtral 8x7B
    - Gemma 2 9B
    - And more

    Features:
    - Industry-leading inference speed (up to 700+ tokens/sec)
    - Chat completion
    - Streaming responses
    - Token usage tracking
    - OpenAI-compatible API

    Example:
        >>> from groknroll.clients.groq import GroqClient
        >>> client = GroqClient(model_name="llama-3.3-70b-versatile")
        >>> response = await client.acompletion("Hello, world!")
        >>> print(response)
    """

    # Popular Groq models
    MODELS = [
        "llama-3.3-70b-versatile",  # Latest Llama 3.3
        "llama-3.1-70b-versatile",  # Llama 3.1 70B
        "llama-3.1-8b-instant",  # Fast Llama 3.1 8B
        "mixtral-8x7b-32768",  # Mixtral 8x7B
        "gemma2-9b-it",  # Google Gemma 2 9B
        "gemma-7b-it",  # Google Gemma 7B
    ]

    def __init__(
        self,
        api_key: Optional[str] = None,
        model_name: str = "llama-3.3-70b-versatile",
        base_url: Optional[str] = None,
        **kwargs,
    ):
        """
        Initialize Groq client

        Args:
            api_key: Groq API key (reads from GROQ_API_KEY env if not provided)
            model_name: Model to use (default: llama-3.3-70b-versatile)
            base_url: Optional base URL override
            **kwargs: Additional Groq-specific options
        """
        super().__init__(model_name=model_name, **kwargs)

        self.api_key = api_key or os.getenv("GROQ_API_KEY")
        if not self.api_key:
            raise ValueError(
                "Groq API key not found. Set GROQ_API_KEY environment variable or pass api_key parameter."
            )

        self.model_name = model_name
        self.base_url = base_url or "https://api.groq.com/openai/v1"

        # Initialize Groq SDK (uses OpenAI-compatible API)
        try:
            from groq import Groq, AsyncGroq
        except ImportError:
            raise ImportError(
                "groq package not found. Install with: pip install groq"
            )

        self.client = Groq(api_key=self.api_key, base_url=self.base_url)
        self.async_client = AsyncGroq(api_key=self.api_key, base_url=self.base_url)

        # Per-model usage tracking
        self.model_call_counts: dict[str, int] = defaultdict(int)
        self.model_input_tokens: dict[str, int] = defaultdict(int)
        self.model_output_tokens: dict[str, int] = defaultdict(int)
        self.model_total_tokens: dict[str, int] = defaultdict(int)

        # Last call tracking
        self.last_prompt_tokens = 0
        self.last_completion_tokens = 0

    def completion(
        self,
        prompt: str | list[dict[str, Any]],
        model: Optional[str] = None,
    ) -> str:
        """
        Synchronous completion

        Args:
            prompt: User prompt (string or messages list)
            model: Model to use (defaults to self.model_name)

        Returns:
            Generated text string
        """
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            messages = prompt
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Groq client.")

        response = self.client.chat.completions.create(
            model=model,
            messages=messages,
        )

        self._track_cost(response, model)
        return response.choices[0].message.content

    async def acompletion(
        self,
        prompt: str | list[dict[str, Any]],
        model: Optional[str] = None,
    ) -> str:
        """
        Asynchronous completion

        Args:
            prompt: User prompt (string or messages list)
            model: Model to use (defaults to self.model_name)

        Returns:
            Generated text string
        """
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            messages = prompt
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Groq client.")

        response = await self.async_client.chat.completions.create(
            model=model,
            messages=messages,
        )

        self._track_cost(response, model)
        return response.choices[0].message.content

    def _track_cost(self, response, model: str):
        """Track token usage from API response"""
        self.model_call_counts[model] += 1

        usage = getattr(response, "usage", None)
        if usage is None:
            raise ValueError("No usage data received. Tracking tokens not possible.")

        self.model_input_tokens[model] += usage.prompt_tokens
        self.model_output_tokens[model] += usage.completion_tokens
        self.model_total_tokens[model] += usage.total_tokens

        # Track last call for handler to read
        self.last_prompt_tokens = usage.prompt_tokens
        self.last_completion_tokens = usage.completion_tokens

    def get_usage_summary(self) -> UsageSummary:
        """Get usage statistics across all models"""
        model_summaries = {}
        for model in self.model_call_counts:
            model_summaries[model] = ModelUsageSummary(
                total_calls=self.model_call_counts[model],
                total_input_tokens=self.model_input_tokens[model],
                total_output_tokens=self.model_output_tokens[model],
            )
        return UsageSummary(model_usage_summaries=model_summaries)

    def get_last_usage(self) -> ModelUsageSummary:
        """Get usage statistics from last completion"""
        return ModelUsageSummary(
            total_calls=1,
            total_input_tokens=self.last_prompt_tokens,
            total_output_tokens=self.last_completion_tokens,
        )
