"""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.

The YouTube surprise draws over a window of the feed scope rather than all of it (`yt_surprise`):
the window is read in SQL, and `shuffled` / `leaning` order it here with the same u.
"""

import hashlib
import math
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 u_of(key: str, seed: int) -> float:
    """`seeded_u` in Python: the same u for the same key and seed."""
    return int(hashlib.md5(f"{key}{seed}".encode()).hexdigest()[:8], 16) / _SPAN


def shuffled(items: list, seed: int, key: Callable[[Any], str]) -> list:
    """`items` in the seed's md5 order — the order `ORDER BY md5(key || seed)` gives."""
    return sorted(items, key=lambda it: hashlib.md5(f"{key(it)}{seed}".encode()).hexdigest())


def leaning(
    items: list, seed: int, key: Callable[[Any], str], weight: Callable[[Any], float]
) -> list:
    """`items` in the weighted draw's order (`weighted_key` in Python), ties in the md5 order."""
    order = {key(it): i for i, it in enumerate(shuffled(items, seed, key))}
    return sorted(
        items,
        key=lambda it: (-math.log(1.0 - u_of(key(it), seed)) / weight(it), order[key(it)]),
    )


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