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

class AiSemanticEngineClient:
    """Client for interacting with the AI Semantic Engine API."""

    def __init__(self, api_key: str, base_url: str = "http://localhost:8200"):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.session.headers.update({"X-API-KEY": self.api_key})

    def embed(self, text_or_texts: Union[str, List[str]], model: str = "fast", store: bool = False) -> Dict[str, Any]:
        """Generate semantic embeddings for a string or a list of texts."""
        if isinstance(text_or_texts, str):
            resp = self.session.post(
                f"{self.base_url}/embed",
                json={"text": text_or_texts, "model": model, "store": store}
            )
        else:
            resp = self.session.post(
                f"{self.base_url}/embed-batch",
                json={"texts": text_or_texts, "model": model, "store": store}
            )
        resp.raise_for_status()
        return resp.json()

    def similarity(self, text1: str, text2: str, model: str = "fast") -> Dict[str, Any]:
        """Compute cosine similarity between two texts."""
        resp = self.session.post(
            f"{self.base_url}/similarity",
            json={"text1": text1, "text2": text2, "model": model}
        )
        resp.raise_for_status()
        return resp.json()

    def store_item(self, text: str, metadata: Optional[Dict[str, Any]] = None, model: str = "fast") -> Dict[str, Any]:
        """Store a text item and its metadata persistently."""
        resp = self.session.post(
            f"{self.base_url}/items",
            json={"text": text, "metadata": metadata or {}, "model": model}
        )
        resp.raise_for_status()
        return resp.json()

    def search(self, query: str, top_k: int = 5, model: str = "fast") -> Dict[str, Any]:
        """Search the persistent storage for texts matching the query."""
        resp = self.session.post(
            f"{self.base_url}/search",
            json={"query": query, "top_k": top_k, "model": model}
        )
        resp.raise_for_status()
        return resp.json()
