horse/cmd/horsebot/toradl.go

package main

import (
	"context"
	"errors"
	"fmt"
	"image"
	"image/png"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/gen2brain/avif"
	"golang.org/x/image/draw"
	"golang.org/x/sync/errgroup"
	"zombiezen.com/go/sqlite"
	"zombiezen.com/go/sqlite/sqlitex"
)

func saveImgs(ctx context.Context, db *sqlitex.Pool, out string) {
	skill, err := skillIconIDs(ctx, db)
	if err != nil {
		slog.Error("skills", slog.Any("err", err))
	}
	chara, err := charaIDs(ctx, db)
	if err != nil {
		slog.Error("characters", slog.Any("err", err))
	}
	rank, err := rankIDs(ctx, db)
	if err != nil {
		slog.Error("ranks", slog.Any("err", err))
	}

	g, ctx := errgroup.WithContext(ctx)
	tick := time.NewTicker(250 * time.Millisecond)
	idouts(ctx, g, tick, skill, filepath.Join(out, "skill"), "https://media.gametora.com/umamusume/skills/icon/%d.png", image.Pt(64, 64), image.Pt(32, 32), image.Pt(16, 16))
	idouts(ctx, g, tick, chara, filepath.Join(out, "chara"), "https://gametora.com/images/umamusume/characters/icons/chr_icon_%d.png", image.Pt(256, 256), image.Pt(64, 64), image.Pt(32, 32), image.Pt(16, 16))
	idouts(ctx, g, tick, rank, filepath.Join(out, "rank"), "https://media.gametora.com/umamusume/ui/rank/%02d.png", image.Pt(128, 128), image.Pt(64, 64), image.Pt(32, 32), image.Pt(16, 16))
	idouts(ctx, g, tick, rank, filepath.Join(out, "rank", "simple"), "https://media.gametora.com/umamusume/ui/rank/simple/%02d.png", image.Pt(50, 50), image.Pt(32, 32), image.Pt(16, 16))
	if err := g.Wait(); err != nil {
		slog.Error("output", slog.Any("err", err))
	}
}

// load scans all results of sql and appends them to r.
func load[T any](ctx context.Context, db *sqlitex.Pool, r []T, sql string, row func(*sqlite.Stmt) T) ([]T, error) {
	conn, err := db.Take(ctx)
	defer db.Put(conn)
	if err != nil {
		return nil, err
	}
	stmt, err := conn.Prepare(sql)
	if err != nil {
		return nil, err
	}

	for {
		ok, err := stmt.Step()
		if err != nil {
			return r, err
		}
		if !ok {
			break
		}
		r = append(r, row(stmt))
	}
	return r, err
}

func skillIconIDs(ctx context.Context, db *sqlitex.Pool) ([]int32, error) {
	return load(ctx, db, nil, "SELECT DISTINCT icon_id FROM skill_data", func(s *sqlite.Stmt) int32 {
		return s.ColumnInt32(0)
	})
}

func charaIDs(ctx context.Context, db *sqlitex.Pool) ([]int32, error) {
	return load(ctx, db, nil, "SELECT id FROM chara_data", func(s *sqlite.Stmt) int32 {
		return s.ColumnInt32(0)
	})
}

func rankIDs(ctx context.Context, db *sqlitex.Pool) ([]int32, error) {
	return load(ctx, db, nil, "SELECT id-1 FROM single_mode_rank", func(s *sqlite.Stmt) int32 {
		return s.ColumnInt32(0)
	})
}

var pngenc = &png.Encoder{CompressionLevel: png.BestCompression}

func idouts(ctx context.Context, g *errgroup.Group, tick *time.Ticker, ids []int32, p, urlf string, sizes ...image.Point) {
	for _, id := range ids {
		g.Go(func() error {
			n := strconv.Itoa(int(id))
			p := filepath.Join(p, n)
			_, err := os.Stat(p)
			if !errors.Is(err, os.ErrNotExist) {
				slog.DebugContext(ctx, "exists", slog.String("path", p))
				return nil
			}
			if err := os.MkdirAll(p, 0777); err != nil {
				return err
			}
			url := fmt.Sprintf(urlf, id)
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-tick.C: // continue on
			}
			img, err := dlimg(ctx, url)
			if err != nil {
				return err
			}
			slog.InfoContext(ctx, "resize", slog.String("url", url))
			rs := resize(img, sizes...)
			for _, r := range rs {
				g.Go(func() error {
					m := r.Bounds().Max
					o := filepath.Join(p, fmt.Sprintf("%dx%d.png", m.X, m.Y))
					slog.InfoContext(ctx, "write", slog.String("path", o))
					f, err := os.Create(o)
					if err != nil {
						return err
					}
					defer f.Close()
					return pngenc.Encode(f, r)
				})
				g.Go(func() error {
					m := r.Bounds().Max
					o := filepath.Join(p, fmt.Sprintf("%dx%d.avif", m.X, m.Y))
					slog.InfoContext(ctx, "write", slog.String("path", o))
					f, err := os.Create(o)
					if err != nil {
						return err
					}
					defer f.Close()
					return avif.Encode(f, r)
				})
			}
			return nil
		})
	}
}

func dlimg(ctx context.Context, url string) (image.Image, error) {
	slog.InfoContext(ctx, "download", slog.String("url", url))
	ctx, stop := context.WithTimeout(ctx, 5*time.Second)
	defer stop()
	req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
	if err != nil {
		return nil, err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	slog.InfoContext(ctx, "decode", slog.String("url", url))
	img, _, err := image.Decode(resp.Body)
	return img, err
}

func resize(img image.Image, sizes ...image.Point) []image.Image {
	rs := make([]image.Image, 0, len(sizes))
	for _, sz := range sizes {
		d := image.NewNRGBA(image.Rectangle{Min: image.Point{}, Max: sz})
		draw.CatmullRom.Scale(d, d.Bounds(), img, img.Bounds(), draw.Src, nil)
		rs = append(rs, d)
	}
	return rs
}