from __future__ import annotations

import json
from abc import ABC, abstractmethod
from dataclasses import dataclass
from threading import Lock
from typing import Protocol, runtime_checkable

from ..models import Usage


@dataclass(slots=True)
class StreamEvent:
    """Streaming event envelope persisted for replay/resume."""

    event_id: int
    data: str
    usage_delta: Usage


class ReplayStore(ABC):
    """Abstract storage for streaming events to support resume semantics."""

    @abstractmethod
    def append(self, request_id: str, event: StreamEvent) -> None:
        """Persist a streaming event for the given request."""

    @abstractmethod
    def replay_from(self, request_id: str, after_event_id: int) -> list[StreamEvent]:
        """Return all events strictly after the provided event id."""

    def replay_all(self, request_id: str) -> list[StreamEvent]:
        """Convenience helper to replay from the beginning."""
        return self.replay_from(request_id, after_event_id=0)

    def prune(self, request_id: str, keep_last: int) -> None:
        """Optional hook for stores that support truncation. In-memory default no-op."""
        # Default implementation is a no-op.
        return None


class InMemoryReplayStore(ReplayStore):
    """Naive replay store keeping events in memory per request id."""

    def __init__(self, max_events: int = 256):
        self._max_events = max_events
        self._events: dict[str, list[StreamEvent]] = {}
        self._lock = Lock()

    def append(self, request_id: str, event: StreamEvent) -> None:
        with self._lock:
            buf = self._events.setdefault(request_id, [])
            buf.append(event)
            if len(buf) > self._max_events:
                del buf[: len(buf) - self._max_events]

    def replay_from(self, request_id: str, after_event_id: int) -> list[StreamEvent]:
        with self._lock:
            buf = list(self._events.get(request_id, ()))
        return [event for event in buf if event.event_id > after_event_id]

    def prune(self, request_id: str, keep_last: int) -> None:
        with self._lock:
            buf = self._events.get(request_id)
            if not buf:
                return
            if keep_last <= 0:
                self._events.pop(request_id, None)
                return
            if len(buf) > keep_last:
                del buf[: len(buf) - keep_last]


class RedisReplayStore(ReplayStore):
    """Redis-backed replay storage using a sorted set per request id."""

    @runtime_checkable
    class _RedisClient(Protocol):
        def zadd(self, name: str, mapping: dict[str, int]) -> int: ...

        def zrangebyscore(self, name: str, min: int, max: str) -> list[object]: ...

        def zcard(self, name: str) -> int: ...

        def zremrangebyrank(self, name: str, start: int, end: int) -> int: ...

        def expire(self, name: str, time: int) -> bool: ...

        def delete(self, *names: str) -> int: ...

    def __init__(self, client: _RedisClient, ttl_seconds: int = 300, prefix: str = "replay"):
        self._client = client
        self._ttl = ttl_seconds
        self._prefix = prefix

    def _key(self, request_id: str) -> str:
        return f"{self._prefix}:{request_id}"

    @staticmethod
    def _encode_event(event: StreamEvent) -> str:
        payload = {
            "event_id": int(event.event_id),
            "data": event.data,
            "usage_delta": {
                "prompt_tokens": int(event.usage_delta.prompt_tokens),
                "completion_tokens": int(event.usage_delta.completion_tokens),
                "total_tokens": int(event.usage_delta.total_tokens),
                "cost_usd": event.usage_delta.cost_usd,
            },
        }
        return json.dumps(payload, separators=(",", ":"), ensure_ascii=False)

    @staticmethod
    def _decode_event(raw: object) -> StreamEvent:
        if isinstance(raw, bytes):
            text = raw.decode("utf-8")
        else:
            text = str(raw)
        payload = json.loads(text)
        if not isinstance(payload, dict):
            raise TypeError("Replay payload must be an object")
        event_id = payload.get("event_id")
        data = payload.get("data")
        usage_raw = payload.get("usage_delta")
        if not isinstance(event_id, int):
            raise TypeError("Replay payload missing integer event_id")
        if not isinstance(data, str):
            raise TypeError("Replay payload missing string data")
        if not isinstance(usage_raw, dict):
            raise TypeError("Replay payload missing usage_delta mapping")
        usage = Usage(
            prompt_tokens=int(usage_raw.get("prompt_tokens", 0)),
            completion_tokens=int(usage_raw.get("completion_tokens", 0)),
            total_tokens=int(usage_raw.get("total_tokens", 0)),
            cost_usd=usage_raw.get("cost_usd"),
        )
        return StreamEvent(event_id=event_id, data=data, usage_delta=usage)

    def append(self, request_id: str, event: StreamEvent) -> None:
        key = self._key(request_id)
        member = self._encode_event(event)
        self._client.zadd(key, {member: int(event.event_id)})
        self._client.expire(key, int(self._ttl))

    def replay_from(
        self, request_id: str, after_event_id: int
    ) -> list[StreamEvent]:
        key = self._key(request_id)
        start = int(after_event_id) + 1
        raw_events = self._client.zrangebyscore(key, start, "+inf")
        return [self._decode_event(raw) for raw in raw_events]

    def prune(self, request_id: str, keep_last: int) -> None:
        key = self._key(request_id)
        keep_last = int(keep_last)
        if keep_last <= 0:
            self._client.delete(key)
            return
        total = int(self._client.zcard(key))
        excess = total - keep_last
        if excess > 0:
            self._client.zremrangebyrank(key, 0, excess - 1)
        self._client.expire(key, int(self._ttl))


__all__ = [
    "StreamEvent",
    "ReplayStore",
    "InMemoryReplayStore",
    "RedisReplayStore",
]
