"""
UI-agnostic retry / failure-state logic.

`stream_with_retries` turns a stream (or its failure) into a sequence of
Events. A UI only has to render events; it never catches exceptions.
"""

import math
import random
import time
from dataclasses import dataclass
from typing import Callable, Iterator, Optional

from providers import (ProviderError, RateLimitError, RefusalError,
                       RequestTimeoutError, StallError, with_watchdog)


@dataclass
class Event:
    kind: str            # start | chunk | retry_wait | done | stalled | refused | error
    text: str = ""       # accumulated response text so far
    message: str = ""
    seconds: int = 0     # for retry_wait: seconds remaining
    attempt: int = 0
    status_code: Optional[int] = None


def backoff_seconds(attempt: int, retry_after: Optional[float] = None,
                    base: float = 2.0, cap: float = 30.0) -> int:
    """Honor the server's Retry-After if present; otherwise exponential
    backoff with jitter (so many clients don't retry in lockstep)."""
    if retry_after:
        return math.ceil(retry_after)
    return math.ceil(min(cap, base * 2 ** (attempt - 1)) * random.uniform(0.75, 1.25))


def stream_with_retries(make_stream: Callable[[], Iterator[str]], *,
                        max_retries: int = 3,
                        stall_seconds: Optional[float] = None,
                        sleep: Callable[[float], None] = time.sleep) -> Iterator[Event]:
    attempt = 0
    while True:
        attempt += 1
        text = ""
        yield Event("start", attempt=attempt)
        try:
            stream = make_stream()
            if stall_seconds:
                stream = with_watchdog(stream, stall_seconds)
            for chunk in stream:
                text += chunk
                yield Event("chunk", text=text, attempt=attempt)
            yield Event("done", text=text, attempt=attempt)
            return

        except RateLimitError as e:
            if attempt > max_retries:
                yield Event("error", text=text, attempt=attempt, status_code=429,
                            message=f"Still rate limited after {attempt} attempts. Giving up.")
                return
            wait = backoff_seconds(attempt, e.retry_after)
            source = "Retry-After header" if e.retry_after else "exponential backoff + jitter"
            yield Event("retry_wait", text=text, attempt=attempt, seconds=wait, status_code=429,
                        message=f"429 on attempt {attempt}. Waiting {wait}s ({source}).")
            for left in range(wait, 0, -1):
                yield Event("retry_wait", text=text, attempt=attempt, seconds=left,
                            status_code=429, message="")
                sleep(1)

        except StallError as e:
            yield Event("stalled", text=text, attempt=attempt, message=str(e))
            return

        except RefusalError as e:
            yield Event("refused", text=text, attempt=attempt, message=str(e))
            return

        except RequestTimeoutError as e:
            yield Event("error", text=text, attempt=attempt,
                        message=f"No response from the API: {e}")
            return

        except ProviderError as e:
            yield Event("error", text=text, attempt=attempt,
                        status_code=e.status_code, message=str(e))
            return