"""Stream diagnostics — per-attempt counters, exception chains, retry logging.

When a streaming request dies mid-response these helpers record WHY (which CF
edge / OpenRouter downstream served it, bytes+chunks before the drop, HTTP
status, underlying httpx error class) to ``agent.log`` in full and to the
user-facing status line compactly. ``run_agent`` keeps thin forwarders so
existing call sites and tests patching ``run_agent.<helper>`` keep working.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Lowercased upstream headers captured per attempt for post-hoc analysis.
STREAM_DIAG_HEADERS = (
    "cf-ray",
    "cf-cache-status",
    "x-openrouter-provider",
    "x-openrouter-model",
    "x-openrouter-id",
    "x-request-id",
    "x-vercel-id",
    "via",
    "server",
    "x-forwarded-for",
)


def stream_diag_init() -> Dict[str, Any]:
    """Fresh per-attempt diagnostic dict; mutated in place by the streaming functions and read by the retry block."""
    return {
        "started_at": time.time(),
        "first_chunk_at": None,
        "chunks": 0,
        "bytes": 0,
        "headers": {},
        "http_status": None,
    }


def stream_diag_capture_response(agent: Any, diag: Dict[str, Any], http_response: Any) -> None:
    """Snapshot headers + HTTP status at stream open (so they survive a drop before the first chunk). Best-effort."""
    if http_response is None or not isinstance(diag, dict):
        return
    try:
        diag["http_status"] = getattr(http_response, "status_code", None)
    except Exception:
        pass
    try:
        headers = getattr(http_response, "headers", None) or {}
        captured: Dict[str, str] = {}
        # Per-agent override of the headers list (back-compat).
        for name in getattr(agent, "_STREAM_DIAG_HEADERS", STREAM_DIAG_HEADERS):
            try:
                val = headers.get(name)
                if val:
                    captured[name] = str(val)[:120]  # keep log lines bounded
            except Exception:
                continue
        diag["headers"] = captured
    except Exception:
        pass


def flatten_exception_chain(error: BaseException) -> str:
    """Compact ``Outer(msg) <- Inner(msg) <- ...`` rendering.

    The OpenAI SDK wraps httpx errors so only the wrapper class is visible at the
    catch site; the inner RemoteProtocolError/ConnectError/ReadError says WHY the
    stream died. Walks ``__cause__`` then ``__context__`` (deduped, max 4 deep).
    """
    seen: List[BaseException] = []
    link: Optional[BaseException] = error
    while link is not None and len(seen) < 4:
        if link in seen:
            break
        seen.append(link)
        nxt = getattr(link, "__cause__", None) or getattr(link, "__context__", None)
        if nxt is None or nxt is link:
            break
        link = nxt
    parts: List[str] = []
    for e in seen:
        msg = str(e).strip().replace("\n", " ")
        if len(msg) > 140:
            msg = msg[:140] + "…"
        parts.append(f"{type(e).__name__}({msg})" if msg else type(e).__name__)
    return " <- ".join(parts) if parts else type(error).__name__


def log_stream_retry(
    agent: Any,
    *,
    kind: str,
    error: BaseException,
    attempt: int,
    max_attempts: int,
    mid_tool_call: bool,
    diag: Optional[Dict[str, Any]] = None,
) -> None:
    """Structured WARNING to ``agent.log`` for a transient stream drop + retry.

    Always logged regardless of UI verbosity (subagent retries no longer spam the
    parent's terminal but keep full detail here). With *diag*, also records upstream
    headers, HTTP status, bytes/chunks streamed, elapsed and TTFB on the dying attempt —
    enough to tell "one CF edge / downstream provider" from "random across runs".
    """
    try:
        try:
            _summary = agent._summarize_api_error(error)
        except Exception:
            _summary = str(error)
        if _summary and len(_summary) > 240:
            _summary = _summary[:240] + "…"
        try:
            _chain = flatten_exception_chain(error)
        except Exception:
            _chain = type(error).__name__

        _now = time.time()
        _bytes = 0
        _chunks = 0
        _elapsed = 0.0
        _ttfb = None
        _headers_repr = "-"
        _http_status = "-"
        if isinstance(diag, dict):
            try:
                _bytes = int(diag.get("bytes") or 0)
                _chunks = int(diag.get("chunks") or 0)
                _started = float(diag.get("started_at") or _now)
                _elapsed = max(0.0, _now - _started)
                _first = diag.get("first_chunk_at")
                if _first is not None:
                    _ttfb = max(0.0, float(_first) - _started)
                headers = diag.get("headers") or {}
                if isinstance(headers, dict) and headers:
                    _headers_repr = " ".join(f"{k}={v}" for k, v in headers.items())
                if diag.get("http_status") is not None:
                    _http_status = str(diag.get("http_status"))
            except Exception:
                pass

        logger.warning(
            "Stream %s on attempt %s/%s — retrying. "
            "subagent_id=%s depth=%s provider=%s base_url=%s "
            "error_type=%s error=%s "
            "chain=%s "
            "http_status=%s bytes=%d chunks=%d elapsed=%.2fs ttfb=%s "
            "upstream=[%s]",
            kind,
            attempt,
            max_attempts,
            getattr(agent, "_subagent_id", None) or "-",
            getattr(agent, "_delegate_depth", 0),
            agent.provider or "-",
            agent.base_url or "-",
            type(error).__name__,
            _summary,
            _chain,
            _http_status,
            _bytes,
            _chunks,
            _elapsed,
            f"{_ttfb:.2f}s" if _ttfb is not None else "-",
            _headers_repr,
            extra={"mid_tool_call": mid_tool_call},
        )
    except Exception:
        logger.debug("stream-retry log emit failed", exc_info=True)


def emit_stream_drop(
    agent: Any,
    *,
    error: BaseException,
    attempt: int,
    max_attempts: int,
    mid_tool_call: bool,
    diag: Optional[Dict[str, Any]] = None,
) -> None:
    """One compact user-visible status line for a stream drop+retry, plus the full WARNING via log_stream_retry.

    Subagent lines get a ``[subagent-N]`` prefix from ``log_prefix``. ``after Xs``
    distinguishes "couldn't connect" (0s) from "died mid-stream" (idle-kill / proxy timeout).
    """
    kind = "drop mid tool-call" if mid_tool_call else "drop"
    log_stream_retry(
        agent, kind=kind, error=error, attempt=attempt, max_attempts=max_attempts, mid_tool_call=mid_tool_call, diag=diag
    )
    provider = agent.provider or "provider"
    _suffix = ""
    if isinstance(diag, dict):
        try:
            started = diag.get("started_at")
            if started is not None:
                _suffix = f" after {max(0.0, time.time() - float(started)):.1f}s"
        except Exception:
            pass
    try:
        agent._buffer_status(
            f"⚠️ {provider} stream {kind} ({type(error).__name__}){_suffix} "
            f"— reconnecting, retry {attempt}/{max_attempts}"
        )
        agent._touch_activity(f"stream retry {attempt}/{max_attempts} after {type(error).__name__}")
    except Exception:
        pass


__all__ = [
    "STREAM_DIAG_HEADERS",
    "stream_diag_init",
    "stream_diag_capture_response",
    "flatten_exception_chain",
    "log_stream_retry",
    "emit_stream_drop",
]