import json
import logging
from typing import AsyncGenerator

import httpx
from agentic_chat_ui.src.models import AgentClient
from agentic_chat_ui.src.settings import ApiAgentClientSettings

_logger = logging.getLogger(__name__)


class ApiAgentClient(AgentClient[ApiAgentClientSettings]):
    """
    Client responsible for talking to the agent-api backend.
    """

    def __init__(self, config: ApiAgentClientSettings) -> None:
        super().__init__(config)

    async def astream_chat(self, query: str) -> AsyncGenerator[str, None]:  # type: ignore[override]
        payload = {
            "model": "local-agent",
            "messages": [
                {"role": "user", "content": query},
            ],
            "stream": True,
        }
        headers = {
            "Accept": "text/event-stream",
            "Content-Type": "application/json",
        }
        async with httpx.AsyncClient(timeout=None) as client:
            async with client.stream(
                "POST",
                self.config.api_url,  # /v1/chat/completions
                json=payload,
                headers=headers,
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if not line or not line.startswith("data:"):
                        continue
                    data = line.removeprefix("data:").strip()
                    if data == "[DONE]":
                        break
                    event = json.loads(data)
                    delta = event["choices"][0]["delta"]
                    if "content" in delta:
                        yield delta["content"]
