import json
import os
import random
import time
from typing import TypeVar, Type

from dotenv import load_dotenv
from groq import Groq, RateLimitError
from pydantic import BaseModel

load_dotenv()

T = TypeVar("T", bound=BaseModel)

MODEL = os.getenv("GROQ_MODEL", "openai/gpt-oss-120b")

client = Groq(
    api_key=os.getenv("GROQ_API_KEY")
)


def call_structured_llm(
    *,
    schema_name: str,
    schema_model: Type[T],
    system_prompt: str,
    feedback: str,
    max_retries: int = 6,
) -> T:

    schema = schema_model.model_json_schema()

    for attempt in range(max_retries):

        try:
            response = client.chat.completions.create(
                model=MODEL,
                messages=[
                    {
                        "role": "system",
                        "content": system_prompt,
                    },
                    {
                        "role": "user",
                        "content": feedback,
                    },
                ],
                response_format={
                    "type": "json_schema",
                    "json_schema": {
                        "name": schema_name,
                        "strict": True,
                        "schema": schema,
                    },
                },
            )

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

            data = json.loads(raw)

            return schema_model.model_validate(data)

        except RateLimitError:

            if attempt == max_retries - 1:
                raise

            wait_time = min(
                (2 ** attempt) + random.uniform(0, 1),
                30,
            )

            print(
                f"[RATE LIMIT] Groq quota reached. "
                f"Retrying in {wait_time:.1f}s..."
            )

            time.sleep(wait_time)