from stores.llm.LLMInterface import LLMInterface
import logging
import requests
import re
import os


class GroqProvider(LLMInterface):
    def __init__(self, url: str = None, model: str = None,
                 default_input_max_characters: int = 1000,
                 default_generation_max_output_tokens: int = 1000,
                 default_generation_temperature: float = 0.1, api_key: str = None):
        self.url = url or "https://api.groq.com/openai/v1"
        self.api_key = api_key or os.getenv("GROQ_API_KEY")
        self.model = model
        self.generation_model_id = None

        self.embedding_model = None
        self.embedding_model_id = None
        self.embedding_size = None

        self.default_input_max_characters = default_input_max_characters
        self.default_generation_max_output_tokens = default_generation_max_output_tokens
        self.default_generation_temperature = default_generation_temperature
        self.logger = logging.getLogger(__name__)

    def set_generation_model(self, model_id: str):
        if model_id:
            self.model = model_id

    def set_embedding_model(self, model_id: str, embedding_size: int):
        if model_id:
            self.embedding_model = model_id
            self.embedding_size = embedding_size
            self.embedding_model_id = model_id

    def process_text(self, text: str):
        if not text:
            return ""
        return str(text).strip()

    def generate_text(self, prompt: str, chat_history: list = None,
                      max_output_tokens: int = None, temperature: float = None):
        try:
            chat_history = chat_history or []
            clean_prompt = self.process_text(prompt)

            messages = []
            for entry in chat_history:
                messages.append({
                    "role": entry.get("role", "user"),
                    "content": entry.get("content", "")
                })
            messages.append({"role": "user", "content": clean_prompt})

            payload = {
                "model": self.model,
                "messages": messages,
                "max_tokens": int(max_output_tokens or self.default_generation_max_output_tokens),
                "temperature": float(temperature or self.default_generation_temperature),
            }

            url = self.url.rstrip("/") + "/chat/completions"
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            }

            resp = requests.post(url, json=payload, headers=headers, timeout=6000)
            if resp.status_code != 200:
                self.logger.error("Groq generate failed: %s %s", resp.status_code, resp.text)
                return None

            data = resp.json()

            try:
                generated_text = data["choices"][0]["message"]["content"].strip()
            except (KeyError, IndexError, TypeError):
                self.logger.error("Unexpected Groq response structure: %s", data)
                return None

            if not generated_text:
                return None

            usage = data.get("usage", {})
            # Groq exposes x_groq.usage.total_time in seconds
            total_time_ms = None
            try:
                total_time_ms = round(data["x_groq"]["usage"]["total_time"] * 1000, 2)
            except (KeyError, TypeError):
                pass

            return {
                "model": data.get("model"),
                "response": generated_text,
                "tokens_generated": usage.get("completion_tokens"),
                "total_duration_ms": total_time_ms,
                "prompt_eval_tokens": usage.get("prompt_tokens"),
            }

        except Exception as e:
            self.logger.exception("Error in GroqProvider.generate_text: %s", e)
            return None

    def embed_text(self, text: str, document_type: str = None):
        """Groq does not support embeddings — returns None."""
        self.logger.warning("GroqProvider does not support embeddings.")
        return None

    def construct_prompt(self, prompt: str, role: str):
        return {
            "role": role,
            "content": self.process_text(prompt)
        }

    def embed_text_batch(self, texts: list[str], batch_size: int = 32):
        """Groq does not support embeddings — returns None."""
        self.logger.warning("GroqProvider does not support embeddings.")
        return None

    def clean_content(self, text: str) -> str:
        text = re.sub(r'\[.*?\]\(.*?\)', '', text)
        text = re.sub(r'\[[^\]]*\]', '', text)
        text = re.sub(r'\n+', '\n', text).strip()
        return text

    def web_search(self, query: str):
        """Groq has no native web search — returns a not-supported notice."""
        self.logger.warning("GroqProvider.web_search is not natively supported.")
        return {
            "text": "Web search is not natively supported by the Groq API.",
            "references": []
        }