"""Greylist retry queue.

Real mail servers often answer the first RCPT for an unknown sender with a
**4xx** (greylisting / "try again later"). That is not a verdict — it is an
instruction to come back. Without a retry, a perfectly valid mailbox on such a
domain stays permanently "risky", which is exactly the false-negative that makes
a validator feel unreliable.

Two rules keep this safe:

  * **Only 4xx is retried.** A 5xx is final, a 250 is done. A *transport*
    failure (connect timeout / unreachable) is systemic to the MX — re-queuing
    every address behind a dead host would flood the queue for no gain, so the
    MX-level circuit breaker handles that instead.
  * **Bounded attempts with exponential backoff.** A queue that retries forever
    is a queue that never drains; entries are dropped after `max_attempts`.

The store is pluggable (Redis in production, memory in tests) and every store
error degrades to "no queue" rather than breaking a validation run.
"""
from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass
from typing import Awaitable, Callable, Optional, Protocol

from ..models import SMTPResult


def is_greylist_result(smtp: Optional[SMTPResult]) -> bool:
    """True only for a genuine SMTP 4xx — the retryable case."""
    if smtp is None:
        return False
    code = smtp.smtp_code
    return code is not None and 400 <= code < 500


@dataclass
class RetryEntry:
    email: str
    mx_host: str
    attempts: int = 0
    due_at: float = 0.0
    last_code: Optional[int] = None

    def as_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> "RetryEntry":
        return cls(
            email=data["email"],
            mx_host=data["mx_host"],
            attempts=int(data.get("attempts", 0)),
            due_at=float(data.get("due_at", 0.0)),
            last_code=data.get("last_code"),
        )


class RetryStore(Protocol):
    async def add(self, entry: RetryEntry) -> None: ...
    async def due(self, now: float, limit: int) -> list[RetryEntry]: ...
    async def remove(self, email: str) -> None: ...


class MemoryRetryStore:
    """In-process store. Adequate for a single worker or tests."""

    def __init__(self) -> None:
        self._items: dict[str, RetryEntry] = {}

    @staticmethod
    def _key(email: str) -> str:
        return email.strip().lower()

    async def add(self, entry: RetryEntry) -> None:
        self._items[self._key(entry.email)] = entry

    async def due(self, now: float, limit: int) -> list[RetryEntry]:
        ready = [e for e in self._items.values() if e.due_at <= now]
        ready.sort(key=lambda e: e.due_at)
        return ready[:limit]

    async def remove(self, email: str) -> None:
        self._items.pop(self._key(email), None)

    async def get(self, email: str) -> Optional[RetryEntry]:
        return self._items.get(self._key(email))

    async def count(self) -> int:
        return len(self._items)


class RedisRetryStore:
    """Durable, worker-shared store: a hash of entries plus a sorted-set index.

    The sorted set is scored by `due_at`, so "what is due now" is one
    ZRANGEBYSCORE rather than a full scan of every pending address.
    """

    HASH_KEY = "vx:retry:entries"
    DUE_KEY = "vx:retry:due"

    def __init__(self, client=None):
        self._client = client

    async def _redis(self):
        if self._client is not None:
            return self._client
        from ..utils.cache import get_redis

        self._client = await get_redis()
        return self._client

    @staticmethod
    def _key(email: str) -> str:
        return email.strip().lower()

    async def add(self, entry: RetryEntry) -> None:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        key = self._key(entry.email)
        await client.hset(self.HASH_KEY, key, json.dumps(entry.as_dict()))
        await client.zadd(self.DUE_KEY, {key: entry.due_at})

    async def due(self, now: float, limit: int) -> list[RetryEntry]:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        members = await client.zrangebyscore(self.DUE_KEY, float("-inf"), now, start=0, num=limit)
        entries: list[RetryEntry] = []
        for member in members:
            raw = await client.hget(self.HASH_KEY, member)
            if not raw:
                # Index without a payload (partial write): drop the orphan.
                await client.zrem(self.DUE_KEY, member)
                continue
            try:
                entries.append(RetryEntry.from_dict(json.loads(raw)))
            except Exception:
                # Corrupt payload must be skipped, never crash the drain.
                await client.zrem(self.DUE_KEY, member)
                await client.hdel(self.HASH_KEY, member)
        return entries

    async def remove(self, email: str) -> None:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        key = self._key(email)
        await client.hdel(self.HASH_KEY, key)
        await client.zrem(self.DUE_KEY, key)

    async def get(self, email: str) -> Optional[RetryEntry]:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        raw = await client.hget(self.HASH_KEY, self._key(email))
        if not raw:
            return None
        try:
            return RetryEntry.from_dict(json.loads(raw))
        except Exception:
            return None

    async def count(self) -> int:
        client = await self._redis()
        if client is None:
            raise RuntimeError("redis unavailable")
        return int(await client.hlen(self.HASH_KEY) or 0)


def build_retry_store():
    """Redis when configured/reachable, else the in-process store."""
    try:
        from ..utils.cache import _HAS_REDIS  # type: ignore[attr-defined]

        if _HAS_REDIS:
            return RedisRetryStore()
    except Exception:
        pass
    return MemoryRetryStore()


class RetryQueue:
    """Schedules greylisted addresses for a later attempt."""

    def __init__(
        self,
        store: Optional[RetryStore] = None,
        backoff_seconds: int = 900,      # 15 min — inside the usual greylist window
        max_attempts: int = 3,
        now: Optional[Callable[[], float]] = None,
    ):
        self.store = store or MemoryRetryStore()
        # 0 is legitimate (immediate retry — tests and hot-path recovery);
        # the bounded max_attempts prevents any runaway.
        self.backoff_seconds = max(0, backoff_seconds)
        self.max_attempts = max(1, max_attempts)
        self._now = now or time.time

    def _delay_for(self, attempts: int) -> float:
        # attempts=0 -> base, 1 -> 2x base, 2 -> 4x base (capped growth).
        return self.backoff_seconds * (2 ** min(attempts, 4))

    async def enqueue(self, email: str, mx_host: str, code: Optional[int] = None) -> bool:
        """Schedule `email` for a later attempt. Returns False if not queued."""
        try:
            existing = None
            if hasattr(self.store, "get"):
                existing = await self.store.get(email)  # type: ignore[attr-defined]
            if existing is not None:
                return False  # already queued; never create duplicates
            entry = RetryEntry(
                email=email,
                mx_host=mx_host,
                attempts=0,
                due_at=self._now() + self._delay_for(0),
                last_code=code,
            )
            await self.store.add(entry)
            return True
        except Exception:
            return False  # a queue outage must not break validation

    async def due(self, limit: int = 50) -> list[RetryEntry]:
        try:
            return await self.store.due(self._now(), limit)
        except Exception:
            return []

    async def mark_attempted(self, entry: RetryEntry) -> None:
        """Record an attempt; reschedule with backoff or drop when exhausted."""
        try:
            attempts = entry.attempts + 1
            if attempts >= self.max_attempts:
                await self.store.remove(entry.email)
                return
            entry.attempts = attempts
            entry.due_at = self._now() + self._delay_for(attempts)
            await self.store.add(entry)
        except Exception:
            return

    async def resolve(self, entry: RetryEntry) -> None:
        """The address got a definitive answer — stop retrying it."""
        try:
            await self.store.remove(entry.email)
        except Exception:
            return

    async def resolve_email(self, email: str) -> None:
        """Convenience: stop retrying by address (definitive verdict arrived)."""
        try:
            await self.store.remove(email)
        except Exception:
            return

    async def size(self) -> int:
        try:
            if hasattr(self.store, "count"):
                return await self.store.count()  # type: ignore[attr-defined]
            return len(await self.store.due(float("inf"), 10_000))
        except Exception:
            return 0


async def drain_retry_queue(
    prober: Callable[[str, str], Awaitable[SMTPResult]],
    queue: Optional[RetryQueue] = None,
    limit: int = 50,
) -> dict:
    """Re-probe addresses that were greylisted earlier.

    A queue that only fills is a leak, so this is the other half of the feature.
    For each due entry:

      * definitive answer (250 or 5xx) -> drop it from the queue
      * still 4xx                       -> re-queue with exponential backoff
      * probe raised                    -> leave queued for a later attempt

    Returns a stats dict including the raw results so the caller can persist them.
    """
    q = queue or retry_queue
    stats: dict = {"checked": 0, "resolved": 0, "requeued": 0, "errors": 0, "results": []}

    for entry in await q.due(limit=limit):
        stats["checked"] += 1
        try:
            result = await prober(entry.email, entry.mx_host)
        except Exception:
            # Transport-level blow-up: keep it queued, try again next drain.
            stats["errors"] += 1
            await q.mark_attempted(entry)
            continue

        stats["results"].append((entry.email, result))

        if is_greylist_result(result):
            stats["requeued"] += 1
            await q.mark_attempted(entry)
        else:
            # 250 or 5xx — a real answer, so stop retrying this address.
            stats["resolved"] += 1
            await q.resolve(entry)

    return stats


# Shared queue used by the bulk pipeline. Prefers Redis so the queue survives a
# restart and is shared across workers; falls back to memory otherwise.
retry_queue = RetryQueue(store=build_retry_store())