# services/llm/groq_service.py
"""
Groq service implementation
"""
from groq import Groq
from typing import Dict, List
import json
from .base import LLMService
from config.api_keys import APIKeyManager
from utils.logger import setup_logger


class GroqService(LLMService):
    def __init__(
        self,
        model: str = "llama3-8b-8192",
        temperature: float = 0.3,
        max_tokens: int = 8000,
    ):
        api_key = APIKeyManager.get_groq_key()
        if not api_key:
            raise ValueError(
                "Groq API key not found. Please set the GROQ_API_KEY environment variable."
            )

        self.client = Groq(api_key=api_key)
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.logger = setup_logger("groq")

    def generate_text(self, prompt: str) -> str:
        self.logger.info(f"--- PROMPT ---\n{prompt}")
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[{"role": "user", "content": prompt}],
                temperature=self.temperature,
                max_tokens=self.max_tokens,
            )
            response_content = response.choices[0].message.content
            self.logger.info(f"--- RESPONSE ---\n{response_content}")
            return response_content
        except Exception as e:
            self.logger.error(f"Error in Groq API call: {e}")
            print(f"Error in Groq API call: {e}")
            return None

    def get_similarity_scores(self, texts_pairs: Dict[str, List[str]]) -> List[float]:
        system_prompt = (
            "Evaluate the semantic similarity between the following table of pairs of texts in json format on a scale from 0 to 1. "
            "Return the similarity scores for every row in JSON format as a list of numbers, without any additional text or formatting."
        )

        request_payload = json.dumps(texts_pairs)

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": request_payload},
                ],
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                response_format={"type": "json_object"},
            )

            response_content = response.choices[0].message.content

            try:
                scores = json.loads(response_content)
                if isinstance(scores, dict) and "similarity_scores" in scores:
                    return scores["similarity_scores"]
                elif isinstance(scores, list):
                    return scores
                else:
                    raise ValueError("Unexpected response format")
            except json.JSONDecodeError:
                raise ValueError("Could not decode response as JSON")

        except Exception as e:
            print(f"Error in Groq similarity calculation: {e}")
            return None
