"""Load GROQ models using LangChain."""

import os
from getpass import getpass
from langchain_groq import ChatGroq
from chemgraph.utils.logging_config import setup_logger

logger = setup_logger(__name__)


def load_groq_model(
    model_name: str,
    temperature: float,
    api_key: str = None,
    prompt: str = None,
    base_url: str = None,
) -> ChatGroq:
    """Load a GROQ chat model into LangChain.
    This function loads a GROQ model and configures it for use with LangChain.
    It handles API key management, including prompting for the key if not provided
    or if the provided key is invalid.
    Parameters
    ----------
    model_name : str
        The name of the GROQ chat model to load. See supported_groq_models for list
        of supported models.
    temperature : float
        Controls the randomness of the generated text. Higher values (e.g., 0.8)
        make the output more random, while lower values (e.g., 0.2) make it more
        deterministic.
    api_key : str, optional
        The GROQ API key. If not provided, the function will attempt to retrieve it
        from the environment variable `GROQ_API_KEY`.
    prompt : str, optional
        Custom prompt to use when requesting the API key from the user.
    base_url : str, optional
        Custom base URL for the GROQ API (currently unused but included for consistency).
    Returns
    -------
    ChatGroq
        An instance of LangChain's ChatGroq model.
    Raises
    ------
    ValueError
        If the model name is not in the list of supported models.
    Exception
        If there is an error loading the model or if the API key is invalid.
    Notes
    -----
    The function will:
    1. Check for the API key in the environment variables
    2. Prompt for the key if not found
    3. Validate the model name against supported models
    4. Attempt to load the model
    5. Handle any authentication errors by prompting for a new key
    """

    # Strip the "groq:" routing prefix before sending to the API.
    if model_name.startswith("groq:"):
        model_name = model_name.removeprefix("groq:")

    if api_key is None:
        api_key = os.getenv("GROQ_API_KEY")
        if not api_key:
            logger.info("GROQ API key not found in environment variables.")
            api_key = getpass("Please enter your GROQ API key: ")
            os.environ["GROQ_API_KEY"] = api_key

    try:
        logger.info(f"Loading GROQ model: {model_name}")
        llm = ChatGroq(
            model=model_name,
            temperature=temperature,
            api_key=api_key,
            max_tokens=6000,
        )
        # No guarantee that api_key is valid, authentication happens only during invocation
        logger.info(f"Requested model: {model_name}")
        logger.info("GROQ model loaded successfully")
        return llm
    except Exception as e:
        # Can remove this since authentication happens only during invocation
        if "AuthenticationError" in str(e) or "invalid_api_key" in str(e):
            logger.warning("Invalid GROQ API key.")
            api_key = getpass("Please enter a valid GROQ API key: ")
            os.environ["GROQ_API_KEY"] = api_key
            # Retry with new API key
            return load_groq_model(model_name, temperature, api_key, prompt)
        else:
            logger.error(f"Error loading GROQ model: {str(e)}")
            raise