import os
import base64
import json
import urllib.request
import urllib.error

API_HOST = os.getenv("OPENAI_API_HOST", "https://chat.matthiasg.dev/ollama")
AI_MODEL = os.getenv("OPENAI_MODEL", "llava:7b-v1.6-mistral-q4_1")
AI_MODEL_THINK = os.getenv("OPENAI_MODEL_THINK", "True").lower() in ("true", "1", "t")
OPENAPI_TOKEN = os.getenv("OPENAI_API_TOKEN", None)
OPENAPI_ENDPOINT = "/api/generate"


def call_api(body: dict, endpoint: str = OPENAPI_ENDPOINT) -> str:
    """Call the API with the given endpoint and body dict."""
    url = f"{API_HOST}{endpoint}"
    headers = {
        "Content-Type": "application/json",
    }
    if OPENAPI_TOKEN:
        headers["Authorization"] = f"Bearer {OPENAPI_TOKEN}"

    json_data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url, data=json_data, headers=headers, method="POST")

    try:
        with urllib.request.urlopen(req) as response:
            return response.read().decode("utf-8")
    except urllib.error.HTTPError as e:
        error_body = e.read().decode("utf-8")
        raise Exception(f"Error calling API: HTTP {e.code} - {error_body}")
    except urllib.error.URLError as e:
        raise Exception(f"Failed to reach server: {e.reason}")


def generate(
    model: str = AI_MODEL,
    prompt: str = "",
    images: list[str] = None,
    output_format: str = None,
    system_message: str = None,
    keep_alive: bool = True,
    should_think: bool = AI_MODEL_THINK,
) -> str:
    """
    Generate a response for a given prompt with a provided model via the Ollama/OpenAI API.
    Handles base64 encoding for local image file paths and structures the request body.
    """
    if images is None:
        images = []

    # Transform image file paths to base64 strings
    encoded_images = []
    for img_path in images:
        if os.path.isfile(img_path):
            with open(img_path, "rb") as image_file:
                encoded_images.append(base64.b64encode(image_file.read()).decode("utf-8"))
        else:
            # If it's already a base64 string or an invalid path, keep as-is
            encoded_images.append(img_path)

    body = {
        "model": model,
        "prompt": prompt,
        "images": encoded_images,
        "think": should_think,
        "stream": False,
    }

    if system_message is not None:
        body["system"] = system_message

    if output_format is not None:
        try:
            body["format"] = json.loads(output_format)
        except json.JSONDecodeError:
            body["format"] = output_format

    if not keep_alive:
        body["keep_alive"] = "0m"

    response_text = call_api(body)

    try:
        decoded_response = json.loads(response_text)
    except json.JSONDecodeError as e:
        raise Exception(f"Error decoding JSON response: {e}")

    return decoded_response.get("response", "")


if __name__ == "__main__":
    # Example usage:
    result = generate(
        prompt="Explain what you see in the image or answer this prompt.",
        should_think=AI_MODEL_THINK,
    )
    print(result)