package main

import (
	"context"
	"fmt"
	"os"
	"time"

	"predictor-refactored/internal/dataset"
	"predictor-refactored/internal/downloader"

	"go.uber.org/zap"
)

// Downloads step 0 of a given run and writes a minimal dataset for comparison.
// Usage: go run ./cmd/compare_step0 <run_YYYYMMDDHH> <output_path>
func main() {
	if len(os.Args) < 3 {
		fmt.Fprintf(os.Stderr, "Usage: %s <run_YYYYMMDDHH> <output_path>\n", os.Args[0])
		os.Exit(1)
	}

	runStr := os.Args[1]
	outPath := os.Args[2]

	run, err := time.Parse("2006010215", runStr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Invalid run time %q: %v\n", runStr, err)
		os.Exit(1)
	}

	log, _ := zap.NewDevelopment()

	// Create a full-size dataset (we only fill step 0)
	fmt.Printf("Creating dataset at %s (%d bytes)...\n", outPath, dataset.DatasetSize)
	ds, err := dataset.Create(outPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Create dataset: %v\n", err)
		os.Exit(1)
	}
	defer ds.Close()

	cfg := &downloader.Config{
		DataDir:  os.TempDir(),
		Parallel: 4,
	}
	dl := downloader.NewDownloader(cfg, log)

	ctx := context.Background()
	date := run.Format("20060102")
	runHour := run.Hour()

	// Download and blit step 0 from pgrb2
	fmt.Println("Downloading pgrb2 step 0...")
	urlA := dataset.GribURL(date, runHour, 0)
	if err := dl.DownloadAndBlit(ctx, ds, urlA, 0, dataset.LevelSetA); err != nil {
		fmt.Fprintf(os.Stderr, "pgrb2: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("  done")

	// Download and blit step 0 from pgrb2b
	fmt.Println("Downloading pgrb2b step 0...")
	urlB := dataset.GribURLB(date, runHour, 0)
	if err := dl.DownloadAndBlit(ctx, ds, urlB, 0, dataset.LevelSetB); err != nil {
		fmt.Fprintf(os.Stderr, "pgrb2b: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("  done")

	if err := ds.Flush(); err != nil {
		fmt.Fprintf(os.Stderr, "Flush: %v\n", err)
		os.Exit(1)
	}

	// Spot-check: print same values as the Python script
	fmt.Println("\n=== Go dataset values (spot check) ===")
	type testPoint struct {
		varName  string
		varIdx   int
		levelIdx int
		lat, lon int
	}

	points := []testPoint{
		{"HGT", 0, 0, 0, 0},       // HGT @ 1000mb, lat=-90, lon=0
		{"HGT", 0, 0, 180, 0},     // HGT @ 1000mb, lat=0, lon=0
		{"HGT", 0, 0, 360, 0},     // HGT @ 1000mb, lat=+90, lon=0
		{"HGT", 0, 20, 180, 360},  // HGT @ 500mb, lat=0, lon=180
		{"UGRD", 1, 0, 180, 0},    // UGRD @ 1000mb, lat=0, lon=0
		{"VGRD", 2, 0, 180, 0},    // VGRD @ 1000mb, lat=0, lon=0
		{"UGRD", 1, 20, 284, 0},   // UGRD @ 500mb, lat=52N, lon=0
	}

	for _, p := range points {
		val := ds.Val(0, p.levelIdx, p.varIdx, p.lat, p.lon)
		actualLat := -90.0 + float64(p.lat)*0.5
		actualLon := float64(p.lon) * 0.5
		fmt.Printf("  %-4s %4dmb lat=%+7.1f lon=%6.1f: %12.4f\n",
			p.varName, dataset.Pressures[p.levelIdx], actualLat, actualLon, val)
	}

	fmt.Printf("\nDataset written to %s\n", outPath)
}