"""Сервис «Моей волны»: собирает пачку из каталога, профиля, сессии и контекста."""

from __future__ import annotations

from typing import Literal

import numpy as np
import psycopg
from pydantic import BaseModel, Field

from pipu_rec.wave.contexts import Contexts, Rule
from pipu_rec.wave.engine import MODEL_VER, Pick, build_batch, taste_centroids, update_session
from pipu_rec.wave.store import IndexHolder, SessionStore, load_profile, seed_tracks

MAX_SERVED_MEMORY = 300

Outcome = Literal["complete", "skip", "skip_early", "like", "dislike"]


class Feedback(BaseModel):
    track_id: str
    outcome: Outcome


class Seed(BaseModel):
    type: Literal["track", "album", "artist"]
    id: str


class BatchRequest(BaseModel):
    user_id: str
    session_id: str
    activity: str | None = None
    mood: str | None = None
    seed: Seed | None = None
    exclude_ids: list[str] = Field(default_factory=list, max_length=1000)
    feedback: list[Feedback] = Field(default_factory=list, max_length=200)
    n: int = Field(default=5, ge=1, le=20)


class BatchTrack(BaseModel):
    id: str
    source: str
    score: float


class BatchResponse(BaseModel):
    tracks: list[BatchTrack]
    model_ver: str


class WaveService:
    def __init__(self, dsn: str, redis_url: str, contexts: Contexts) -> None:
        self.dsn = dsn
        self.index = IndexHolder(dsn)
        self.sessions = SessionStore(redis_url)
        self.contexts = contexts

    def rules(self, activity: str | None, mood: str | None) -> list[Rule]:
        """Занятие и настроение комбинируются — пересечение профилей; неизвестный ключ игнорируется."""
        rules: list[Rule] = []
        if activity and activity in self.contexts.activities:
            rules += self.contexts.activities[activity].rules
        if mood and mood in self.contexts.moods:
            rules += self.contexts.moods[mood].rules
        return rules

    def batch(self, req: BatchRequest) -> BatchResponse:
        index = self.index.get()
        with psycopg.connect(self.dsn) as conn:
            profile = load_profile(conn, req.user_id)
            seeds = seed_tracks(conn, req.seed.type, req.seed.id) if req.seed else None
        session = self.sessions.get(req.session_id)
        update_session(
            index, session, [(f.track_id, f.outcome) for f in req.feedback], taste_centroids(index, profile)
        )
        picks: list[Pick] = build_batch(
            index,
            profile,
            session,
            self.rules(req.activity, req.mood),
            n=req.n,
            exclude=set(req.exclude_ids),
            seed_ids=seeds,
            rng=np.random.default_rng(),
        )
        session.served = (session.served + [p.id for p in picks])[-MAX_SERVED_MEMORY:]
        artists = [index.artist[index.pos[p.id]] for p in picks]
        session.recent_artists = (session.recent_artists + artists)[-10:]
        self.sessions.put(req.session_id, session)
        return BatchResponse(
            tracks=[BatchTrack(id=p.id, source=p.source, score=p.score) for p in picks], model_ver=MODEL_VER
        )