import httpx
import os
from typing import AsyncIterator

OLLAMA = os.getenv("OLLAMA_URL", "http://192.168.20.47:11434")
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
CHAT_MODEL  = os.getenv("CHAT_MODEL",  "qwen2.5-coder:7b-instruct-q4_K_M")

async def embed(text: str) -> list[float]:
    async with httpx.AsyncClient(timeout=60) as c:
        r = await c.post(f"{OLLAMA}/api/embeddings",
                         json={"model": EMBED_MODEL, "prompt": text})
        r.raise_for_status()
        return r.json()["embedding"]

async def chat(messages: list[dict], model: str = None, stream: bool = False):
    async with httpx.AsyncClient(timeout=300) as c:
        r = await c.post(f"{OLLAMA}/api/chat",
                         json={"model": model or CHAT_MODEL,
                               "messages": messages,
                               "stream": False})
        r.raise_for_status()
        return r.json()["message"]["content"]

async def chat_stream(messages: list[dict], model: str = None) -> AsyncIterator[str]:
    async with httpx.AsyncClient(timeout=300) as c:
        async with c.stream("POST", f"{OLLAMA}/api/chat",
                            json={"model": model or CHAT_MODEL,
                                  "messages": messages,
                                  "stream": True}) as r:
            async for line in r.aiter_lines():
                if line:
                    import json
                    data = json.loads(line)
                    if not data.get("done"):
                        yield data.get("message", {}).get("content", "")