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 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(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
	}

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

	// Get populated works or bookmarks based on the category
	if err := getPopulatedWorks(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(r *http.Request, user *User, id string) error {
	// Get counts for illustrations and manga, as well as novels
	_, countInfo, _, err := fetchWorkIDsAndSeriesData(r, id, CategoryAny, 1)
	if err != nil {
		return err
	}
	user.WorkCounts = countInfo

	// Get count for bookmarks
	_, bookmarksCount, err := fetchBookmarks(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(r *http.Request, user *User, id string, category UserWorkCategory, page int, getTags bool) error {
	if category == CategoryBookmarks {
		// Fetch bookmarks
		works, count, err := fetchBookmarks(r, id, "show", page)
		if err != nil {
			return err
		}

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

	// Fetch work IDs and series data
	//
	// At this stage, these are only simple string slices that are not populated,
	// as pixiv only returns the work IDs at the GetUserWorksURL endpoint.
	userWorkIDs, countInfo, series, err := fetchWorkIDsAndSeriesData(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
	}

	println(illustrationIDs == "")

	// Fetch and populate []ArtworkBrief for the IDs
	if illustrationIDs != "" || mangaIDs != "" {
		// Combine illustration and manga IDs
		ids := illustrationIDs + mangaIDs
		artworks, err := fetchPopulatedBriefs[ArtworkBrief](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](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

		// TODO: frequent tags for novels are *not* fetched when CategoryAny or CategoryAnyAlt

		// Since category is pre-validated, we don't need to handle unexpected categories
		if category == CategoryAny || category == CategoryAnyAlt || category == CategoryIllustration || category == CategoryManga {
			// Append illustrationIDs and mangaIDs if they are not empty
			if illustrationIDs != "" {
				tagsIDs = append(tagsIDs, illustrationIDs)
			}
			if mangaIDs != "" {
				tagsIDs = append(tagsIDs, mangaIDs)
			}
		} else if category == CategoryNovels {
			// Append novelIDs if not empty
			if novelIDs != "" {
				tagsIDs = append(tagsIDs, novelIDs)
			}
		}

		if len(tagsIDs) == 0 {
			// No tags to fetch, set FrequentTags to empty slice
			user.FrequentTags = []FrequentTag{}
		} else {
			// Concatenate IDs form a single string
			ids := strings.Join(tagsIDs, "")
			user.FrequentTags, err = fetchFrequentTags(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](
	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
}

// 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(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(), 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,
		},
	}

	for _, cat := range idCategories {
		totalItems := len(*cat.ids)
		start, end, err := computeSliceBounds(page, worksPerPage, totalItems, cat.name)
		if err != nil {
			fmt.Printf("Error computing slice bounds for category '%s': %v\n", cat.name, err)
			// Set idStringField to an empty string so that it's caught by getPopulatedWorks correctly
			*(cat.idStringField) = ""
			continue
		}
		idsToUse := (*cat.ids)[start:end]
		var idsBuilder strings.Builder
		for _, k := range idsToUse {
			idsBuilder.WriteString(fmt.Sprintf("&ids[]=%d", k))
		}
		*(cat.idStringField) = idsBuilder.String()
	}

	return userWorkIDs, workCounts, 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
}

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

	// TODO: If overall, request both?
	if category != "novels" {
		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(), 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](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
}

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

	works, err := fetchWorkIDs[ArtworkBrief](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(r *http.Request, id, ids string) ([]NovelBrief, error) {
	URL := GetUserFullNovelURL(id, ids)

	works, err := fetchWorkIDs[NovelBrief](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(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, categoryName string) (start, end int, err error) {
	// fmt.Printf("[DEBUG] Entering computeSliceBounds - page: %d, worksPerPage: %.2f, totalItems: %d, category: %s\n", page, worksPerPage, totalItems, categoryName)

	if totalItems == 0 {
		// fmt.Printf("[DEBUG] totalItems is 0 for category '%s', returning 0, 0, nil\n", categoryName)
		return 0, 0, nil
	}

	maxPages := int(math.Ceil(float64(totalItems) / worksPerPage))
	// fmt.Printf("[DEBUG] Calculated maxPages: %d for category '%s'\n", maxPages, categoryName)

	if page < 1 || page > maxPages {
		// fmt.Printf("[DEBUG] Invalid page number for category '%s'. page: %d, maxPages: %d\n", categoryName, page, maxPages)
		return 0, 0, fmt.Errorf("Invalid page number for category '%s'", categoryName)
	}

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

	// fmt.Printf("[DEBUG] Calculated start: %d, end: %d for category '%s'\n", start, end, categoryName)

	// fmt.Printf("[DEBUG] Exiting computeSliceBounds successfully for category '%s'\n", categoryName)
	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
}