"""SSE transport helpers for the kernel model client (docs/kernel_spec.md §2).

Idle-deadline iteration and retry classification ported from agent_loop_stream.py.
"""
from __future__ import annotations

from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from typing import Any

import httpx

# ─── Timeouts / retry ────────────────────────────────────────────────────────
RETRY_ATTEMPTS = 3
RETRY_BASE_DELAY = 1.0   # seconds; exponential backoff
RETRY_MAX_DELAY = 30.0

IDLE_TIMEOUT = 180            # base max silence between SSE chunks
IDLE_TIMEOUT_PER_10K_TOKENS = 40
IDLE_TIMEOUT_MAX = 600
FIRST_CHUNK_TIMEOUT = 120     # time-to-first-token cap

SSE_DATA_PREFIX = "data:"


class SSEIdleTimeout(Exception):
    """No SSE chunk arrived within the idle deadline."""


def is_retryable_error(exc: BaseException) -> bool:
    """True for connection failures, timeouts, 429 and 5xx."""
    if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout, httpx.ReadTimeout,
                        httpx.PoolTimeout, httpx.RemoteProtocolError)):
        return True
    status = getattr(exc, "status_code", None)
    if status is not None:
        return status == 429 or (500 <= int(status) < 600)
    text = str(exc).lower()
    return any(s in text for s in ("timeout", "connection", "429", "502", "503"))


def scaled_idle_timeout(estimated_tokens: int) -> int:
    """Idle timeout scaled to context size (large prompts eval slowly locally)."""
    extra = int((max(0, estimated_tokens) / 10_000) * IDLE_TIMEOUT_PER_10K_TOKENS)
    return min(IDLE_TIMEOUT + extra, IDLE_TIMEOUT_MAX)


def iter_with_timeout(
    it: Iterator[Any],
    *,
    idle_timeout: float,
    first_chunk_timeout: float | None = None,
) -> Iterator[Any]:
    """Yield from *it*; raise SSEIdleTimeout when a chunk takes too long.

    *first_chunk_timeout* applies only to the first chunk (time-to-first-token);
    later chunks use *idle_timeout*.
    """
    executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="kernel_stream")
    first = True
    try:
        while True:
            deadline = (
                first_chunk_timeout
                if first and first_chunk_timeout is not None
                else idle_timeout
            )
            future = executor.submit(next, it)
            try:
                yield future.result(timeout=deadline)
                first = False
            except StopIteration:
                return
            except TimeoutError:
                future.cancel()
                raise SSEIdleTimeout(
                    f"LLM stream stalled — no chunk received for {deadline}s"
                ) from None
    finally:
        executor.shutdown(wait=False)
