package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
)

type CliFlags struct {
	outputFileName   string
	manifestFileName string
	userAgent        string
	verbose          bool
}

type UgoiraManifestFrame struct {
	File  string `json:"file"`
	Delay int    `json:"delay"`
}

type UgoiraManifestBodyDesc struct {
	Src         string                `json:"src"`
	OriginalSrc string                `json:"originalSrc"`
	MimeType    string                `json:"mime_type"`
	Frames      []UgoiraManifestFrame `json:"frames"`
}

type UgoiraManifest struct {
	Error   bool                   `json: "error"`
	Message string                 `json:"message"`
	Body    UgoiraManifestBodyDesc `json:"body"`
}

type PostId uint64

const ArtworkUrlTemplate = "https://www.pixiv.net/en/artworks/%d"
const ManifestUrlTemplate = "https://www.pixiv.net/ajax/illust/%d/ugoira_meta"

func doPixivRequest(client *http.Client, url, referer, userAgent string) io.ReadCloser {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		die("couldn't create request: %s", err)
	}

	req.Header.Set("Authority", "www.pixiv.net")
	req.Header.Set("Sec-Fetch-Dest", "empty")
	req.Header.Set("Sec-Fetch-Mode", "cors")
	req.Header.Set("Sec-Fetch-Site", "same-origin")
	req.Header.Set("Referer", referer)
	if len(userAgent) > 0 {
		req.Header.Set("User-Agent", userAgent)
	}

	response, err := client.Do(req)
	if err != nil {
		die("GET request failed: %s", err)
	}
	if response.StatusCode != http.StatusOK {
		die("Got non-200 status code: %d", response.StatusCode)
	}
	return response.Body
}

func decodeUgoiraManifest(jsonData []byte) UgoiraManifest {
	var manifest UgoiraManifest
	if err := json.Unmarshal(jsonData, &manifest); err != nil {
		die("Could not parse manifest JSON: %s", err)
	}
	if manifest.Error {
		die("Manifest returned error: %s", manifest.Message)
	}
	return manifest
}

func getUgoira(postId PostId, flags CliFlags) {
	client := &http.Client{}
	artworkUrl := fmt.Sprintf(ArtworkUrlTemplate, postId)
	manifestUrl := fmt.Sprintf(ManifestUrlTemplate, postId)

	say(flags.verbose, "Fetching and decoding manifest")
	manifestData := doPixivRequest(client, manifestUrl, artworkUrl, flags.userAgent)
	defer manifestData.Close()
	// Read the entire JSON in memory, since we need to dump it to a file later
	manifestJsonData, err := io.ReadAll(manifestData)
	if err != nil {
		die("error while fetching manifest: %s", err)
	}
	// Now decode the JSON...
	manifest := decodeUgoiraManifest(manifestJsonData)
	if len(manifest.Body.OriginalSrc) == 0 {
		die("Ugoira url is empty")
	}
	say(flags.verbose, "Ugoira file URL: %s", manifest.Body.OriginalSrc)

	// Download ugoira data
	var ugoiraFileName string
	if len(flags.outputFileName) > 0 {
		ugoiraFileName = flags.outputFileName
	} else {
		ugoiraFileName = fmt.Sprintf("./%d_ugoira.zip", postId)
	}
	say(flags.verbose, "Downloading ugoira (saving to: %s)", ugoiraFileName)

	ugoiraData := doPixivRequest(client, manifest.Body.OriginalSrc, artworkUrl, flags.userAgent)
	ugoiraFile, err := os.Create(ugoiraFileName)
	if err != nil {
		die("Could not create ugoira file: %s", err)
	}
	defer ugoiraFile.Close()
	if _, err := io.Copy(ugoiraFile, ugoiraData); err != nil {
		die("Ugoira download failed: %s", err)
	}

	// Dump manifest JSON
	var ugoiraManifestFileName string
	if len(flags.manifestFileName) > 0 {
		ugoiraManifestFileName = flags.manifestFileName
	} else {
		ugoiraManifestFileName = fmt.Sprintf("./%d_ugoira.json")
	}
	say(flags.verbose, "Dumping manifest (to: %s)", ugoiraManifestFileName)

	ugoiraManifestFile, err := os.Create(ugoiraManifestFileName)
	if err != nil {
		die("Could not create ugoira manifest file: %s", err)
	}
	defer ugoiraManifestFile.Close()
	if _, err := ugoiraManifestFile.Write(manifestJsonData); err != nil {
		die("Ugoira manifest dump failed: %s", err)
	}

	say(flags.verbose, "Successfully downloaded post %d", postId)
}

func die(format string, args ...any) {
	fmt.Fprintf(os.Stderr, os.Args[0]+": "+format+"\n", args...)
	os.Exit(1)
}

func say(enabled bool, format string, args ...any) {
	if !enabled {
		return
	}
	fmt.Printf(format+"\n", args...)
}

func main() {
	var flags CliFlags
	flag.StringVar(&flags.outputFileName, "o", "", "Ugoira output file name")
	flag.StringVar(&flags.manifestFileName, "M", "", "Ugoira manifest output file name")
	flag.StringVar(&flags.userAgent, "A", "", "User agent to use when doing GET requests")
	flag.BoolVar(&flags.verbose, "v", false, "Enable verbosity")
	flag.Parse()

	args := flag.Args()
	if len(args) == 0 {
		die("expected post ID")
	}

	postIdStr := args[0]
	postId, err := strconv.ParseUint(postIdStr, 10, 0)
	if err != nil {
		die("Invalid post ID (has to be numerical): '%s'", postIdStr)
	}
	say(flags.verbose, "Working on post %d", postId)
	getUgoira(PostId(postId), flags)
}