import requests
from domains.Credenciales.ollama_credencial import OllamaCredencial

class OllamaCliente:
    def __init__(self, credencial: OllamaCredencial):
        """
        Inicializa el cliente de Ollama con una instancia de OllamaCredencial.
        """
        self.credencial = credencial
        self.base_url = self.credencial.base_url

    # --- Chat Completions ---
    def chat_completion(self, model: str, messages: list, stream: bool = False, **kwargs):
        url = f"{self.base_url}/api/chat"
        payload = {
            "model": model,
            "messages": messages,
            "stream": stream,
            **kwargs
        }
        response = requests.post(url, json=payload, stream=stream)
        response.raise_for_status()

        return self._handle_stream(response) if stream else response.json()

    def _handle_stream(self, response):
        for line in response.iter_lines():
            if line:
                try:
                    parsed = line.decode("utf-8")
                    # Extraer contenido si está en JSON como {"message":{"content":"..."},...}
                    if parsed.startswith("{"):
                        import json
                        data = json.loads(parsed)
                        if "message" in data and "content" in data["message"]:
                            yield data["message"]["content"]
                except Exception:
                    continue

    # --- Text Completion (legacy) ---
    def completion(self, model: str, prompt: str, **kwargs):
        url = f"{self.base_url}/api/generate"
        payload = {
            "model": model,
            "prompt": prompt,
            **kwargs
        }
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return response.json()

    # --- Embeddings ---
    def embedding(self, model: str, input: str | list[str], **kwargs):
        url = f"{self.base_url}/api/embeddings"
        payload = {
            "model": model,
            "prompt": input,
            **kwargs
        }
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return response.json()