package main

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

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

type PostId uint64

const (
	ArtworkUrlTemplate  = "https://www.pixiv.net/en/artworks/%d"
	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 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, err := decodeUgoiraManifest(manifestJsonData)
	if err != nil {
		die("%s", err)
	}
	if manifest.Error {
		die("Manifest returned error: %s", manifest.Message)
	}
	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", postId)
	}
	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 downloadUgoiraCmd(args []string, flags CliFlags) {
	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)
}

func ugoira2videoCmd(args []string, flags CliFlags) {
	if len(args) < 2 {
		die("expected ugoira zip file and manifest JSON")
	}
	ugoiraFileName, manifestFileName := args[0], args[1]
	mfp, err := os.Open(manifestFileName)
	if err != nil {
		die("Could not open manifest file: %s", err)
	}
	defer mfp.Close()

	manifestData, err := io.ReadAll(mfp)
	if err != nil {
		die("Could not read manifest: %s", err)
	}

	manifest, err := decodeUgoiraManifest(manifestData)
	if err != nil {
		die("%s", err)
	}
	err = ugoira2video(manifest, ugoiraFileName, flags.outputFileName, flags.ffmpegArgs)
	if err != nil {
		die("Conversion failed: %s", err)
	}
}

func main() {
	var flags CliFlags

	downloadCmd := flag.NewFlagSet("download", flag.ExitOnError)
	downloadCmd.StringVar(&flags.outputFileName, "o", "", "Ugoira output file name")
	downloadCmd.StringVar(&flags.manifestFileName, "M", "", "Ugoira manifest output file name")
	downloadCmd.StringVar(&flags.userAgent, "A", "", "User agent to use when doing GET requests")
	downloadCmd.BoolVar(&flags.verbose, "v", false, "Enable verbosity")

	u2vCmd := flag.NewFlagSet("u2v", flag.ExitOnError)
	u2vCmd.StringVar(&flags.outputFileName, "o", "", "Converted Ugoira output file")
	u2vCmd.StringVar(
		&flags.ffmpegArgs,
		"ffmpeg-args",
		"",
		"Arguments to be supplied to FFmpeg (a single `%s` acts as a placeholder for the output file name)",
	)

	if len(os.Args) < 2 {
		die("expected command")
	}

	switch os.Args[1] {
	case "download", "down", "d":
		downloadCmd.Parse(os.Args[2:])
		downloadUgoiraCmd(downloadCmd.Args(), flags)
	case "u2v":
		u2vCmd.Parse(os.Args[2:])
		if len(flags.outputFileName) == 0 {
			die("Expected filename output for video")
		}
		ugoira2videoCmd(u2vCmd.Args(), flags)
	case "-h":
		fmt.Printf("Commands: download, u2v\nUse '%s <command> -h' to see the command-specific flags.\n", os.Args[0])
	default:
		die("invalid command '%s'", os.Args[1])
	}
}