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/audit"
	"codeberg.org/vnpower/pixivfe/v2/server/session"
	"github.com/goccy/go-json"
	"go.uber.org/zap"
)

// UserWorkCategory is a struct that represents a category of user work with a page limit.
//
// pixivfe internal data type. not used by pixiv.
type UserWorkCategory struct {
	Value     string
	PageLimit int
}

// NewUserWorkCategory creates UserWorkCategory instances
func NewUserWorkCategory(value string) UserWorkCategory {
	return UserWorkCategory{Value: value, PageLimit: 0}
}

func (c *UserWorkCategory) SetPageLimit(limit int) {
	c.PageLimit = limit
}

var (
	CategoryAny          = NewUserWorkCategory("")
	CategoryAnyAlt       = NewUserWorkCategory("artworks")
	CategoryIllustration = NewUserWorkCategory("illustrations")
	CategoryManga        = NewUserWorkCategory("manga")
	CategoryBookmarks    = NewUserWorkCategory("bookmarks")
	CategoryNovels       = NewUserWorkCategory("novels")
)

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

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

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

type WorkCounts 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"`
	// the Artworks slice includes both illustrations and manga, but not novels
	// it's also used to hold the []ArtworkBrief when viewing a user's bookmarks (TODO to separate?)
	Artworks          []ArtworkBrief `json:"artworks"`
	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

	// Used to display the number of works for a given category
	WorkCounts WorkCounts
}

// GetUserProfile retrieves the user profile, including counts, artworks/bookmarks, and social data.
func GetUserProfile(auditor *audit.Auditor, 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(), auditor, 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
	}

	// Populate work counts
	if err := populateWorkCounts(auditor, r, &user, id); err != nil {
		return user, err
	}

	// Get populated works or bookmarks based on the category
	if err := getPopulatedWorks(auditor, 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
}

// populateWorkCounts populates the WorkCounts struct.
func populateWorkCounts(auditor *audit.Auditor, r *http.Request, user *User, id string) error {
	// Get counts for illustrations and manga, as well as novels
	_, countInfo, _, err := fetchWorkIDsAndSeriesData(auditor, r, id, &CategoryAny, 1)
	if err != nil {
		return err
	}
	user.WorkCounts = countInfo

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

	return nil
}

// getPopulatedWorks fetches then populates information for a user's works or bookmarks, based on the category.
//
// It also handles a user's frequently used tags and series data.
func getPopulatedWorks(auditor *audit.Auditor, r *http.Request, user *User, id string, category *UserWorkCategory, page int, getTags bool) error {
	if category.Value == CategoryBookmarks.Value {
		// Fetch bookmarks
		works, count, err := fetchBookmarks(auditor, r, id, "show", page)
		if err != nil {
			return err
		}

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

		// Calculate PageLimit for bookmarks
		worksPerPage := 48.0
		pageLimit := int(math.Ceil(float64(count) / worksPerPage))
		category.SetPageLimit(pageLimit)

		return nil
	}

	// Fetch work IDs and series data
	userWorkIDs, countInfo, series, err := fetchWorkIDsAndSeriesData(auditor, r, id, category, page)
	if err != nil {
		return err
	}

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

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

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

	// println(illustrationIDs == "")

	// Fetch and populate []ArtworkBrief for the IDs
	if illustrationIDs != "" || mangaIDs != "" {
		// Combine illustration and manga IDs
		ids := illustrationIDs + mangaIDs
		artworks, err := fetchPopulatedBriefs[ArtworkBrief](auditor, r, id, ids, fetchArtworkIDs)
		if err != nil {
			return err
		}

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

	// Fetch and populate []NovelBrief for the IDs
	if novelIDs != "" {
		ids := novelIDs
		novels, err := fetchPopulatedBriefs[NovelBrief](auditor, r, id, ids, fetchNovelIDs)
		if err != nil {
			return err
		}
		user.Novels = novels
	}

	// Fetch a user's frequently used tags if requested
	if getTags {
		var tagsIDs []string

		// Since category is pre-validated, we don't need to handle unexpected categories
		if category.Value == CategoryAny.Value ||
			category.Value == CategoryAnyAlt.Value ||
			category.Value == CategoryIllustration.Value ||
			category.Value == CategoryManga.Value {
			// Append illustrationIDs and mangaIDs if they are not empty
			if illustrationIDs != "" {
				auditor.Logger.Debug("illustrationIDs found, appending to tagsIDs")
				tagsIDs = append(tagsIDs, illustrationIDs)
			}
			if mangaIDs != "" {
				auditor.Logger.Debug("mangaIDs found, appending to tagsIDs")
				tagsIDs = append(tagsIDs, mangaIDs)
			}
		} else if category.Value == CategoryNovels.Value {
			auditor.Logger.Debug("Using CategoryNovels to fetch tags")
			if novelIDs != "" {
				auditor.Logger.Debug("novelIDs found, appending to tagsIDs")
				tagsIDs = append(tagsIDs, novelIDs)
			}
		} else {
			auditor.Logger.Warn("No recognized category found")
			return nil
		}

		if len(tagsIDs) == 0 {
			auditor.Logger.Warn("No tags found to fetch, setting FrequentTags to empty slice",
				zap.String("category", category.Value),
				zap.Strings("tagsIDs", tagsIDs),
				zap.String("illustrationIDs", illustrationIDs),
				zap.String("mangaIDs", mangaIDs),
				zap.String("novelIDs", novelIDs),
			)
			user.FrequentTags = []FrequentTag{}
		} else {
			// Concatenate IDs to form a single string
			ids := strings.Join(tagsIDs, "")
			user.FrequentTags, err = fetchFrequentTags(auditor, r, ids, category)
			if err != nil {
				return fmt.Errorf("failed to get user frequent tags: %w", 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
}

// fetchPopulatedBriefs is a generic function that fetches and processes items of type T,
// where T must implement the HasID interface.
func fetchPopulatedBriefs[T HasID](
	auditor *audit.Auditor,
	r *http.Request,
	id, ids string,
	fetchFunc func(*audit.Auditor, *http.Request, string, string) ([]T, error),
) ([]T, error) {
	items, err := fetchFunc(auditor, 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
}

// extractIDs is a helper function to extract IDs and count from a raw JSON message
func extractIDs(rawMessage *json.RawMessage, categoryName string) ([]int, int, error) {
	var dataMap IntStringMap
	if rawMessage != nil && len(*rawMessage) > 0 {
		if err := json.Unmarshal(*rawMessage, &dataMap); err != nil {
			fmt.Printf("Error unmarshalling %s: %v\n", categoryName, err)
			dataMap = make(IntStringMap)
		}
	} else {
		dataMap = make(IntStringMap)
	}

	ids := make([]int, 0, len(dataMap))
	for k := range dataMap {
		ids = append(ids, k)
	}
	return ids, len(dataMap), nil
}

// fetchWorkIDsAndSeriesData retrieves artwork IDs and series information for a user.
//
// It returns separate ID strings for illustrations, manga, and novels encapsulated in UserWorkIDs.
func fetchWorkIDsAndSeriesData(auditor *audit.Auditor, r *http.Request, id string, category *UserWorkCategory, page int) (userWorkIDs UserWorkIDs, workCounts WorkCounts, series json.RawMessage, err error) {
	URL := GetUserWorksURL(id)

	resp, err := API_GET_UnwrapJson(r.Context(), auditor, URL, "", r.Header)
	if err != nil {
		return UserWorkIDs{}, WorkCounts{}, 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{}, WorkCounts{}, nil, fmt.Errorf("failed to unmarshal response body: %w", err)
	}

	workCounts = WorkCounts{}

	var illustIDs, mangaIDs, novelIDs []int

	// Define the categories to process
	categories := []struct {
		name           string
		rawMessage     *json.RawMessage
		seriesMessage  *json.RawMessage
		ids            *[]int
		countInfoField *int
	}{
		{
			name:           "illusts",
			rawMessage:     body.Illusts,
			seriesMessage:  nil, // illustrations don't have series in this context
			ids:            &illustIDs,
			countInfoField: &workCounts.Illustrations,
		},
		{
			name:           "manga",
			rawMessage:     body.Manga,
			seriesMessage:  body.MangaSeries,
			ids:            &mangaIDs,
			countInfoField: &workCounts.Manga,
		},
		{
			name:           "novels",
			rawMessage:     body.Novels,
			seriesMessage:  body.NovelSeries,
			ids:            &novelIDs,
			countInfoField: &workCounts.Novels,
		},
	}

	// Process each category
	for _, cat := range categories {
		ids, count, err := extractIDs(cat.rawMessage, cat.name)
		if err != nil {
			return UserWorkIDs{}, WorkCounts{}, nil, err
		}
		*(cat.ids) = ids
		workCounts.All += count
		*(cat.countInfoField) = count

		// Handle series data
		if cat.seriesMessage != nil && len(*cat.seriesMessage) > 0 {
			series = *cat.seriesMessage
		}
	}

	// Sort IDs in descending order for each category
	idSlices := []*[]int{&illustIDs, &mangaIDs, &novelIDs}
	for _, ids := range idSlices {
		sort.Sort(sort.Reverse(sort.IntSlice(*ids)))
	}

	worksPerPage := 30.0

	// Build the UserWorkIDs by computing slice bounds and constructing ID strings
	idCategories := []struct {
		name          string
		ids           *[]int
		idStringField *string
	}{
		{
			name:          "illusts",
			ids:           &illustIDs,
			idStringField: &userWorkIDs.IllustrationIDs,
		},
		{
			name:          "manga",
			ids:           &mangaIDs,
			idStringField: &userWorkIDs.MangaIDs,
		},
		{
			name:          "novels",
			ids:           &novelIDs,
			idStringField: &userWorkIDs.NovelIDs,
		},
	}

	var totalItemsForCategory int

	for _, cat := range idCategories {
		totalItems := len(*cat.ids)
		subCategory := determineSubCategory(cat.name)

		auditor.Logger.Debug("Entering computeSliceBounds",
			zap.String("category", cat.name),
		)

		start, end, pageLimit, err := computeSliceBounds(auditor, page, worksPerPage, totalItems)
		if err != nil {
			fmt.Printf("Error computing slice bounds for category '%s': %v\n", cat.name, err)
			*(cat.idStringField) = ""
			continue // skip the broken category and continue
		}

		// Set PageLimit on the correct category
		if subCategory.Value == category.Value {
			category.SetPageLimit(pageLimit)
		} else if category.Value == "" || category.Value == "artworks" {
			// For "any" or "artworks", accumulate total items and compute PageLimit later
			totalItemsForCategory += totalItems
		}

		// Build IDs
		idsToUse := (*cat.ids)[start:end]
		var idsBuilder strings.Builder
		for _, k := range idsToUse {
			idsBuilder.WriteString(fmt.Sprintf("&ids[]=%d", k))
		}
		*(cat.idStringField) = idsBuilder.String()
	}

	// If category is "any" or "artworks", set the PageLimit after summing total items
	// if category.Value == "" || category.Value == "artworks" {
	// 	pageLimit := int(math.Ceil(float64(totalItemsForCategory) / worksPerPage))
	// 	category.SetPageLimit(pageLimit)
	// }

	return userWorkIDs, workCounts, series, nil
}

func determineSubCategory(name string) UserWorkCategory {
	switch name {
	case "illusts":
		return CategoryIllustration
	case "manga":
		return CategoryManga
	case "novels":
		return CategoryNovels
	default:
		return CategoryAny
	}
}

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
}

// fetchFrequentTags fetches a user's frequently used tags, based on category.
func fetchFrequentTags(auditor *audit.Auditor, r *http.Request, ids string, category *UserWorkCategory) ([]FrequentTag, error) {
	var tags []FrequentTag
	var URL string

	// TODO: If overall, request both?
	if category != &CategoryNovels {
		URL = GetUserFrequentArtworkTagsURL(ids)
	} else {
		URL = GetUserFrequentNovelTagsURL(ids)
	}

	// Return early if there are no IDs
	//
	// NOTE: theoretically, this check should never be reached as fetchUserWorks doesn't call
	// getUserFrequentTags when len(tagsIDs) == 0, instead returning an empty []FrequentTag
	if ids == "" {
		return tags, nil
	}

	response, err := API_GET_UnwrapJson(r.Context(), auditor, 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{}

// fetchWorkIDs is a generic helper function to fetch work IDs.
func fetchWorkIDs[T Work](auditor *audit.Auditor, r *http.Request, url string) ([]T, error) {
	resp, err := API_GET_UnwrapJson(r.Context(), auditor, 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
}

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

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

	return works, nil
}

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

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

	return works, nil
}

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

	URL := GetUserBookmarksURL(id, mode, page)

	resp, err := API_GET_UnwrapJson(r.Context(), auditor, 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(auditor *audit.Auditor, page int, worksPerPage float64, totalItems int) (start, end, pageLimit int, err error) {
	auditor.Logger.Debug("Entered computeSliceBounds",
		zap.Int("page", page),
		zap.Float64("worksPerPage", worksPerPage),
		zap.Int("totalItems", totalItems),
	)

	if totalItems == 0 {
		auditor.Logger.Debug("totalItems is 0, returning 0, 0, nil")
		return 0, 0, 0, nil
	}

	pageLimit = int(math.Ceil(float64(totalItems) / worksPerPage))
	auditor.Logger.Debug("Calculated pageLimit",
		zap.Int("pageLimit", pageLimit),
	)

	if page < 1 || page > pageLimit {
		auditor.Logger.Debug("Invalid page number",
			zap.Int("page", page),
			zap.Int("pageLimit", pageLimit),
		)
		return 0, 0, 0, fmt.Errorf("invalid page number")
	}

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

	auditor.Logger.Debug("Calculated slice bounds",
		zap.Int("start", start),
		zap.Int("end", end),
	)
	auditor.Logger.Debug("Set pageLimit",
		zap.Int("pageLimit", pageLimit),
	)

	auditor.Logger.Debug("Exited computeSliceBounds successfully")

	return start, end, pageLimit, nil
}

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