PixivFE/core/rankingCalendar.go

package core

import (
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"time"

	"codeberg.org/vnpower/pixivfe/v2/audit"
	"github.com/andybalholm/cascadia"
	"golang.org/x/net/html"

	"codeberg.org/vnpower/pixivfe/v2/server/session"
)

// DayCalendar represents the data for a single day in the ranking calendar
type DayCalendar struct {
	DayNumber   int    // The day of the month
	DateString  string // Pixiv-compatible string that represents the date (format: YYYYMMDD)
	ImageURL    string // Proxy URL to the image (optional, can be empty when no image is available)
	ArtworkLink string // The link to the artwork page for this day
}

// Precompiled CSS selector and regex
var (
	selectorImg    = cascadia.MustCompile("img")
	artworkIDRegex = regexp.MustCompile(`/(\d+)_p0_(custom|square)1200\.jpg`)
)

// GetRankingCalendar retrieves and processes the ranking calendar data from Pixiv.
// It returns a slice of DayCalendar structs and any error encountered.
//
// iacore: so the funny thing about Pixiv is that they will return this month's data for a request of a future date. is it a bug or a feature?
func GetRankingCalendar(auditor *audit.Auditor, r *http.Request, mode string, year, month int) ([]DayCalendar, error) {
	token := session.GetUserToken(r)
	URL := GetRankingCalendarURL(mode, year, month)

	resp, err := API_GET(r.Context(), auditor, URL, token, r.Header)
	if err != nil {
		return nil, err
	}

	doc, err := html.Parse(strings.NewReader(resp.Body))
	if err != nil {
		return nil, err
	}

	links, err := extractImageLinks(r, doc)
	if err != nil {
		return nil, err
	}

	calendar := []DayCalendar{}
	dayCount := 0

	// Get the first day of the month
	firstDayOfMonth := time.Date(year, time.Month(month), 1, 0, 0, 0, 0, time.UTC)

	// Add empty days before the first day of the month
	calendar, dayCount = addEmptyDaysBefore(calendar, firstDayOfMonth, dayCount)

	// Get the number of days in the month
	numDays := daysIn(time.Month(month), year)

	// Add days of the month
	calendar, dayCount = addDaysOfMonth(calendar, links, dayCount, numDays, year, month)

	// Add empty days after the last day to complete the week
	calendar, dayCount = addEmptyDaysAfter(calendar, dayCount)

	return calendar, nil
}

// extractImageLinks extracts image links from the parsed HTML document.
func extractImageLinks(r *http.Request, doc *html.Node) ([]string, error) {
	var links []string
	for _, node := range cascadia.QueryAll(doc, selectorImg) {
		for _, attr := range node.Attr {
			if attr.Key == "data-src" {
				url, err := session.ProxyImageUrlNoEscape(r, attr.Val)
				if err != nil {
					return nil, err
				}
				links = append(links, url)
			}
		}
	}
	return links, nil
}

// addEmptyDaysBefore adds empty days to the calendar before the first day of the month.
func addEmptyDaysBefore(calendar []DayCalendar, firstDay time.Time, dayCount int) ([]DayCalendar, int) {
	emptyDays := int(firstDay.Weekday())
	for i := 0; i < emptyDays; i++ {
		calendar = append(calendar, DayCalendar{DayNumber: 0})
		dayCount++
	}
	return calendar, dayCount
}

// addDaysOfMonth adds the actual days of the month to the calendar.
func addDaysOfMonth(calendar []DayCalendar, links []string, dayCount, numDays, year, month int) ([]DayCalendar, int) {
	for i := 0; i < numDays; i++ {
		day := DayCalendar{
			DayNumber: i + 1,
		}
		if len(links) > i {
			day.ImageURL = links[i]
			artworkID := extractArtworkID(links[i])
			if artworkID != "" {
				day.ArtworkLink = fmt.Sprintf("/artworks/%s", artworkID)
			}
		}
		day.DateString = fmt.Sprintf("%d%02d%02d", year, month, i+1)
		calendar = append(calendar, day)
		dayCount++
	}
	return calendar, dayCount
}

// addEmptyDaysAfter adds empty days to the calendar after the last day of the month to complete the week.
func addEmptyDaysAfter(calendar []DayCalendar, dayCount int) ([]DayCalendar, int) {
	for dayCount%7 != 0 {
		calendar = append(calendar, DayCalendar{DayNumber: 0})
		dayCount++
	}
	return calendar, dayCount
}

// daysIn returns the number of days in a given month and year.
func daysIn(month time.Month, year int) int {
	return time.Date(year, month+1, 0, 0, 0, 0, 0, time.UTC).Day()
}

// extractArtworkID extracts the artwork ID from the image URL.
func extractArtworkID(imageURL string) string {
	matches := artworkIDRegex.FindStringSubmatch(imageURL)
	if len(matches) > 1 {
		return matches[1]
	}
	return ""
}