PixivFE/core/user.go

package core

import (
	"fmt"
	"math"
	"net/http"
	"sort"
	"strconv"
	"strings"

	"codeberg.org/vnpower/pixivfe/v2/i18n"
	"codeberg.org/vnpower/pixivfe/v2/server/session"
	"github.com/goccy/go-json"
)

// pixivfe internal data type. not used by pixiv.
type UserWorkCategory string

const (
	CategoryAny          UserWorkCategory = ""
	CategoryAnyAlt       UserWorkCategory = "artworks"
	CategoryIllustration UserWorkCategory = "illustrations"
	CategoryManga        UserWorkCategory = "manga"
	CategoryBookmarks    UserWorkCategory = "bookmarks" // what this user has bookmarked; not art by this user
	CategoryNovels       UserWorkCategory = "novels"
)

// UserWorkIDs holds separate ID strings for illustrations, manga, and novels.
type UserWorkIDs struct {
	IllustrationIDs string
	MangaIDs        string
	NovelIDs        string
}

func (s UserWorkCategory) Validate() error {
	if s != CategoryAny &&
		s != CategoryAnyAlt &&
		s != CategoryIllustration &&
		s != CategoryManga &&
		s != CategoryBookmarks &&
		s != CategoryNovels {
		return i18n.Errorf(`Invalid work category: %#v. Only "%s", "%s", "%s", "%s", "%s" and "%s" are available`, s, CategoryAny, CategoryAnyAlt, CategoryIllustration, CategoryManga, CategoryBookmarks, CategoryNovels)
	} else {
		return nil
	}
}

type FrequentTag struct {
	Name           string `json:"tag"`
	TranslatedName string `json:"tag_translation"`
}

type CountInfo struct {
	All           int
	Illustrations int
	Manga         int
	Novels        int
	Bookmarks     int
}

type User struct {
	ID                string          `json:"userId"`
	Name              string          `json:"name"`
	Avatar            string          `json:"imageBig"`
	Following         int             `json:"following"`
	MyPixiv           int             `json:"mypixivCount"`
	Comment           HTML            `json:"commentHtml"`
	Webpage           string          `json:"webpage"`
	SocialRaw         json.RawMessage `json:"social"`
	Artworks          []ArtworkBrief  `json:"artworks"` // this slice includes both illustrations and manga, but not novels
	Illustrations     []ArtworkBrief
	Manga             []ArtworkBrief
	Novels            []NovelBrief   `json:"novels"`
	Background        map[string]any `json:"background"`
	CategoryItemCount int
	FrequentTags      []FrequentTag
	Social            map[string]map[string]string
	BackgroundImage   string
	NovelSeries       []NovelSeries
	MangaSeries       []MangaSeries

	IsFollowed bool `json:"isFollowed"` // Denotes whether the logged in user currently following the given user

	// The following fields are internal to PixivFE, used to display the number of works for a given category
	CountInfo CountInfo
}

// GetUserProfile retrieves the user profile, including counts, artworks/bookmarks, and social data.
func GetUserProfile(r *http.Request, id string, category UserWorkCategory, page int, getTags bool) (User, error) {
	var user User

	// Fetch user information
	userInfoURL := GetUserInformationURL(id)
	token := session.GetUserToken(r)
	resp, err := API_GET_UnwrapJson(r.Context(), userInfoURL, token, r.Header)
	if err != nil {
		return user, err
	}

	resp = session.ProxyImageUrl(r, resp)
	if err := json.Unmarshal([]byte(resp), &user); err != nil {
		return user, err
	}

	// Fetch counts
	if err := fetchUserCounts(r, &user, id); err != nil {
		return user, err
	}

	// Fetch artworks or bookmarks based on the category
	if err := fetchUserWorks(r, &user, id, category, page, getTags); err != nil {
		return user, err
	}

	// Parse social data
	if err := user.parseSocial(); err != nil {
		return User{}, err
	}

	// Set background image if available
	if user.Background != nil {
		user.BackgroundImage = user.Background["url"].(string)
	}

	return user, nil
}

// fetchUserCounts retrieves and sets the CountInfo field of the User.
func fetchUserCounts(r *http.Request, user *User, id string) error {
	// Get counts for illustrations and manga
	_, countInfo, _, err := getUserArtworksIDAndSeries(r, id, CategoryAny, 1)
	if err != nil {
		return err
	}
	user.CountInfo = countInfo

	// Get count for novels separately
	_, novelCountInfo, _, err := getUserArtworksIDAndSeries(r, id, CategoryNovels, 1)
	if err != nil {
		return err
	}
	user.CountInfo.Novels = novelCountInfo.Novels

	// Get bookmarks count
	_, bookmarksCount, err := getUserBookmarks(r, id, "show", 1)
	if err != nil {
		return err
	}
	user.CountInfo.Bookmarks = bookmarksCount

	return nil
}

// fetchUserWorks retrieves and sets artworks or bookmarks based on the category.
//
// It also handles frequent tags and series data.
func fetchUserWorks(r *http.Request, user *User, id string, category UserWorkCategory, page int, getTags bool) error {
	if category == CategoryBookmarks {
		// Fetch bookmarks
		works, count, err := getUserBookmarks(r, id, "show", page)
		if err != nil {
			return err
		}

		user.Artworks = works
		user.NovelSeries = nil
		user.MangaSeries = nil
		user.CategoryItemCount = count
		return nil
	}

	// Fetch artworks or novels
	userWorkIDs, countInfo, series, err := getUserArtworksIDAndSeries(r, id, category, page)
	if err != nil {
		return err
	}

	// Initialize ID strings for different work types
	var illustrationIDs, mangaIDs, novelIDs string

	switch category {
	case CategoryIllustration:
		illustrationIDs = userWorkIDs.IllustrationIDs
	case CategoryManga:
		mangaIDs = userWorkIDs.MangaIDs
	case CategoryNovels:
		novelIDs = userWorkIDs.NovelIDs
	case CategoryAny, CategoryAnyAlt:
		// Combine all relevant ID strings
		illustrationIDs = userWorkIDs.IllustrationIDs
		mangaIDs = userWorkIDs.MangaIDs
		novelIDs = userWorkIDs.NovelIDs
	default:
		return fmt.Errorf("unsupported category: %v", category)
	}

	// Early return if there are no works to fetch
	if countInfo.All == 0 {
		user.Artworks = nil
		user.CategoryItemCount = 0
		user.NovelSeries = nil
		user.MangaSeries = nil
		return nil
	}

	// Fetch and process artworks if they exist
	if illustrationIDs != "" || mangaIDs != "" {
		// Combine Illustration and Manga IDs
		ids := illustrationIDs + mangaIDs
		artworks, err := fetchAndProcessItems[ArtworkBrief](r, id, ids, getUserArtworkIDs)
		if err != nil {
			return err
		}

		user.Artworks = artworks
		user.Illustrations = filterArtworksByType(artworks, 0)
		user.Manga = filterArtworksByType(artworks, 1)
	}

	// Fetch and process novels if they exist
	if novelIDs != "" {
		ids := novelIDs
		novels, err := fetchAndProcessItems[NovelBrief](r, id, ids, getUserNovelIDs)
		if err != nil {
			return err
		}
		user.Novels = novels
	}

	// Fetch frequent tags if requested
	if getTags {
		var tagsIDs string

		// TODO: frequent tags for novels are *not* fetched when CategoryAny or CategoryAnyAlt
		if category == CategoryAny || category == CategoryAnyAlt || category == CategoryIllustration || category == CategoryManga {
			tagsIDs = illustrationIDs + mangaIDs
		} else {
			tagsIDs = novelIDs
		}

		user.FrequentTags, err = getUserFrequentTags(r, tagsIDs, category)
		if err != nil {
			return err
		}
	}

	// Handle series data
	user.NovelSeries, user.MangaSeries = handleSeriesData(series, category)
	user.CategoryItemCount = countInfo.All

	return nil
}

// The HasID interface allows for generic handling of the ArtworkBrief
// and NovelBrief structs, which both contain an ID string.
type HasID interface {
	GetID() string
}

// GetID implements the HasID interface for the ArtworkBrief struct.
func (a ArtworkBrief) GetID() string {
	return a.ID
}

// GetID implements the HasID interface for the NovelBrief struct.
func (n NovelBrief) GetID() string {
	return n.ID
}

// fetchAndProcessItems is a generic function that fetches and processes items of type T,
// where T must implement the HasID interface.
func fetchAndProcessItems[T HasID](
	r *http.Request,
	id, ids string,
	fetchFunc func(*http.Request, string, string) ([]T, error),
) ([]T, error) {
	items, err := fetchFunc(r, id, ids)
	if err != nil {
		return nil, err
	}

	sort.Slice(items, func(i, j int) bool {
		return numberGreaterThan(items[i].GetID(), items[j].GetID())
	})

	return items, nil
}

// filterArtworksByType filters artworks based on the IllustType.
func filterArtworksByType(artworks []ArtworkBrief, illustType int) []ArtworkBrief {
	filtered := make([]ArtworkBrief, 0)
	for _, artwork := range artworks {
		if artwork.IllustType == illustType {
			filtered = append(filtered, artwork)
		}
	}
	return filtered
}

// IntStringMap is a custom type that can handle both JSON objects and empty arrays.
//
// Required as the Pixiv API returns empty arrays for work types that
// don't exist for a user.
type IntStringMap map[int]string

// UnmarshalJSON implements custom unmarshalling for IntStringMap.
func (m *IntStringMap) UnmarshalJSON(data []byte) error {
	// Check if the data is an empty array.
	if string(data) == "[]" {
		*m = make(map[int]string)
		return nil
	}

	// Otherwise, expect an object.
	var temp map[string]*string
	if err := json.Unmarshal(data, &temp); err != nil {
		return err
	}

	result := make(map[int]string)
	for keyStr, value := range temp {
		key, err := strconv.Atoi(keyStr)
		if err != nil {
			return fmt.Errorf("invalid key '%s': %v", keyStr, err)
		}
		if value != nil {
			result[key] = *value
		} else {
			result[key] = ""
		}
	}

	*m = result
	return nil
}

// getUserArtworksIDAndSeries retrieves artwork IDs and series information for a user.
// It now returns separate ID strings for illustrations, manga, and novels encapsulated in UserWorkIDs.
func getUserArtworksIDAndSeries(r *http.Request, id string, category UserWorkCategory, page int) (userWorkIDs UserWorkIDs, countInfo CountInfo, series json.RawMessage, err error) {
	URL := GetUserArtworksURL(id)

	resp, err := API_GET_UnwrapJson(r.Context(), URL, "", r.Header)
	if err != nil {
		return UserWorkIDs{}, CountInfo{}, nil, err
	}

	resp = session.ProxyImageUrl(r, resp)

	var body struct {
		Illusts     *json.RawMessage `json:"illusts"`
		Manga       *json.RawMessage `json:"manga"`
		MangaSeries *json.RawMessage `json:"mangaSeries"`
		Novels      *json.RawMessage `json:"novels"`
		NovelSeries *json.RawMessage `json:"novelSeries"`
	}

	err = json.Unmarshal([]byte(resp), &body)
	if err != nil {
		return UserWorkIDs{}, CountInfo{}, nil, fmt.Errorf("failed to unmarshal response body: %w", err)
	}

	countInfo = CountInfo{}

	var illusts IntStringMap
	var manga IntStringMap
	var novels IntStringMap

	var illustIDs []int
	var mangaIDs []int
	var novelIDs []int

	// Process Illustrations
	if body.Illusts != nil && len(*body.Illusts) > 0 {
		if err = json.Unmarshal(*body.Illusts, &illusts); err != nil {
			fmt.Printf("Error unmarshalling illusts: %v\n", err)
			illusts = make(IntStringMap)
		}
		for k := range illusts {
			illustIDs = append(illustIDs, k)
		}
		countInfo.All += len(illusts)
		countInfo.Illustrations = len(illusts)
	}

	// Process Manga
	if body.Manga != nil && len(*body.Manga) > 0 {
		if err = json.Unmarshal(*body.Manga, &manga); err != nil {
			fmt.Printf("Error unmarshalling manga: %v\n", err)
			manga = make(IntStringMap)
		}
		for k := range manga {
			mangaIDs = append(mangaIDs, k)
		}
		countInfo.All += len(manga)
		countInfo.Manga = len(manga)
		if body.MangaSeries != nil {
			series = *body.MangaSeries
		}
	}

	// Process Novels
	if body.Novels != nil && len(*body.Novels) > 0 {
		if err = json.Unmarshal(*body.Novels, &novels); err != nil {
			fmt.Printf("Error unmarshalling novels: %v\n", err)
			novels = make(IntStringMap)
		}
		for k := range novels {
			novelIDs = append(novelIDs, k)
		}
		countInfo.All += len(novels)
		countInfo.Novels = len(novels)
		if body.NovelSeries != nil {
			series = *body.NovelSeries
		}
	}

	// Sort IDs in descending order
	sort.Sort(sort.Reverse(sort.IntSlice(illustIDs)))
	sort.Sort(sort.Reverse(sort.IntSlice(mangaIDs)))
	sort.Sort(sort.Reverse(sort.IntSlice(novelIDs)))

	worksPerPage := 30.0

	// Compute slice bounds for each category
	startIllust, endIllust, err := computeSliceBounds(page, worksPerPage, len(illustIDs))
	if err != nil {
		return UserWorkIDs{}, CountInfo{}, nil, err
	}

	startManga, endManga, err := computeSliceBounds(page, worksPerPage, len(mangaIDs))
	if err != nil {
		return UserWorkIDs{}, CountInfo{}, nil, err
	}

	startNovel, endNovel, err := computeSliceBounds(page, worksPerPage, len(novelIDs))
	if err != nil {
		return UserWorkIDs{}, CountInfo{}, nil, err
	}

	// Build the ID strings for each category
	var idsBuilder strings.Builder

	for _, k := range illustIDs[startIllust:endIllust] {
		idsBuilder.WriteString(fmt.Sprintf("&ids[]=%d", k))
	}
	userWorkIDs.IllustrationIDs = idsBuilder.String()

	idsBuilder.Reset()
	for _, k := range mangaIDs[startManga:endManga] {
		idsBuilder.WriteString(fmt.Sprintf("&ids[]=%d", k))
	}
	userWorkIDs.MangaIDs = idsBuilder.String()

	idsBuilder.Reset()
	for _, k := range novelIDs[startNovel:endNovel] {
		idsBuilder.WriteString(fmt.Sprintf("&ids[]=%d", k))
	}
	userWorkIDs.NovelIDs = idsBuilder.String()

	return userWorkIDs, countInfo, series, nil
}

func handleSeriesData(series json.RawMessage, category UserWorkCategory) ([]NovelSeries, []MangaSeries) {
	var novelSeries []NovelSeries
	var mangaSeries []MangaSeries

	if series != nil {
		if category == CategoryNovels {
			_ = json.Unmarshal(series, &novelSeries)
		} else {
			_ = json.Unmarshal(series, &mangaSeries)
		}
	}

	return novelSeries, mangaSeries
}

// getUserFrequentTags retrieves frequent tags for a user based on category.
func getUserFrequentTags(r *http.Request, ids string, category UserWorkCategory) ([]FrequentTag, error) {
	var tags []FrequentTag
	var URL string

	if category != "novels" {
		URL = GetUserFrequentArtworkTagsURL(ids)
	} else {
		URL = GetUserFrequentNovelTagsURL(ids)
	}

	response, err := API_GET_UnwrapJson(r.Context(), URL, "", r.Header)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal([]byte(response), &tags)
	if err != nil {
		return nil, err
	}

	return tags, nil
}

// Work is a generic type constraint.
type Work interface{}

// getWorkIDs is a generic helper function to fetch work IDs.
func getWorkIDs[T Work](r *http.Request, url string) ([]T, error) {
	resp, err := API_GET_UnwrapJson(r.Context(), url, "", r.Header)
	if err != nil {
		return nil, err
	}

	resp = session.ProxyImageUrl(r, resp)

	// Define a generic body structure.
	var body struct {
		Works map[int]json.RawMessage `json:"works"`
	}

	// Unmarshal the response into the body.
	if err := json.Unmarshal([]byte(resp), &body); err != nil {
		return nil, err
	}

	// Initialize the slice to hold the works.
	var works []T

	// Iterate over each work and unmarshal into the specific type.
	for _, v := range body.Works {
		var work T
		if err := json.Unmarshal(v, &work); err != nil {
			return nil, err
		}
		works = append(works, work)
	}

	return works, nil
}

// getUserArtworkIDs fetches the list of artwork IDs for a user (without metadata).
func getUserArtworkIDs(r *http.Request, id, ids string) ([]ArtworkBrief, error) {
	URL := GetUserFullArtworkURL(id, ids)

	works, err := getWorkIDs[ArtworkBrief](r, URL)
	if err != nil {
		return nil, err
	}

	return works, nil
}

// getUserNovelIDs fetches the list of novel IDs for a user (without metadata).
func getUserNovelIDs(r *http.Request, id, ids string) ([]NovelBrief, error) {
	URL := GetUserFullNovelURL(id, ids)

	works, err := getWorkIDs[NovelBrief](r, URL)
	if err != nil {
		return nil, err
	}

	return works, nil
}

// getUserBookmarks fetches the list of bookmarks for a user (with metadata).
//
// This function cannot be neatly refactored to use getWorkIDs due
// to having a different API response structure
func getUserBookmarks(r *http.Request, id, mode string, page int) ([]ArtworkBrief, int, error) {
	page--

	URL := GetUserBookmarksURL(id, mode, page)

	resp, err := API_GET_UnwrapJson(r.Context(), URL, "", r.Header)
	if err != nil {
		return nil, -1, err
	}
	resp = session.ProxyImageUrl(r, resp)

	var body struct {
		Artworks []json.RawMessage `json:"works"`
		Total    int               `json:"total"`
	}

	err = json.Unmarshal([]byte(resp), &body)
	if err != nil {
		return nil, -1, err
	}

	artworks := make([]ArtworkBrief, len(body.Artworks))

	for index, value := range body.Artworks {
		var artwork ArtworkBrief

		err = json.Unmarshal([]byte(value), &artwork)
		if err != nil {
			artworks[index] = ArtworkBrief{
				ID:        "#",
				Title:     "Deleted or Private",
				Thumbnail: "https://s.pximg.net/common/images/limit_unknown_360.png",
			}
			continue
		}
		artworks[index] = artwork
	}

	return artworks, body.Total, nil
}

func (s *User) parseSocial() error {
	if string(s.SocialRaw[:]) == "[]" {
		// Fuck Pixiv
		return nil
	}

	err := json.Unmarshal(s.SocialRaw, &s.Social)
	if err != nil {
		return err
	}
	return nil
}

// computeSliceBounds is a utility function to compute slice bounds safely
func computeSliceBounds(page int, worksPerPage float64, totalItems int) (start, end int, err error) {
	if totalItems == 0 {
		return 0, 0, nil
	}

	maxPages := int(math.Ceil(float64(totalItems) / worksPerPage))
	if page < 1 || page > maxPages {
		return 0, 0, i18n.Error("Invalid page number.")
	}

	start = (page - 1) * int(worksPerPage)
	end = min(start+int(worksPerPage), totalItems)

	return start, end, nil
}

func numberGreaterThan(l, r string) bool {
	if len(l) > len(r) {
		return true
	}
	if len(l) < len(r) {
		return false
	}
	return l > r
}