import os
import requests

EMBEDDING_PROVIDER = {
    "text-embedding-3-small": "openai",
    "text-embedding-3-large": "openai",
    "text-embedding-ada-002": "openai",
    "all-minilm-l6-v2": "ibm-watsonx",
    "granite-embedding-107m-multilingual": "ibm-watsonx",
    "granite-embedding-278m-multilingual": "ibm-watsonx",
    "ms-marco-minilm-l-12-v2": "ibm-watsonx",
    "multilingual-e5-large": "ibm-watsonx",
    "slate-30m-english-rtrvr-v2": "ibm-watsonx",
    "slate-30m-english-rtrvr": "ibm-watsonx",
    "slate-125m-english-rtrvr-v2": "ibm-watsonx",
    "slate-125m-english-rtrvr": "ibm-watsonx"
}


def run_prompt(prompt: str) -> str:
    base_url = os.getenv("OPENAI_BASE_URL") or "https://api.openai.com/v1"
    api_key = os.getenv("OPENAI_API_KEY")

    if not base_url or not api_key:
        raise EnvironmentError("Missing OPENAI_BASE_URL or OPENAI_API_KEY environment variables.")

    url = f"{base_url.rstrip('/')}/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    payload = {
        "model": "gpt-4o-mini",
        "messages": [{"role": "user", "content": prompt}],
    }
    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()
    return response.json()
    
def generate_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]:
    if EMBEDDING_PROVIDER[model] == "openai":
        return generate_embedding_openai(text, model)
    elif EMBEDDING_PROVIDER[model] == "ibm-watsonx":
        return generate_embedding_ibm_watsonx(text, model)
    else:
        raise ValueError(f"Unsupported embedding model: {model}")

def generate_embedding_ibm_watsonx(text: str, model: str = "granite-embedding-278m-multilingual") -> list[float]:
    """ 
    Generate an embedding using the IBM Watsonx REST API without using the SDK.
    Environment variables required:
      - IBM_WATSONX_BASE_URL: the base URL of the IBM Watsonx endpoint (e.g., https://{cluster_url}/ml/v1/text/embeddings?version=2023-10-25)
      - IBM_WATSONX_API_KEY: your API key
      - IBM_WATSONX_PROJECT_ID: your project ID
    """
    base_url = os.getenv("IBM_WATSONX_BASE_URL")
    api_key = os.getenv("IBM_WATSONX_API_KEY")
    project_id = os.getenv("IBM_WATSONX_PROJECT_ID")

    if not base_url or not api_key or not project_id:
        raise EnvironmentError("Missing IBM_WATSONX_BASE_URL or IBM_WATSONX_API_KEY or IBM_WATSONX_PROJECT_ID environment variables.")

    url = f"{base_url.rstrip('/')}"

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    payload = {
        "model_id": model,
        "inputs": [text],
        "project_id": os.getenv("IBM_WATSONX_PROJECT_ID")
    }

    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()

    data = response.json()
    return data["results"][0]["embedding"]

def generate_embedding_openai(text: str, model: str = "text-embedding-3-small") -> list[float]:
    """
    Generate an embedding using the OpenAI REST API without using the SDK.

    Environment variables required:
      - OPENAI_API_KEY: your API key
      - OPENAI_BASE_URL: the base URL of the OpenAI endpoint (e.g., https://api.openai.com/v1)
    """
    base_url = os.getenv("OPENAI_BASE_URL") or "https://api.openai.com/v1"
    api_key = os.getenv("OPENAI_API_KEY")

    if not base_url or not api_key:
        raise EnvironmentError("Missing OPENAI_BASE_URL or OPENAI_API_KEY environment variables.")

    url = f"{base_url.rstrip('/')}/embeddings"

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }

    payload = {
        "model": model,
        "input": text,
    }

    response = requests.post(url, headers=headers, json=payload)
    response.raise_for_status()

    data = response.json()
    return data["data"][0]["embedding"]


