"""Keiro client — low-level hosted gateway access for external Keiro users.

Usage::

    from keiro import models

    # Preferred prompt-first facade
    print(models("eb1-preview", "What is machine learning?"))

    # Lower-level HTTP client for raw chat/responses calls
    from keiro import Client
    client = Client()
    response = client.chat(
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Explain quantum computing."},
        ],
        model="eb1-preview",
    )
    print(response["choices"][0]["message"]["content"])

    # Batch concurrent chat calls
    prompts = ["What is ML?", "What is AI?", "What is DL?"]
    results = client.batch(prompts, model="eb1-preview", max_workers=3)

The lightweight package prefers ``keiro.models`` for prompt-first testing and
keeps ``Client`` as the lower-level escape hatch. No Ember dependency
required. Credentials are read from ``~/.keiro/credentials`` (written by
``keiro setup``).
"""

from __future__ import annotations

import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Any, Iterator, Mapping

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from keiro.credentials import load_credentials
from keiro.defaults import DEFAULT_API_KEY, DEFAULT_GATEWAY_URL, DEFAULT_MODEL


class KeirError(Exception):
    """Error returned by the Keiro gateway."""

    def __init__(
        self,
        message: str,
        *,
        status_code: int,
        error_type: str | None = None,
        code: str | None = None,
    ) -> None:
        self.status_code = status_code
        self.error_type = error_type
        self.code = code
        super().__init__(message)


@dataclass(frozen=True)
class BatchFailure:
    """Single failed batch request."""

    index: int
    prompt: str
    error: Exception


class BatchError(Exception):
    """Aggregate error for ``Client.batch()`` failures."""

    def __init__(
        self,
        failures: list[BatchFailure],
        partial_results: list[str | None],
    ) -> None:
        self.failures = failures
        self.partial_results = partial_results
        message = ", ".join(f"index={failure.index}: {failure.error}" for failure in failures)
        super().__init__(f"One or more batch requests failed: {message}")


def _raise_for_error(resp: requests.Response) -> None:
    """Raise KeirError with structured info from gateway error responses.

    Handles both nested (OpenAI-compatible) and flat error formats:
    - Nested: ``{"error": {"message": ..., "type": ..., "code": ...}}``
    - Flat:   ``{"error": "...", "message": "..."}``
    """
    if resp.status_code < 400:
        return
    try:
        body = resp.json()
        err = body.get("error", {})
        if isinstance(err, dict):
            # Nested OpenAI-compatible format
            message = err.get("message", resp.text)
            error_type = err.get("type")
            code = err.get("code")
        else:
            # Flat format: {"error": "type_string", "message": "..."}
            message = body.get("message", str(err) or resp.text)
            error_type = str(err) if err else None
            code = str(err) if err else None
    except (ValueError, KeyError):
        message = resp.text or f"HTTP {resp.status_code}"
        error_type = None
        code = None
    raise KeirError(
        message,
        status_code=resp.status_code,
        error_type=error_type,
        code=code,
    )


def _build_session(retries: int = 3, pool_maxsize: int = 100) -> requests.Session:
    """Build a requests session with retry on transient failures.

    Only retries on 429 (rate limit) and 503/504 (transient server issues).
    Does NOT retry on 500/502 because those indicate upstream provider
    failures that retrying won't resolve.
    """
    session = requests.Session()
    retry = Retry(
        total=retries,
        backoff_factor=1.0,
        status_forcelist=[429, 503, 504],
        allowed_methods=["POST", "GET"],
        respect_retry_after_header=True,
    )
    adapter = HTTPAdapter(
        max_retries=retry,
        pool_connections=pool_maxsize,
        pool_maxsize=pool_maxsize,
    )
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session


# EB1 ensemble runs 3+ models in parallel + judge synthesis.
# 120s is too short for complex queries. Match defaults.yaml read_standard.
_DEFAULT_TIMEOUT: float = 600.0


class Client:
    """Keiro API client with retry, streaming, and batch support.

    Args:
        api_key: API key. Falls back to ``KEIRO_API_KEY`` env var,
            then ``~/.keiro/credentials``.
        base_url: Gateway URL. Falls back to ``KEIRO_BASE_URL`` env var,
            then ``~/.keiro/credentials``, then the default gateway URL.
        timeout: Request timeout in seconds. Default 600s (10 min) to
            accommodate EB1 ensemble latency on complex queries.
        retries: Number of automatic retries on transient errors (429, 5xx).
        pool_maxsize: Maximum size of the underlying HTTP connection pool.
    """

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = _DEFAULT_TIMEOUT,
        retries: int = 3,
        pool_maxsize: int = 100,
    ) -> None:
        resolved_key, resolved_url = _resolve_credentials(api_key, base_url)
        self._api_key = resolved_key
        self._base_url = resolved_url
        self._timeout = timeout
        self._session = _build_session(retries, pool_maxsize)

    def close(self) -> None:
        """Close the underlying HTTP session."""
        self._session.close()

    def __enter__(self) -> "Client":
        return self

    def __exit__(self, *_args: object) -> None:
        self.close()

    def _headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

    def _url(self, path: str) -> str:
        return _build_url(self._base_url, path)

    def chat(
        self,
        messages: list[dict[str, Any]],
        *,
        model: str = DEFAULT_MODEL,
        stream: bool = False,
        timeout: float | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Send a chat completion request.

        Args:
            messages: Conversation messages in OpenAI format
                (``[{"role": "user", "content": "..."}]``).
            model: Model name (default ``"eb1"``).
            stream: If True, returns after collecting the full streamed
                response. Use ``chat_stream()`` for incremental chunks.
            timeout: Override per-request timeout.
            **kwargs: Additional API parameters (``max_tokens``,
                ``temperature``, etc.).

        Returns:
            Chat completion response dict with ``choices``, ``usage``, etc.

        Raises:
            KeirError: On gateway errors (4xx/5xx).
            ValueError: If no API key is configured.
        """
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            **kwargs,
        }
        if stream:
            return self._collect_stream(payload, timeout)

        resp = self._session.post(
            self._url("/v1/chat/completions"),
            headers=self._headers(),
            json=payload,
            timeout=timeout or self._timeout,
        )
        _raise_for_error(resp)
        return resp.json()

    def chat_stream(
        self,
        messages: list[dict[str, Any]],
        *,
        model: str = DEFAULT_MODEL,
        timeout: float | None = None,
        **kwargs: Any,
    ) -> Iterator[dict[str, Any]]:
        """Stream a chat completion, yielding incremental chunks.

        Args:
            messages: Conversation messages.
            model: Model name.
            timeout: Override timeout.
            **kwargs: Additional API parameters.

        Yields:
            SSE chunk dicts. The final chunk has ``finish_reason`` set.

        Raises:
            KeirError: On gateway errors.
        """
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": True,
            **kwargs,
        }
        resp = self._session.post(
            self._url("/v1/chat/completions"),
            headers=self._headers(),
            json=payload,
            timeout=timeout or self._timeout,
            stream=True,
        )
        _raise_for_error(resp)
        yield from _iter_parsed_sse(resp)

    def responses(
        self,
        payload: Mapping[str, Any],
        *,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        """Send a Responses API request and return the JSON payload."""
        if not isinstance(payload, Mapping):
            raise TypeError("payload must be a mapping")

        resp = self._session.post(
            self._url("/v1/responses"),
            headers=self._headers(),
            json=dict(payload),
            timeout=timeout or self._timeout,
        )
        _raise_for_error(resp)
        return resp.json()

    def responses_stream(
        self,
        payload: Mapping[str, Any],
        *,
        timeout: float | None = None,
    ) -> Iterator[dict[str, Any]]:
        """Stream a Responses API request, yielding parsed SSE events."""
        if not isinstance(payload, Mapping):
            raise TypeError("payload must be a mapping")

        resp = self._session.post(
            self._url("/v1/responses"),
            headers=self._headers(),
            json={**dict(payload), "stream": True},
            timeout=timeout or self._timeout,
            stream=True,
        )
        _raise_for_error(resp)
        yield from _iter_parsed_sse(resp)

    def get_models(
        self,
        *,
        provider: str | None = None,
        status: str | None = None,
        legacy: bool = False,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        """Fetch the gateway model catalog."""
        params: dict[str, str] = {}
        if provider:
            params["provider"] = provider
        if status:
            params["status"] = status
        if legacy:
            params["legacy"] = "true"

        resp = self._session.get(
            self._url("/v1/models"),
            headers=self._headers(),
            params=params or None,
            timeout=timeout or self._timeout,
        )
        _raise_for_error(resp)
        return resp.json()

    def batch(
        self,
        prompts: list[str],
        *,
        model: str = DEFAULT_MODEL,
        system: str | None = None,
        max_workers: int = 5,
        timeout: float | None = None,
        raise_on_error: bool = True,
        **kwargs: Any,
    ) -> list[str | None]:
        """Run multiple prompts concurrently and return responses.

        Args:
            prompts: List of user prompts.
            model: Model name applied to all requests.
            system: Optional system prompt applied to all requests.
            max_workers: Maximum concurrent requests.
            timeout: Per-request timeout override.
            **kwargs: Additional API parameters.

        Returns:
            List of response texts, in the same order as prompts.
            When ``raise_on_error=False``, failed requests return ``None`` at
            their index. Otherwise a ``BatchError`` is raised with partial
            results and structured failure details.
        """
        results: list[str | None] = [None] * len(prompts)
        failures: list[BatchFailure] = []

        def _call(index: int, prompt: str) -> tuple[int, str | None, Exception | None]:
            messages: list[dict[str, str]] = []
            if system:
                messages.append({"role": "system", "content": system})
            messages.append({"role": "user", "content": prompt})
            try:
                resp = self.chat(messages, model=model, timeout=timeout, **kwargs)
                text = resp["choices"][0]["message"]["content"]
                return (index, text, None)
            except Exception as exc:
                return (index, None, exc)

        workers = min(max_workers, len(prompts))
        with ThreadPoolExecutor(max_workers=workers) as pool:
            futures = {pool.submit(_call, i, p): i for i, p in enumerate(prompts)}
            for future in as_completed(futures):
                idx, text, error = future.result()
                results[idx] = text
                if error is not None:
                    failures.append(BatchFailure(index=idx, prompt=prompts[idx], error=error))

        if failures and raise_on_error:
            raise BatchError(
                sorted(failures, key=lambda failure: failure.index),
                results,
            )

        return results

    def _collect_stream(
        self,
        payload: dict[str, Any],
        timeout: float | None,
    ) -> dict[str, Any]:
        """Stream a request but collect and return the full response."""
        payload["stream"] = True
        chunks: list[str] = []
        stream_payload = dict(payload)
        model = str(stream_payload.pop("model", ""))
        messages = stream_payload.pop("messages", [])
        stream_payload.pop("stream", None)
        usage: dict[str, int] = {}

        for chunk in self.chat_stream(
            messages,
            model=model,
            timeout=timeout,
            **stream_payload,
        ):
            event_name = chunk.get("event")
            if event_name == "delta":
                text = chunk.get("text")
                if isinstance(text, str) and text:
                    chunks.append(text)
                continue
            if event_name == "done":
                usage_payload = chunk.get("usage")
                if isinstance(usage_payload, dict):
                    usage = usage_payload
                done_text = chunk.get("text")
                if isinstance(done_text, str) and done_text and not chunks:
                    chunks.append(done_text)
                continue
            if event_name == "header":
                continue
            choices = chunk.get("choices", [])
            if choices:
                delta = choices[0].get("delta", {})
                content = delta.get("content", "")
                if content:
                    chunks.append(content)
            if "usage" in chunk:
                usage = chunk["usage"]

        text = "".join(chunks)
        return {
            "choices": [{"index": 0, "message": {"role": "assistant", "content": text}}],
            "model": model,
            "usage": usage,
        }


def complete(
    prompt: str,
    **kwargs: Any,
) -> str:
    """Send a prompt to the Keiro gateway and return the response text.

    This is the simplest interface — one prompt in, one response out.

    Args:
        prompt: The user message to send.
        **kwargs: Passed to :class:`Client.chat`. Supported keyword
            arguments: ``model``, ``api_key``, ``base_url``, ``timeout``,
            ``system``, ``max_tokens``.

    Returns:
        The model's response text.

    Raises:
        ValueError: If no API key is configured.
        KeirError: On gateway errors.
    """
    api_key = kwargs.pop("api_key", None)
    base_url = kwargs.pop("base_url", None)
    timeout = kwargs.pop("timeout", None)
    system = kwargs.pop("system", None)

    messages: list[dict[str, str]] = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    with Client(api_key=api_key, base_url=base_url) as client:
        resp = client.chat(messages, timeout=timeout, **kwargs)
    return resp["choices"][0]["message"]["content"]


def _iter_sse_events(resp: requests.Response) -> Iterator[tuple[str | None, str]]:
    """Parse an SSE response into ``(event, data)`` tuples."""
    event_name: str | None = None
    data_lines: list[str] = []

    for raw_line in resp.iter_lines(decode_unicode=True):
        if raw_line is None:
            continue
        line = raw_line if isinstance(raw_line, str) else raw_line.decode("utf-8")
        if line == "":
            if data_lines:
                yield event_name, "\n".join(data_lines)
            event_name = None
            data_lines = []
            continue
        if line.startswith(":"):
            continue
        if line.startswith("event:"):
            event_name = line[6:].strip() or None
            continue
        if line.startswith("data:"):
            data_lines.append(line[5:].lstrip())

    if data_lines:
        yield event_name, "\n".join(data_lines)


def _raise_stream_error(payload: dict[str, Any], *, status_code: int) -> None:
    """Raise ``KeirError`` from a streamed error payload."""
    message = str(payload.get("message") or "streaming request failed")
    error_type = None
    code = payload.get("code")

    error_obj = payload.get("error")
    if isinstance(error_obj, dict):
        message = str(error_obj.get("message") or message)
        if error_obj.get("type") is not None:
            error_type = str(error_obj.get("type"))
        if code is None and error_obj.get("code") is not None:
            code = str(error_obj.get("code"))
    elif isinstance(error_obj, str):
        error_type = error_obj
        if code is None:
            code = error_obj

    raise KeirError(
        message,
        status_code=status_code,
        error_type=error_type,
        code=str(code) if code is not None else None,
    )


def _iter_parsed_sse(resp: requests.Response) -> Iterator[dict[str, Any]]:
    """Parse an SSE response into typed JSON payloads.

    Handles ``[DONE]`` sentinel, malformed lines, and streamed error
    events. Closes the response when iteration completes or is abandoned.
    """
    try:
        for event_name, data in _iter_sse_events(resp):
            if data == "[DONE]":
                return
            try:
                payload_obj = json.loads(data)
            except json.JSONDecodeError:
                continue
            if not isinstance(payload_obj, dict):
                continue
            if event_name:
                payload_obj.setdefault("event", event_name)
            if event_name == "error":
                _raise_stream_error(payload_obj, status_code=resp.status_code)
            yield payload_obj
    finally:
        resp.close()


def _resolve_credentials(
    explicit_key: str | None,
    explicit_url: str | None,
) -> tuple[str, str]:
    """Resolve API key and base URL from args > env > credentials file > defaults."""
    key = explicit_key or os.environ.get("KEIRO_API_KEY")
    url = explicit_url or os.environ.get("KEIRO_BASE_URL")

    if not key or not url:
        creds = load_credentials()
        key = key or creds.get("api_key")
        url = url or creds.get("base_url")

    key = key or DEFAULT_API_KEY

    return key, (url or DEFAULT_GATEWAY_URL).rstrip("/")


def _build_url(base: str, path: str) -> str:
    """Build a full endpoint URL, avoiding duplicate ``/v1`` prefixes."""
    normalized_base = base.rstrip("/")
    if path.startswith("/v1/") and normalized_base.endswith("/v1"):
        path = path[3:]
    return f"{normalized_base}{path}"
