"""
Voyage AI API Client
Generated by vai v{{vaiVersion}} on {{generatedAt}}

Model: {{model}}
Dimensions: {{dimensions}}
"""

import os
import requests
from typing import List, Optional, Dict, Any

VOYAGE_API_URL = os.getenv("VOYAGE_API_URL", "https://api.voyageai.com/v1")
VOYAGE_API_KEY = os.getenv("VOYAGE_API_KEY")

if not VOYAGE_API_KEY:
    raise ValueError("VOYAGE_API_KEY environment variable is required")


def embed(
    texts: List[str],
    model: str = "{{model}}",
    input_type: str = "{{inputType}}",
    output_dimension: int = {{dimensions}},
) -> Dict[str, Any]:
    """
    Generate embeddings for a list of texts using Voyage AI.
    
    Args:
        texts: List of texts to embed
        model: Embedding model name
        input_type: 'document' or 'query'
        output_dimension: Output vector dimensions
        
    Returns:
        Dict with 'embeddings' (list of vectors) and 'usage' (token counts)
    """
    response = requests.post(
        f"{VOYAGE_API_URL}/embeddings",
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {VOYAGE_API_KEY}",
        },
        json={
            "model": model,
            "input": texts,
            "input_type": input_type,
            "output_dimension": output_dimension,
        },
    )
    
    response.raise_for_status()
    data = response.json()
    
    return {
        "embeddings": [d["embedding"] for d in data["data"]],
        "usage": data["usage"],
    }


def embed_query(query: str, **kwargs) -> List[float]:
    """Embed a single query, returning the embedding vector."""
    result = embed([query], input_type="query", **kwargs)
    return result["embeddings"][0]


def embed_documents(documents: List[str], **kwargs) -> List[List[float]]:
    """Embed multiple documents, returning list of embedding vectors."""
    result = embed(documents, input_type="document", **kwargs)
    return result["embeddings"]


{{#if rerank}}
def rerank(
    query: str,
    documents: List[str],
    model: str = "{{rerankModel}}",
    top_k: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Rerank documents by relevance to a query.
    
    Args:
        query: The query to rank against
        documents: List of documents to rerank
        model: Rerank model name
        top_k: Number of top results to return
        
    Returns:
        Dict with 'results' containing index, relevance_score, and document
    """
    payload = {
        "model": model,
        "query": query,
        "documents": documents,
        "return_documents": True,
    }
    if top_k:
        payload["top_k"] = top_k
    
    response = requests.post(
        f"{VOYAGE_API_URL}/rerank",
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {VOYAGE_API_KEY}",
        },
        json=payload,
    )
    
    response.raise_for_status()
    data = response.json()
    
    return {
        "results": [
            {
                "index": d["index"],
                "relevance_score": d["relevance_score"],
                "document": d.get("document"),
            }
            for d in data["data"]
        ],
        "usage": data.get("usage"),
    }
{{/if}}