"""Diffbot LLM RAG API: stream a chat completion."""

import json
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List

if TYPE_CHECKING:
    from .client import Diffbot, DiffbotAsync


def _build_payload(client: Any, messages: List[Dict[str, str]]) -> tuple:
    headers = {"Authorization": f"Bearer {client.token}"}
    payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
    return headers, payload


def _parse_chunk(line: str):
    try:
        chunk = json.loads(line.replace("data: ", ""))
    except json.JSONDecodeError:
        return None
    choices = chunk.get("choices")
    if choices and choices[0].get("delta", {}).get("content"):
        return choices[0]["delta"]["content"]
    return None


def ask(client: "Diffbot", messages: List[Dict[str, str]]) -> Iterator[str]:
    headers = {"Authorization": f"Bearer {client.token}"}
    payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
    with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
        client._raise_for_status(response)
        for line in response.iter_lines():
            if line:
                content = _parse_chunk(line)
                if content:
                    yield content


async def ask_async(client: "DiffbotAsync", messages: List[Dict[str, str]]) -> AsyncIterator[str]:
    headers = {"Authorization": f"Bearer {client.token}"}
    payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
    async with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
        client._raise_for_status(response)
        async for line in response.aiter_lines():
            if line:
                content = _parse_chunk(line)
                if content:
                    yield content