"""AnalogAI SDK helpers for chat completions and streaming."""

from __future__ import annotations

import json
import os
from typing import Any, Dict, Generator, Iterable, Optional
from uuid import uuid4

import requests

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

def _resolve_base_url() -> str:
    explicit_base_url = os.getenv("CHAT_BASE_URL")
    if explicit_base_url:
        return explicit_base_url
    return "https://chat.analogai.net:8001"


def _normalize_messages(
    message_text: str,
    auth_type: str,
    messages: Optional[Iterable[Dict[str, Any]]],
) -> list[Dict[str, Any]]:
    source_messages = (
        list(messages)
        if messages is not None
        else [{"role": auth_type, "content": message_text}]
    )
    normalized: list[Dict[str, Any]] = []
    for message in source_messages:
        if not isinstance(message, dict):
            continue
        message_copy = dict(message)
        message_copy.setdefault("id", uuid4().hex)
        message_copy.setdefault("role", auth_type)
        normalized.append(message_copy)
    return normalized


def _build_payload(
    *,
    message_text: str,
    auth_type: str,
    messages: Optional[Iterable[Dict[str, str]]],
    state: Optional[Dict[str, Any]],
    tools: Optional[Iterable[Dict[str, Any]]],
    thread_id: Optional[str],
    run_id: Optional[str],
    mode: Optional[str],
) -> Dict[str, Any]:
    payload_messages = _normalize_messages(message_text, auth_type, messages)
    payload: Dict[str, Any] = {
        "messages": payload_messages,
        "thread_id": thread_id or uuid4().hex,
        "run_id": run_id or uuid4().hex,
        "state": dict(state or {}),
        "tools": list(tools or []),
        "context": [],
        "forwarded_props": {},
    }
    if mode:
        payload["state"]["mode"] = mode
    return payload


def _iter_sse_payloads(response: requests.Response) -> Generator[str, None, None]:
    data_lines: list[str] = []
    for raw_line in response.iter_lines(decode_unicode=True):
        if raw_line is None:
            continue
        line = raw_line.rstrip("\r")
        if not line:
            if data_lines:
                yield "\n".join(data_lines)
                data_lines = []
            continue
        if line.startswith(":"):
            continue
        if line.startswith("data:"):
            data_lines.append(line[5:].lstrip())
            continue
    if data_lines:
        yield "\n".join(data_lines)


def _iter_sse_events(response: requests.Response) -> Generator[Dict[str, Any], None, None]:
    for payload in _iter_sse_payloads(response):
        trimmed = payload.strip()
        if not trimmed:
            continue
        if trimmed in {"[DONE]", "DONE", "done"}:
            break
        try:
            parsed = json.loads(trimmed)
        except json.JSONDecodeError:
            continue
        if isinstance(parsed, dict):
            yield parsed

class AnalogAIClient:
    def __init__(
        self,
        api_key: Optional[str] = None,
        agent_id: Optional[str] = None,
        base_url: Optional[str] = None,
        timeout: int = 60,
    ) -> None:
        self.api_key = api_key or os.getenv("ANALOGAI_API_KEY")
        self.agent_id = agent_id or os.getenv("ANALOGAI_AGENT_ID")
        self.base_url = base_url or _resolve_base_url()
        self.timeout = timeout

        if not self.agent_id:
            raise ValueError("Missing agent id. Set ANALOGAI_AGENT_ID in environment or pass agent_id.")

        self.base_url = self.base_url.rstrip("/")

    def _headers(self, extra: Optional[Dict[str, str]] = None) -> dict:
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        if extra:
            headers.update(extra)
        return headers

    def health(self) -> Dict[str, Any]:
        url = f"{self.base_url}/health"
        response = requests.get(url, headers=self._headers(), timeout=self.timeout)
        response.raise_for_status()
        return response.json()

    def generate_completion(
        self,
        message_text: str,
        auth_type: str = "user",
        messages: Optional[Iterable[Dict[str, str]]] = None,
        *,
        state: Optional[Dict[str, Any]] = None,
        tools: Optional[Iterable[Dict[str, Any]]] = None,
        thread_id: Optional[str] = None,
        run_id: Optional[str] = None,
        mode: Optional[str] = None,
    ) -> str:
        url = f"{self.base_url}/agent/{self.agent_id}/chat"
        payload = _build_payload(
            message_text=message_text,
            auth_type=auth_type,
            messages=messages,
            state=state,
            tools=tools,
            thread_id=thread_id,
            run_id=run_id,
            mode=mode,
        )
        headers = self._headers(
            {
                "Accept": "text/event-stream",
                **({"x-mode": mode} if mode else {}),
            }
        )
        collected: list[str] = []
        with requests.post(
            url,
            headers=headers,
            json=payload,
            stream=True,
            timeout=self.timeout,
        ) as response:
            response.raise_for_status()
            for event in _iter_sse_events(response):
                event_type = str(event.get("type") or "").lower()
                if event_type == "run_error":
                    raise RuntimeError(str(event.get("message") or "Run error"))
                if event_type == "text_message_chunk":
                    delta = str(event.get("delta") or "")
                    if delta:
                        collected.append(delta)
                elif event_type == "text_message_content":
                    content = str(event.get("content") or "")
                    if content:
                        collected.append(content)

        content = "".join(collected).strip()
        return content

    def generate_streaming_completion(
        self,
        message_text: str,
        auth_type: str = "user",
        messages: Optional[Iterable[Dict[str, str]]] = None,
        *,
        state: Optional[Dict[str, Any]] = None,
        tools: Optional[Iterable[Dict[str, Any]]] = None,
        thread_id: Optional[str] = None,
        run_id: Optional[str] = None,
        mode: Optional[str] = None,
    ) -> Generator[str, None, None]:
        url = f"{self.base_url}/agent/{self.agent_id}/chat"
        payload = _build_payload(
            message_text=message_text,
            auth_type=auth_type,
            messages=messages,
            state=state,
            tools=tools,
            thread_id=thread_id,
            run_id=run_id,
            mode=mode,
        )
        headers = self._headers(
            {
                "Accept": "text/event-stream",
                **({"x-mode": mode} if mode else {}),
            }
        )
        with requests.post(
            url,
            headers=headers,
            json=payload,
            stream=True,
            timeout=self.timeout,
        ) as response:
            response.raise_for_status()
            for event in _iter_sse_events(response):
                event_type = str(event.get("type") or "").lower()
                if event_type == "run_error":
                    raise RuntimeError(str(event.get("message") or "Run error"))
                if event_type == "text_message_chunk":
                    delta = str(event.get("delta") or "")
                    if delta:
                        yield delta
                elif event_type == "text_message_content":
                    content = str(event.get("content") or "")
                    if content:
                        yield content


__all__ = ["AnalogAIClient"]
