import os

import requests

# Get Ollama server address from environment variable, default to local address
ollama_addr = os.environ.get("OLLAMA_ADDR", "http://host.docker.internal:11434")


def prompt(model, text, allow_pull=False):
    """
    Send a prompt to Ollama server and return the response.

    Args:
        model (str): The model to use
        text (str): The prompt text to send
        allow_pull (bool): If True, automatically download the model if not available

    Returns:
        str: The response from the Ollama server
    """
    # Check if model is available
    if allow_pull and not is_model_available(model):
        # Try to pull the model
        pull_result = pull_model(model)
        if not pull_result.startswith("Model downloaded successfully"):
            return pull_result  # Return error from pull

    url = f"{ollama_addr}/api/generate"

    payload = {"model": model, "prompt": text, "stream": False}

    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return response.json()["response"]
    except requests.exceptions.RequestException as e:
        return f"Error: {str(e)}"


def pull_model(model_name, force=False):
    """
    Download a model from Ollama.

    Args:
        model_name (str): The name of the model to download
        force (bool): If True, forces download even if model already exists

    Returns:
        str: Success message or error message
    """
    url = f"{ollama_addr}/api/pull"

    payload = {"name": model_name}
    if force:
        payload["force"] = True

    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return "Model downloaded successfully"
    except requests.exceptions.RequestException as e:
        return f"Error downloading model: {str(e)}"


def is_model_available(model_name):
    """
    Check if a model is available in Ollama.

    Args:
        model_name (str): The name of the model to check

    Returns:
        bool: True if model is available, False otherwise
    """
    url = f"{ollama_addr}/api/tags"

    try:
        response = requests.get(url)
        response.raise_for_status()
        models = response.json().get("models", [])
        return any(model.get("name") == model_name for model in models)
    except requests.exceptions.RequestException:
        return False