"""
Groq API Model Wrapper

This module handles communication with Groq API for LLaMA inference.
Using Groq API is much faster than local inference and doesn't require GPU.
"""

import logging
from groq import Groq
from typing import Optional
import os
from ..config import settings

logger = logging.getLogger(__name__)


class LLaMAModel:
    """
    Wrapper for LLaMA model via Groq API

    Handles:
    - API client initialization
    - Text generation via Groq
    - Error handling
    """

    def __init__(self):
        self.api_key = settings.GROQ_API_KEY
        self.model_name = settings.GROQ_MODEL
        self.max_tokens = settings.MAX_TOKENS

        self.client: Optional[Groq] = None

        self._initialize_client()

    def _initialize_client(self):
        """Initialize Groq API client"""
        try:
            if not self.api_key:
                raise ValueError("GROQ_API_KEY not set in environment variables")

            logger.info(f"Initializing Groq API client with model: {self.model_name}")

            self.client = Groq(api_key=self.api_key)

            logger.info("Groq API client initialized successfully!")

        except Exception as e:
            logger.error(f"Failed to initialize Groq client: {e}")
            raise

    def generate(self, prompt: str, max_new_tokens: int = 150, temperature: float = 0.7) -> str:
        """
        Generate text from prompt using Groq API with retry logic

        Args:
            prompt: Input prompt
            max_new_tokens: Maximum number of tokens to generate
            temperature: Sampling temperature (lower = more deterministic)

        Returns:
            Generated text
        """
        import time
        import random
        from groq import RateLimitError

        if self.client is None:
            raise RuntimeError("Groq client not initialized")

        max_retries = 5
        base_delay = 2  # seconds

        for attempt in range(max_retries):
            try:
                # Call Groq API
                chat_completion = self.client.chat.completions.create(
                    messages=[
                        {
                            "role": "user",
                            "content": prompt,
                        }
                    ],
                    model=self.model_name,
                    temperature=temperature,
                    max_tokens=max_new_tokens,
                    top_p=0.9,
                    stream=False,
                )

                # Extract response
                response = chat_completion.choices[0].message.content.strip()
                logger.debug(f"Groq API response: {response[:100]}...")
                return response

            except RateLimitError as e:
                wait_time = (base_delay * (2 ** attempt)) + (random.random() * 2)
                logger.warning(f"Rate limit hit (attempt {attempt + 1}/{max_retries}). Waiting {wait_time:.2f}s...")
                if attempt == max_retries - 1:
                    logger.error(f"Max retries reached for RateLimitError: {e}")
                    raise
                time.sleep(wait_time)
            except Exception as e:
                logger.error(f"Groq API generation failed: {e}")
                raise

    def is_loaded(self) -> bool:
        """Check if client is initialized"""
        return self.client is not None

