"""Бандит по источникам кандидатов (docs/recommendations.md#6.6): Beta(α, β) на пару (пользователь, источник).

Награда — трек из источника доигран или лайкнут; скип и дизлайк — неудача. Наблюдения затухают ×0.99 в день,
чтобы бандит подстраивался под смену вкуса. Доля «нового» (exploration) следует за успехом источника `fresh`.
"""

from __future__ import annotations

import json
from collections import defaultdict
from typing import Any

import psycopg

DAILY_DECAY = 0.99
REWARD = {"complete": 1.0, "like": 1.0, "skip": 0.0, "skip_early": 0.0, "dislike": 0.0}
EXPLORATION_MIN, EXPLORATION_MAX = 0.05, 0.4

Params = dict[str, list[float]]  # источник → [α, β]


def compute(rows: list[tuple[str, str, str, float]], now: float) -> dict[str, Params]:
    """rows: (user, source, outcome, served_at) → user → source → [α, β] с приором Beta(1, 1)."""
    out: dict[str, Params] = defaultdict(dict)
    for user, source, outcome, ts in rows:
        r = REWARD.get(outcome)
        if r is None:
            continue
        w = DAILY_DECAY ** max((now - ts) / 86400.0, 0.0)
        ab = out[user].setdefault(source, [1.0, 1.0])
        ab[0] += w * r
        ab[1] += w * (1.0 - r)
    return {u: {s: [round(a, 4), round(b, 4)] for s, (a, b) in p.items()} for u, p in out.items()}


def exploration(params: Params, default: float = 0.15) -> float:
    """Доля «нового»: успешнее `fresh` относительно остальных источников — больше исследования."""
    fresh = params.get("fresh")
    others = [a / (a + b) for s, (a, b) in params.items() if s != "fresh"]
    if fresh is None or not others:
        return default
    rate = fresh[0] / (fresh[0] + fresh[1])
    base = sum(others) / len(others)
    value = default * (rate / base if base > 0 else 1.0)
    return round(min(max(value, EXPLORATION_MIN), EXPLORATION_MAX), 3)


def update_bandits(conn: psycopg.Connection[Any], now: float) -> int:
    """Пересчитывает параметры по показам волны за 90 дней. Возвращает число пользователей."""
    rows = conn.execute(
        """
        SELECT ws.user_id::text, wi.candidate_source, wi.outcome, extract(epoch FROM wi.served_at)::float8
        FROM wave_impressions wi JOIN wave_sessions ws ON ws.id = wi.wave_session_id
        WHERE wi.outcome IS NOT NULL AND wi.served_at > now() - interval '90 days'
        """
    ).fetchall()
    params = compute([(u, s, o, float(ts)) for u, s, o, ts in rows], now)
    with conn.transaction(), conn.cursor() as cur:
        cur.executemany(
            """
            INSERT INTO user_profiles_rec (user_id, bandit, exploration, updated_at)
            VALUES (%s, %s, %s, now())
            ON CONFLICT (user_id) DO UPDATE SET bandit = EXCLUDED.bandit, exploration = EXCLUDED.exploration,
                updated_at = now()
            """,
            [(u, json.dumps(p), exploration(p)) for u, p in params.items()],
        )
    return len(params)