weatherreporter/internal/promptinput/package.go

// Package promptinput builds prompt data packages from module snapshots.
package promptinput

import (
	"bytes"
	"encoding/json"
	"fmt"
	"time"

	"gitea.maximumdirect.net/eric/weatherreporter/internal/changes"
	"gitea.maximumdirect.net/eric/weatherreporter/internal/fileutil"
	"gitea.maximumdirect.net/eric/weatherreporter/internal/module"
	"gitea.maximumdirect.net/eric/weatherreporter/internal/report"
	"gitea.maximumdirect.net/eric/weatherreporter/internal/timeutil"
	"gitea.maximumdirect.net/eric/weatherreporter/internal/weatherdata"
	"gopkg.in/yaml.v3"
)

const SchemaVersion = "weatherreporter.data_package.v3"

const (
	metadataStanza                 = "metadata"
	categoryApplicableRiskProducts = "applicable_risk_products"
	categoryDerivedSummaries       = "derived_summaries"
	categoryNarrativeProducts      = "narrative_products"
	categoryRawData                = "raw_data"
)

var briefingCategoryOrder = []string{
	categoryApplicableRiskProducts,
	categoryDerivedSummaries,
	categoryNarrativeProducts,
	categoryRawData,
}

var briefingStanzaCategories = map[string]string{
	string(module.AlertDigest):             categoryApplicableRiskProducts,
	string(module.SPCConvectiveOutlooks):   categoryApplicableRiskProducts,
	string(module.DerivedDailySummary):     categoryDerivedSummaries,
	string(module.DerivedDaypartSummaries): categoryDerivedSummaries,
	string(module.PrecipTiming):            categoryDerivedSummaries,
	string(module.OutdoorWindows):          categoryDerivedSummaries,
	string(module.DailyPlanning):           categoryDerivedSummaries,
	string(module.TodayPlanning):           categoryDerivedSummaries,
	string(module.TomorrowPlanning):        categoryDerivedSummaries,
	string(module.NarrativeForecast):       categoryNarrativeProducts,
	string(module.AreaForecastDiscussion):  categoryNarrativeProducts,
	string(module.SPCConvectiveDiscussion): categoryNarrativeProducts,
	string(module.WeatherStory):            categoryNarrativeProducts,
	string(module.CurrentConditions):       categoryRawData,
	string(module.HourlyForecast):          categoryRawData,
}

type BuildRequest struct {
	Metadata      Metadata
	Modules       module.Snapshot
	RecentChanges []changes.Change
}

type Metadata struct {
	RunID          string
	ReportID       report.ID
	Variant        string
	PromptID       string
	GeneratedAt    time.Time
	Timezone       string
	ValidPeriod    timeutil.Period
	SourceWarnings []weatherdata.SourceWarning
}

type Package struct {
	SchemaVersion  string                      `json:"schemaVersion" yaml:"schema_version"`
	RunID          string                      `json:"runId" yaml:"run_id"`
	Report         Report                      `json:"report" yaml:"report"`
	Briefing       BriefingStanzas             `json:"briefing" yaml:"briefing"`
	RecentChanges  RecentChanges               `json:"recentChanges" yaml:"recent_changes"`
	SourceWarnings []weatherdata.SourceWarning `json:"sourceWarnings,omitempty" yaml:"source_warnings,omitempty"`
}

type Report struct {
	ID               report.ID       `json:"id" yaml:"id"`
	Variant          string          `json:"variant,omitempty" yaml:"variant,omitempty"`
	PromptID         string          `json:"promptId" yaml:"prompt_id"`
	GeneratedAt      time.Time       `json:"generatedAt" yaml:"generated_at"`
	Timezone         string          `json:"timezone" yaml:"timezone"`
	CurrentLocalDate string          `json:"currentLocalDate" yaml:"current_local_date"`
	ValidPeriod      timeutil.Period `json:"validPeriod" yaml:"valid_period"`
}

type BriefingStanzas struct {
	Order  []string       `json:"-" yaml:"-"`
	Values map[string]any `json:"-" yaml:"-"`
}

type RecentChanges struct {
	Items []changes.Change `json:"items" yaml:"items"`
}

func Build(req BuildRequest) (Package, error) {
	localDate, err := currentLocalDate(req.Metadata.GeneratedAt, req.Metadata.Timezone)
	if err != nil {
		return Package{}, err
	}
	items := make([]changes.Change, len(req.RecentChanges))
	copy(items, req.RecentChanges)
	if items == nil {
		items = []changes.Change{}
	}
	pkg := Package{
		SchemaVersion: SchemaVersion,
		RunID:         req.Metadata.RunID,
		Report: Report{
			ID:               req.Metadata.ReportID,
			Variant:          req.Metadata.Variant,
			PromptID:         req.Metadata.PromptID,
			GeneratedAt:      req.Metadata.GeneratedAt,
			Timezone:         req.Metadata.Timezone,
			CurrentLocalDate: localDate,
			ValidPeriod:      req.Metadata.ValidPeriod,
		},
		Briefing:       stanzasFromSnapshot(req.Modules),
		RecentChanges:  RecentChanges{Items: items},
		SourceWarnings: append([]weatherdata.SourceWarning(nil), req.Metadata.SourceWarnings...),
	}
	if err := Validate(pkg); err != nil {
		return Package{}, err
	}
	return pkg, nil
}

func stanzasFromSnapshot(snapshot module.Snapshot) BriefingStanzas {
	values := map[string]any{}
	order := make([]string, 0, len(snapshot.Outputs))
	for _, output := range snapshot.Outputs {
		order = append(order, output.StanzaName)
		values[output.StanzaName] = output.DataPackageValue()
	}
	return BriefingStanzas{Order: order, Values: values}
}

func currentLocalDate(generatedAt time.Time, timezone string) (string, error) {
	location, err := timeutil.LoadLocation(timezone)
	if err != nil {
		return "", fmt.Errorf("load report timezone %q: %w", timezone, err)
	}
	return generatedAt.In(location).Format(timeutil.DateLayout), nil
}

func Validate(pkg Package) error {
	if pkg.SchemaVersion == "" {
		return fmt.Errorf("schemaVersion is required")
	}
	if pkg.SchemaVersion != SchemaVersion {
		return fmt.Errorf("schemaVersion must be %s", SchemaVersion)
	}
	if pkg.RunID == "" {
		return fmt.Errorf("runId is required")
	}
	if pkg.Report.ID == "" {
		return fmt.Errorf("report.id is required")
	}
	if pkg.Report.PromptID == "" {
		return fmt.Errorf("report.promptId is required")
	}
	if pkg.Report.GeneratedAt.IsZero() {
		return fmt.Errorf("report.generatedAt is required")
	}
	if pkg.Report.Timezone == "" {
		return fmt.Errorf("report.timezone is required")
	}
	if pkg.Report.CurrentLocalDate == "" {
		return fmt.Errorf("report.currentLocalDate is required")
	}
	if !pkg.Report.ValidPeriod.IsValid() {
		return fmt.Errorf("report.validPeriod must be valid")
	}
	if len(pkg.Briefing.Order) == 0 {
		return fmt.Errorf("briefing stanzas are required")
	}
	seen := map[string]struct{}{}
	for _, name := range pkg.Briefing.Order {
		if name == "" {
			return fmt.Errorf("briefing stanza name is required")
		}
		if _, ok := seen[name]; ok {
			return fmt.Errorf("duplicate briefing stanza %q", name)
		}
		seen[name] = struct{}{}
		if _, ok := pkg.Briefing.Values[name]; !ok {
			return fmt.Errorf("briefing stanza %q is missing", name)
		}
		if name == metadataStanza {
			continue
		}
		if _, ok := briefingStanzaCategories[name]; !ok {
			return fmt.Errorf("briefing stanza %q has no prompt-input category", name)
		}
	}
	return nil
}

func Save(path string, pkg Package) error {
	data, err := MarshalYAML(pkg)
	if err != nil {
		return err
	}
	if err := fileutil.WriteFileAtomic(path, data); err != nil {
		return fmt.Errorf("save data package: %w", err)
	}
	return nil
}

func MarshalYAML(pkg Package) ([]byte, error) {
	if err := Validate(pkg); err != nil {
		return nil, err
	}
	data, err := yaml.Marshal(pkg)
	if err != nil {
		return nil, fmt.Errorf("marshal data package: %w", err)
	}
	return data, nil
}

func LoadYAML(data []byte) (Package, error) {
	var pkg Package
	decoder := yaml.NewDecoder(bytes.NewReader(data))
	if err := decoder.Decode(&pkg); err != nil {
		return Package{}, fmt.Errorf("decode data package: %w", err)
	}
	if err := Validate(pkg); err != nil {
		return Package{}, err
	}
	return pkg, nil
}

func (b BriefingStanzas) MarshalYAML() (any, error) {
	node := &yaml.Node{Kind: yaml.MappingNode}
	categoryNames := map[string][]string{}
	for _, name := range b.Order {
		value, ok := b.Values[name]
		if !ok {
			continue
		}
		if name == metadataStanza {
			if err := appendYAMLMappingValue(node, name, value); err != nil {
				return nil, err
			}
			continue
		}
		category, ok := briefingStanzaCategories[name]
		if !ok {
			return nil, fmt.Errorf("briefing stanza %q has no prompt-input category", name)
		}
		categoryNames[category] = append(categoryNames[category], name)
	}
	for _, category := range briefingCategoryOrder {
		names := categoryNames[category]
		if len(names) == 0 {
			continue
		}
		categoryNode := &yaml.Node{Kind: yaml.MappingNode}
		for _, name := range names {
			value := b.Values[name]
			if err := appendYAMLMappingValue(categoryNode, name, value); err != nil {
				return nil, err
			}
		}
		node.Content = append(node.Content,
			&yaml.Node{Kind: yaml.ScalarNode, Value: category},
			categoryNode,
		)
	}
	return node, nil
}

func (b *BriefingStanzas) UnmarshalYAML(value *yaml.Node) error {
	if value.Kind != yaml.MappingNode {
		return fmt.Errorf("briefing must be a mapping")
	}
	values := map[string]any{}
	order := make([]string, 0, len(value.Content)/2)
	seen := map[string]struct{}{}
	seenCategories := map[string]struct{}{}
	categoryOrder := map[string][]string{}
	for i := 0; i < len(value.Content); i += 2 {
		name := value.Content[i].Value
		if name == metadataStanza {
			if err := decodeBriefingStanza(value.Content[i+1], name, values, &order, seen); err != nil {
				return err
			}
			continue
		}
		if !knownBriefingCategory(name) {
			return fmt.Errorf("unknown briefing category %q", name)
		}
		if _, ok := seenCategories[name]; ok {
			return fmt.Errorf("duplicate briefing category %q", name)
		}
		seenCategories[name] = struct{}{}
		categoryNode := value.Content[i+1]
		if categoryNode.Kind != yaml.MappingNode {
			return fmt.Errorf("briefing category %q must be a mapping", name)
		}
		var names []string
		for j := 0; j < len(categoryNode.Content); j += 2 {
			stanzaName := categoryNode.Content[j].Value
			category, ok := briefingStanzaCategories[stanzaName]
			if !ok {
				return fmt.Errorf("briefing stanza %q has no prompt-input category", stanzaName)
			}
			if category != name {
				return fmt.Errorf("briefing stanza %q belongs under category %q, not %q", stanzaName, category, name)
			}
			if err := decodeBriefingStanza(categoryNode.Content[j+1], stanzaName, values, &names, seen); err != nil {
				return err
			}
		}
		categoryOrder[name] = names
	}
	for _, category := range briefingCategoryOrder {
		order = append(order, categoryOrder[category]...)
	}
	b.Order = order
	b.Values = values
	return nil
}

func (b BriefingStanzas) MarshalJSON() ([]byte, error) {
	out := map[string]any{}
	for _, name := range b.Order {
		if value, ok := b.Values[name]; ok {
			out[name] = value
		}
	}
	return json.Marshal(out)
}

func (b *BriefingStanzas) UnmarshalJSON(data []byte) error {
	var values map[string]any
	if err := json.Unmarshal(data, &values); err != nil {
		return err
	}
	order := make([]string, 0, len(values))
	for name := range values {
		order = append(order, name)
	}
	b.Order = order
	b.Values = values
	return nil
}

func appendYAMLMappingValue(node *yaml.Node, name string, value any) error {
	keyNode := &yaml.Node{Kind: yaml.ScalarNode, Value: name}
	valueNode, err := yamlNode(value)
	if err != nil {
		return fmt.Errorf("marshal briefing stanza %q: %w", name, err)
	}
	node.Content = append(node.Content, keyNode, valueNode)
	return nil
}

func decodeBriefingStanza(node *yaml.Node, name string, values map[string]any, order *[]string, seen map[string]struct{}) error {
	if _, ok := seen[name]; ok {
		return fmt.Errorf("duplicate briefing stanza %q", name)
	}
	var stanza any
	if err := node.Decode(&stanza); err != nil {
		return err
	}
	seen[name] = struct{}{}
	*order = append(*order, name)
	values[name] = stanza
	return nil
}

func knownBriefingCategory(name string) bool {
	for _, category := range briefingCategoryOrder {
		if name == category {
			return true
		}
	}
	return false
}

func yamlNode(value any) (*yaml.Node, error) {
	data, err := json.Marshal(value)
	if err != nil {
		return nil, err
	}
	var normalized any
	if err := json.Unmarshal(data, &normalized); err != nil {
		return nil, err
	}
	data, err = yaml.Marshal(normalized)
	if err != nil {
		return nil, err
	}
	var node yaml.Node
	if err := yaml.Unmarshal(data, &node); err != nil {
		return nil, err
	}
	if len(node.Content) == 0 {
		return &yaml.Node{Kind: yaml.MappingNode}, nil
	}
	return node.Content[0], nil
}