package dataset

import (
	"encoding/binary"
	"fmt"
	"math"
	"os"
	"time"

	mmap "github.com/edsrzf/mmap-go"
)

// File represents an mmap-backed wind dataset file.
type File struct {
	mm       mmap.MMap
	file     *os.File
	writable bool
	DSTime   time.Time // forecast run time (UTC)
}

// Open opens an existing dataset file for reading.
func Open(path string, dsTime time.Time) (*File, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open dataset: %w", err)
	}

	info, err := f.Stat()
	if err != nil {
		f.Close()
		return nil, fmt.Errorf("stat dataset: %w", err)
	}
	if info.Size() != DatasetSize {
		f.Close()
		return nil, fmt.Errorf("dataset should be %d bytes (was %d)", DatasetSize, info.Size())
	}

	mm, err := mmap.Map(f, mmap.RDONLY, 0)
	if err != nil {
		f.Close()
		return nil, fmt.Errorf("mmap dataset: %w", err)
	}

	return &File{mm: mm, file: f, writable: false, DSTime: dsTime}, nil
}

// Create creates a new dataset file for writing.
// The file is truncated to the correct size and mmap'd read-write.
func Create(path string) (*File, error) {
	f, err := os.Create(path)
	if err != nil {
		return nil, fmt.Errorf("create dataset: %w", err)
	}

	if err := f.Truncate(DatasetSize); err != nil {
		f.Close()
		return nil, fmt.Errorf("truncate dataset: %w", err)
	}

	mm, err := mmap.MapRegion(f, int(DatasetSize), mmap.RDWR, 0, 0)
	if err != nil {
		f.Close()
		return nil, fmt.Errorf("mmap dataset: %w", err)
	}

	return &File{mm: mm, file: f, writable: true}, nil
}

// offset computes the byte offset for element [hour][level][variable][lat][lon].
// Row-major C-order indexing matching the reference implementation:
// shape = (65, 47, 3, 361, 720)
func offset(hour, level, variable, lat, lon int) int64 {
	idx := int64(hour)
	idx = idx*int64(NumLevels) + int64(level)
	idx = idx*int64(NumVariables) + int64(variable)
	idx = idx*int64(NumLatitudes) + int64(lat)
	idx = idx*int64(NumLongitudes) + int64(lon)
	return idx * int64(ElementSize)
}

// Val reads a float32 value from the dataset at [hour][level][variable][lat][lon].
func (d *File) Val(hour, level, variable, lat, lon int) float32 {
	off := offset(hour, level, variable, lat, lon)
	bits := binary.LittleEndian.Uint32(d.mm[off : off+4])
	return math.Float32frombits(bits)
}

// SetVal writes a float32 value to the dataset at [hour][level][variable][lat][lon].
// Only valid on writable (created) datasets.
func (d *File) SetVal(hour, level, variable, lat, lon int, val float32) {
	off := offset(hour, level, variable, lat, lon)
	binary.LittleEndian.PutUint32(d.mm[off:off+4], math.Float32bits(val))
}

// BlitGribData copies decoded GRIB grid data into the dataset at the given position.
// gribData is 361*720 float64 values in GRIB scan order (north-to-south, west-to-east).
// This function flips the latitude so that dataset index 0 = -90 (south) and 360 = +90 (north).
func (d *File) BlitGribData(hourIdx, levelIdx, varIdx int, gribData []float64) error {
	expected := NumLatitudes * NumLongitudes
	if len(gribData) != expected {
		return fmt.Errorf("grib data has %d values, expected %d", len(gribData), expected)
	}

	for lat := 0; lat < NumLatitudes; lat++ {
		for lon := 0; lon < NumLongitudes; lon++ {
			// GRIB scans north-to-south: row 0 = 90N, row 360 = 90S
			// Dataset stores south-to-north: index 0 = -90 (90S), index 360 = +90 (90N)
			gribIdx := (360-lat)*NumLongitudes + lon
			val := float32(gribData[gribIdx])
			d.SetVal(hourIdx, levelIdx, varIdx, lat, lon, val)
		}
	}

	return nil
}

// Flush flushes the mmap to disk.
func (d *File) Flush() error {
	if d.mm != nil {
		return d.mm.Flush()
	}
	return nil
}

// Close unmaps and closes the dataset file.
func (d *File) Close() error {
	if d.mm != nil {
		if err := d.mm.Unmap(); err != nil {
			d.file.Close()
			return fmt.Errorf("unmap: %w", err)
		}
		d.mm = nil
	}
	if d.file != nil {
		err := d.file.Close()
		d.file = nil
		return err
	}
	return nil
}