predictor/internal/downloader/idx.go

package downloader

import (
	"fmt"
	"strconv"
	"strings"
)

// IdxEntry represents a single parsed line from a GRIB .idx file.
// Example line: "15:1207405:d=2024010100:HGT:1000 mb:0 hour fcst:"
type IdxEntry struct {
	Index     int
	Offset    int64
	Variable  string // "HGT", "UGRD", "VGRD", etc.
	LevelMB   int    // pressure level in mb (0 if not a pressure level)
	Hour      int    // forecast hour
	EndOffset int64  // byte after this message (from next entry's offset, or -1 if last)
}

// Length returns the byte length of this GRIB message, or -1 if unknown.
func (e *IdxEntry) Length() int64 {
	if e.EndOffset <= 0 {
		return -1
	}
	return e.EndOffset - e.Offset
}

// ParseIdx parses a .idx file body and returns all entries.
// Lines that can't be parsed are silently skipped.
func ParseIdx(body []byte) []IdxEntry {
	lines := strings.Split(string(body), "\n")
	var entries []IdxEntry

	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}

		parts := strings.Split(line, ":")
		if len(parts) < 7 {
			continue
		}

		idx, err := strconv.Atoi(parts[0])
		if err != nil {
			continue
		}

		offset, err := strconv.ParseInt(parts[1], 10, 64)
		if err != nil {
			continue
		}

		variable := parts[3]
		levelStr := parts[4]
		hourStr := parts[5]

		levelMB := parseLevelMB(levelStr)
		hour := parseHour(hourStr)

		entries = append(entries, IdxEntry{
			Index:     idx,
			Offset:    offset,
			Variable:  variable,
			LevelMB:   levelMB,
			Hour:      hour,
			EndOffset: -1, // filled in below
		})
	}

	// Fill in EndOffset from the next entry's Offset.
	for i := 0; i < len(entries)-1; i++ {
		entries[i].EndOffset = entries[i+1].Offset
	}

	return entries
}

// FilterIdx returns entries matching the given variables at pressure levels.
// Only entries with a recognized pressure level (levelMB > 0) are returned.
func FilterIdx(entries []IdxEntry, variables map[string]bool) []IdxEntry {
	var filtered []IdxEntry
	for _, e := range entries {
		if !variables[e.Variable] {
			continue
		}
		if e.LevelMB <= 0 {
			continue
		}
		// Must have a known length (not the last entry) or be handled specially
		if e.Length() <= 0 {
			continue
		}
		filtered = append(filtered, e)
	}
	return filtered
}

// parseLevelMB parses a level string like "1000 mb" and returns the pressure in mb.
// Returns 0 if not a pressure level.
func parseLevelMB(s string) int {
	s = strings.TrimSpace(s)
	if !strings.HasSuffix(s, " mb") {
		return 0
	}
	numStr := strings.TrimSuffix(s, " mb")
	n, err := strconv.Atoi(numStr)
	if err != nil {
		return 0
	}
	return n
}

// parseHour parses a forecast hour string like "0 hour fcst" or "anl".
// Returns -1 if it can't be parsed.
func parseHour(s string) int {
	s = strings.TrimSpace(s)
	if s == "anl" {
		return 0
	}
	s = strings.TrimSuffix(s, " hour fcst")
	n, err := strconv.Atoi(s)
	if err != nil {
		return -1
	}
	return n
}

// GroupByRange groups idx entries into byte ranges suitable for HTTP Range downloads.
// Each range covers one contiguous GRIB message.
type ByteRange struct {
	Start int64
	End   int64 // inclusive
	Entry IdxEntry
}

// EntriesToRanges converts filtered idx entries to byte ranges.
func EntriesToRanges(entries []IdxEntry) []ByteRange {
	ranges := make([]ByteRange, 0, len(entries))
	for _, e := range entries {
		if e.Length() <= 0 {
			continue
		}
		ranges = append(ranges, ByteRange{
			Start: e.Offset,
			End:   e.EndOffset - 1, // inclusive
			Entry: e,
		})
	}
	return ranges
}

// FormatRange returns an HTTP Range header value for a byte range.
func (r ByteRange) FormatRange() string {
	return fmt.Sprintf("bytes=%d-%d", r.Start, r.End)
}