pipu-music/services/api/internal/wave/wave.go

// Package wave — «Моя волна»: сессии, вызов рекомендателя и запасная волна на Go (docs/recommendations.md).
package wave

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"pipu/api/internal/api/gen"
	"pipu/api/internal/auth"
	"pipu/api/internal/library"
	"pipu/api/internal/platform/apperr"
)

// BatchSize — треков в пачке волны.
const BatchSize = 5

// RecommenderTimeout — сколько ждать рекомендатель, прежде чем отдать запасную волну.
const RecommenderTimeout = 800 * time.Millisecond

// Занятия и настроения — ключи совпадают с contexts.yaml рекомендателя, тексты — в i18n клиента.
var (
	activities = []gen.WaveOption{
		{Key: "running", Icon: "directions_run"}, {Key: "workout", Icon: "fitness_center"},
		{Key: "waking_up", Icon: "wb_sunny"}, {Key: "focus", Icon: "headphones"},
		{Key: "commute", Icon: "directions_car"}, {Key: "chill", Icon: "weekend"}, {Key: "sleep", Icon: "bedtime"},
	}
	moods = []gen.WaveOption{
		{Key: "energetic", Icon: "bolt"}, {Key: "happy", Icon: "sentiment_satisfied"},
		{Key: "calm", Icon: "spa"}, {Key: "sad", Icon: "rainy"},
	}
)

func known(list []gen.WaveOption, key *string) bool {
	if key == nil || *key == "" {
		return true
	}
	for _, o := range list {
		if o.Key == *key {
			return true
		}
	}
	return false
}

// Pick — трек, выбранный рекомендателем.
type Pick struct {
	ID     string  `json:"id"`
	Source string  `json:"source"`
	Score  float64 `json:"score"`
}

// Feedback — реакция на выданный трек.
type Feedback struct {
	TrackID string `json:"track_id"`
	Outcome string `json:"outcome"`
}

// BatchRequest — запрос к рекомендателю.
type BatchRequest struct {
	UserID     string     `json:"user_id"`
	SessionID  string     `json:"session_id"`
	Activity   *string    `json:"activity"`
	Mood       *string    `json:"mood"`
	Seed       *Seed      `json:"seed,omitempty"`
	ExcludeIDs []string   `json:"exclude_ids"`
	Feedback   []Feedback `json:"feedback"`
	N          int        `json:"n"`
}

// Seed — от чего запущена волна.
type Seed struct {
	Type string `json:"type"`
	ID   string `json:"id"`
}

// Recommender — клиент внутреннего API рекомендаций.
type Recommender struct {
	URL    string
	Token  string
	Client *http.Client
}

// Batch запрашивает пачку.
func (r *Recommender) Batch(ctx context.Context, req BatchRequest) ([]Pick, string, error) {
	raw, err := json.Marshal(req)
	if err != nil {
		return nil, "", err
	}
	ctx, cancel := context.WithTimeout(ctx, RecommenderTimeout)
	defer cancel()
	hr, err := http.NewRequestWithContext(ctx, http.MethodPost, r.URL+"/internal/wave/batch", bytes.NewReader(raw))
	if err != nil {
		return nil, "", err
	}
	hr.Header.Set("Authorization", "Bearer "+r.Token)
	hr.Header.Set("Content-Type", "application/json")
	res, err := r.Client.Do(hr)
	if err != nil {
		return nil, "", fmt.Errorf("recommender: %w", err)
	}
	defer func() { _ = res.Body.Close() }()
	if res.StatusCode != http.StatusOK {
		return nil, "", fmt.Errorf("recommender: status %d", res.StatusCode)
	}
	var out struct {
		Tracks   []Pick `json:"tracks"`
		ModelVer string `json:"model_ver"`
	}
	if err := json.NewDecoder(res.Body).Decode(&out); err != nil {
		return nil, "", fmt.Errorf("recommender: %w", err)
	}
	return out.Tracks, out.ModelVer, nil
}

// session — строка wave_sessions.
type session struct {
	ID       uuid.UUID
	UserID   uuid.UUID
	Activity *string
	Mood     *string
	SeedType *string
	SeedID   *uuid.UUID
	Served   int
	Ended    bool
}

// Handler реализует wave-часть gen.StrictServerInterface.
type Handler struct {
	pool    *pgxpool.Pool
	rec     *Recommender
	catalog *library.Store
	dto     func() library.DTO
	log     *slog.Logger
	// Warm заранее готовит транскоды выданной пачки (docs/playback.md).
	Warm func(ctx context.Context, trackIDs []uuid.UUID)
}

// NewHandler создаёт обработчик.
func NewHandler(pool *pgxpool.Pool, rec *Recommender, catalog *library.Store, dto func() library.DTO, log *slog.Logger) *Handler {
	return &Handler{pool: pool, rec: rec, catalog: catalog, dto: dto, log: log}
}

func principal(ctx context.Context) (auth.Principal, error) {
	p, ok := auth.PrincipalFrom(ctx)
	if !ok {
		return auth.Principal{}, apperr.Unauthorized(auth.CodeUnauthorized, "not authenticated")
	}
	return p, nil
}

// GetWaveOptions реализует gen.StrictServerInterface.
func (h *Handler) GetWaveOptions(context.Context, gen.GetWaveOptionsRequestObject) (gen.GetWaveOptionsResponseObject, error) {
	return gen.GetWaveOptions200JSONResponse{Activities: activities, Moods: moods}, nil
}

func validateContext(activity, mood *string) error {
	if !known(activities, activity) {
		return apperr.Validation("activity", "validation.wave_activity_invalid", "unknown activity", nil)
	}
	if !known(moods, mood) {
		return apperr.Validation("mood", "validation.wave_mood_invalid", "unknown mood", nil)
	}
	return nil
}

func emptyToNil(s *string) *string {
	if s == nil || *s == "" {
		return nil
	}
	return s
}

// StartWave реализует gen.StrictServerInterface.
func (h *Handler) StartWave(ctx context.Context, req gen.StartWaveRequestObject) (gen.StartWaveResponseObject, error) {
	p, err := principal(ctx)
	if err != nil {
		return nil, err
	}
	s := session{ID: uuid.Must(uuid.NewV7()), UserID: p.UserID}
	if b := req.Body; b != nil {
		s.Activity, s.Mood = emptyToNil(b.Activity), emptyToNil(b.Mood)
		if b.Seed != nil {
			t, id := string(b.Seed.Type), b.Seed.Id
			s.SeedType, s.SeedID = &t, &id
		}
	}
	if err := validateContext(s.Activity, s.Mood); err != nil {
		return nil, err
	}
	if _, err := h.pool.Exec(ctx, `
		INSERT INTO wave_sessions (id, user_id, activity, mood, seed_type, seed_id) VALUES ($1, $2, $3, $4, $5, $6)`,
		s.ID, s.UserID, s.Activity, s.Mood, s.SeedType, s.SeedID); err != nil {
		return nil, fmt.Errorf("wave: start: %w", err)
	}
	out, err := h.serve(ctx, s, nil)
	if err != nil {
		return nil, err
	}
	return gen.StartWave201JSONResponse(out), nil
}

func (h *Handler) load(ctx context.Context, p auth.Principal, id uuid.UUID) (session, error) {
	var s session
	err := h.pool.QueryRow(ctx, `
		SELECT id, user_id, activity, mood, seed_type, seed_id, tracks_served, ended_at IS NOT NULL
		FROM wave_sessions WHERE id = $1 AND user_id = $2`, id, p.UserID).
		Scan(&s.ID, &s.UserID, &s.Activity, &s.Mood, &s.SeedType, &s.SeedID, &s.Served, &s.Ended)
	if errors.Is(err, pgx.ErrNoRows) {
		return s, apperr.NotFound("wave.session_not_found", "wave session not found")
	}
	if err != nil {
		return s, fmt.Errorf("wave: load: %w", err)
	}
	return s, nil
}

// NextWave реализует gen.StrictServerInterface.
func (h *Handler) NextWave(ctx context.Context, req gen.NextWaveRequestObject) (gen.NextWaveResponseObject, error) {
	p, err := principal(ctx)
	if err != nil {
		return nil, err
	}
	s, err := h.load(ctx, p, req.Id)
	if err != nil {
		return nil, err
	}
	var fb []Feedback
	if req.Body != nil && req.Body.Feedback != nil {
		for _, f := range *req.Body.Feedback {
			fb = append(fb, Feedback{TrackID: f.TrackId.String(), Outcome: string(f.Outcome)})
		}
	}
	out, err := h.serve(ctx, s, fb)
	if err != nil {
		return nil, err
	}
	return gen.NextWave200JSONResponse(out), nil
}

// UpdateWave реализует gen.StrictServerInterface.
func (h *Handler) UpdateWave(ctx context.Context, req gen.UpdateWaveRequestObject) (gen.UpdateWaveResponseObject, error) {
	p, err := principal(ctx)
	if err != nil {
		return nil, err
	}
	s, err := h.load(ctx, p, req.Id)
	if err != nil {
		return nil, err
	}
	s.Activity, s.Mood = emptyToNil(req.Body.Activity), emptyToNil(req.Body.Mood)
	if err := validateContext(s.Activity, s.Mood); err != nil {
		return nil, err
	}
	if _, err := h.pool.Exec(ctx, `UPDATE wave_sessions SET activity = $2, mood = $3 WHERE id = $1`,
		s.ID, s.Activity, s.Mood); err != nil {
		return nil, fmt.Errorf("wave: update: %w", err)
	}
	out, err := h.serve(ctx, s, nil)
	if err != nil {
		return nil, err
	}
	return gen.UpdateWave200JSONResponse(out), nil
}

// EndWave реализует gen.StrictServerInterface.
func (h *Handler) EndWave(ctx context.Context, req gen.EndWaveRequestObject) (gen.EndWaveResponseObject, error) {
	p, err := principal(ctx)
	if err != nil {
		return nil, err
	}
	if _, err := h.pool.Exec(ctx, `UPDATE wave_sessions SET ended_at = now() WHERE id = $1 AND user_id = $2 AND ended_at IS NULL`,
		req.Id, p.UserID); err != nil {
		return nil, fmt.Errorf("wave: end: %w", err)
	}
	return gen.EndWave204Response{}, nil
}

// serve собирает пачку (рекомендатель или запасная волна), пишет показы и отдаёт треки.
func (h *Handler) serve(ctx context.Context, s session, fb []Feedback) (gen.WaveBatch, error) {
	served, err := h.servedIDs(ctx, s.ID)
	if err != nil {
		return gen.WaveBatch{}, err
	}
	req := BatchRequest{
		UserID: s.UserID.String(), SessionID: s.ID.String(), Activity: s.Activity, Mood: s.Mood,
		ExcludeIDs: served, Feedback: fb, N: BatchSize,
	}
	if s.SeedType != nil && s.SeedID != nil {
		req.Seed = &Seed{Type: *s.SeedType, ID: s.SeedID.String()}
	}
	if req.Feedback == nil {
		req.Feedback = []Feedback{}
	}
	picks, modelVer, err := h.rec.Batch(ctx, req)
	fallback := false
	if err != nil || len(picks) == 0 {
		if err != nil {
			h.log.WarnContext(ctx, "recommender unavailable, serving fallback wave", "err", err)
		}
		fallback, modelVer = true, "fallback"
		if picks, err = h.fallback(ctx, s, served); err != nil {
			return gen.WaveBatch{}, err
		}
	}

	ids := make([]uuid.UUID, 0, len(picks))
	for _, p := range picks {
		if id, err := uuid.Parse(p.ID); err == nil {
			ids = append(ids, id)
		}
	}
	rows, err := h.catalog.TracksByIDs(ctx, ids)
	if err != nil {
		return gen.WaveBatch{}, err
	}
	tracks, err := h.dto().Tracks(ctx, rows)
	if err != nil {
		return gen.WaveBatch{}, err
	}
	if err := h.record(ctx, s, picks, modelVer); err != nil {
		return gen.WaveBatch{}, err
	}
	if h.Warm != nil && len(ids) > 0 {
		go h.Warm(context.WithoutCancel(ctx), ids)
	}
	return gen.WaveBatch{SessionId: s.ID, Activity: s.Activity, Mood: s.Mood, Tracks: tracks, Fallback: fallback}, nil
}

func (h *Handler) servedIDs(ctx context.Context, sessionID uuid.UUID) ([]string, error) {
	rows, err := h.pool.Query(ctx, `SELECT track_id::text FROM wave_impressions WHERE wave_session_id = $1`, sessionID)
	if err != nil {
		return nil, fmt.Errorf("wave: served: %w", err)
	}
	ids, err := pgx.CollectRows(rows, pgx.RowTo[string])
	if ids == nil {
		ids = []string{}
	}
	return ids, err
}

// record пишет показы (для метрик и будущего обучения ранжировщика).
func (h *Handler) record(ctx context.Context, s session, picks []Pick, modelVer string) error {
	return pgx.BeginFunc(ctx, h.pool, func(tx pgx.Tx) error {
		var pos int
		if err := tx.QueryRow(ctx, `
			UPDATE wave_sessions SET tracks_served = tracks_served + $2 WHERE id = $1 RETURNING tracks_served - $2`,
			s.ID, len(picks)).Scan(&pos); err != nil {
			return fmt.Errorf("wave: impressions: %w", err)
		}
		for i, p := range picks {
			if _, err := tx.Exec(ctx, `
				INSERT INTO wave_impressions (wave_session_id, position, track_id, candidate_source, score, model_ver)
				VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT DO NOTHING`,
				s.ID, pos+i, p.ID, p.Source, p.Score, modelVer); err != nil {
				return fmt.Errorf("wave: impressions: %w", err)
			}
		}
		return nil
	})
}

// fallback — запасная волна без рекомендателя: лайки, любимые исполнители и случайные доступные треки,
// без дизлайков, недавно игравших и уже выданных; один трек исполнителя на пачку.
func (h *Handler) fallback(ctx context.Context, s session, served []string) ([]Pick, error) {
	rows, err := h.pool.Query(ctx, `
		WITH liked AS (
			SELECT track_id FROM track_reactions WHERE user_id = $1 AND reaction = 'like'
		), fav_artists AS (
			SELECT DISTINCT ta.artist_id FROM track_artists ta JOIN liked l ON l.track_id = ta.track_id
		), pool AS (
			SELECT t.id,
				COALESCE((SELECT ta.artist_id FROM track_artists ta WHERE ta.track_id = t.id ORDER BY ta.position LIMIT 1),
					al.album_artist_id) AS artist_id,
				CASE WHEN t.id IN (SELECT track_id FROM liked) THEN 'liked'
					WHEN EXISTS (SELECT 1 FROM track_artists ta WHERE ta.track_id = t.id
						AND ta.artist_id IN (SELECT artist_id FROM fav_artists)) THEN 'artist'
					ELSE 'random' END AS source
			FROM tracks t
			JOIN albums al ON al.id = t.album_id
			JOIN media_files f ON f.id = t.primary_file_id AND f.missing_since IS NULL
			JOIN library_sources src ON src.id = f.source_id AND src.status = 'mounted'
			WHERE NOT EXISTS (SELECT 1 FROM track_reactions r WHERE r.user_id = $1 AND r.track_id = t.id AND r.reaction = 'dislike')
				AND NOT (t.id::text = ANY($2))
				AND NOT EXISTS (SELECT 1 FROM play_events e WHERE e.user_id = $1 AND e.track_id = t.id
					AND e.type = 'play_start' AND e.occurred_at > now() - interval '3 hours')
		)
		SELECT DISTINCT ON (artist_id) id::text, source FROM (
			SELECT * FROM pool ORDER BY CASE source WHEN 'liked' THEN 0 WHEN 'artist' THEN 1 ELSE 2 END, random() LIMIT 200
		) c
		ORDER BY artist_id, random()
		LIMIT 50`, s.UserID, served)
	if err != nil {
		return nil, fmt.Errorf("wave: fallback: %w", err)
	}
	picks, err := pgx.CollectRows(rows, func(r pgx.CollectableRow) (Pick, error) {
		var p Pick
		return p, r.Scan(&p.ID, &p.Source)
	})
	if err != nil {
		return nil, fmt.Errorf("wave: fallback: %w", err)
	}
	// DISTINCT ON упорядочил по исполнителю — перемешиваем, предпочитая лайки и любимых исполнителей.
	rank := map[string]int{"liked": 0, "artist": 1, "random": 2}
	shuffled := shuffle(picks)
	stableSort(shuffled, func(a, b Pick) bool { return rank[a.Source] < rank[b.Source] })
	// Половина пачки — любимое, половина — случайное: запасная волна не должна быть одними лайками.
	var out []Pick
	var fav, rnd []Pick
	for _, p := range shuffled {
		if p.Source == "random" {
			rnd = append(rnd, p)
		} else {
			fav = append(fav, p)
		}
	}
	for len(out) < BatchSize && (len(fav) > 0 || len(rnd) > 0) {
		if len(fav) > 0 && (len(out)%2 == 0 || len(rnd) == 0) {
			out, fav = append(out, fav[0]), fav[1:]
		} else if len(rnd) > 0 {
			out, rnd = append(out, rnd[0]), rnd[1:]
		}
	}
	return out, nil
}