"""A weighted draw that stays a surprise: what "Surprise me" deals from.

A title's chance to come first is its weight over the pool's (Efraimidis–Spirakis: order by
-ln(1-u)/w, u uniform per row). The weight is 1 + the title's affinity to what the user watches
(`plex.affinity`, `discover.queries`), so a title nobody leans towards still comes up — only less
often. Two things keep it a surprise rather than a ranking:

- **u is the seed's md5 of the row's key**, the feed's shuffle idiom. With every weight 1 the order
  is exactly the md5 order the draw had before it was weighted, so a user with no history sees the
  old draw, and one seed still moves both sources of the Home hero together. A draw that is not
  seeded (the Plex Discover rail) takes u from random().
- **One slot of three is blind** (`blend`): dealt from the unweighted order, so the same few
  favourites cannot fill the rail.
"""

from collections.abc import Callable, Hashable
from typing import Any

from sqlalchemy import BigInteger, Float, cast, func, literal
from sqlalchemy.dialects.postgresql import BIT
from sqlalchemy.orm import InstrumentedAttribute
from sqlalchemy.sql.elements import ColumnElement

# 2**32: the first 8 hex digits of an md5 as a fraction of one.
_SPAN = 4294967296.0


def seeded_u(
    key: ColumnElement[Any] | InstrumentedAttribute[Any], seed: int
) -> ColumnElement[float]:
    """A uniform u in [0, 1) per row, from md5(key || seed) — the order `md5(...)` ascending gives."""
    digest = func.md5(func.concat(key, str(seed)))
    head = cast(cast(literal("x") + func.substr(digest, 1, 8), BIT(32)), BigInteger)
    return cast(head, Float) / _SPAN


def weighted_key(u: ColumnElement[float], weight: ColumnElement[float]) -> ColumnElement[float]:
    """The row's place in a weighted draw, ascending: -ln(1-u)/w. Monotonic in u when every w is 1,
    so the unweighted order is unchanged."""
    return -func.ln(1.0 - u) / weight


def random_u() -> ColumnElement[float]:
    return func.random()


def blend(
    weighted: list, blind: list, offset: int, n: int, key: Callable[[Any], Hashable] = lambda x: x
) -> list:
    """`n` picks: slot i comes from `blind` when (i + offset) % 3 == 0, else from `weighted`, each
    skipping what an earlier slot took (`key` tells two picks apart). A side that runs out leaves the
    rest to the other."""
    out: list = []
    taken: set = set()
    sources = (iter(weighted), iter(blind))

    def nxt(source):
        for item in source:
            if key(item) not in taken:
                return item
        return None

    for i in range(n):
        first, second = (sources[1], sources[0]) if (i + offset) % 3 == 0 else sources
        item = nxt(first)
        if item is None:
            item = nxt(second)
        if item is None:
            break
        taken.add(key(item))
        out.append(item)
    return out