// This file defines, fetches, validates, and loads continuous historical replay datasets.
package evaluation

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"sigs.k8s.io/yaml"
)

const maxReplayDatasetSpecBytes = 1 << 20

type ReplayDatasetSpec struct {
	Version          string                      `json:"version"`
	ID               string                      `json:"id"`
	Component        string                      `json:"component"`
	Profile          string                      `json:"profile"`
	Target           O11YTargetSpec              `json:"target"`
	Window           O11YQueryWindow             `json:"window"`
	Endpoints        map[O11YMetricSource]string `json:"endpoints"`
	ReferenceSources []O11YSourceSpec            `json:"referenceSources,omitempty"`
}

type O11YDatasetExportSpec struct {
	Version   string           `json:"version"`
	ID        string           `json:"id"`
	Component string           `json:"component"`
	Target    O11YTargetSpec   `json:"target"`
	Window    O11YQueryWindow  `json:"window"`
	Sources   []O11YSourceSpec `json:"sources"`
}

type O11YDatasetExportOptions struct {
	Client             *http.Client
	BearerToken        string
	RequireBearerToken bool
	ArchivePath        string
	CacheDir           string
}

type O11YDatasetArchive struct {
	Version   string              `json:"version"`
	DatasetID string              `json:"datasetId"`
	Component string              `json:"component"`
	Target    O11YTargetSpec      `json:"target"`
	Window    O11YQueryWindow     `json:"window"`
	Queries   []O11YArchivedQuery `json:"queries"`
}

func DecodeReplayDatasetSpec(r io.Reader) (*ReplayDatasetSpec, error) {
	data, err := io.ReadAll(io.LimitReader(r, maxReplayDatasetSpecBytes+1))
	if err != nil {
		return nil, fmt.Errorf("evaluation: read replay dataset spec: %w", err)
	}
	if len(data) > maxReplayDatasetSpecBytes {
		return nil, fmt.Errorf("evaluation: replay dataset spec exceeds %d bytes", maxReplayDatasetSpecBytes)
	}

	var spec ReplayDatasetSpec
	if err := yaml.UnmarshalStrict(data, &spec); err != nil {
		return nil, fmt.Errorf("evaluation: decode replay dataset spec: %w", err)
	}
	if err := validateReplayDatasetSpec(&spec); err != nil {
		return nil, err
	}
	return &spec, nil
}

func BuildO11YDatasetExportSpec(spec *ReplayDatasetSpec, profile *MetricsProfile) (*O11YDatasetExportSpec, error) {
	if err := validateReplayDatasetSpec(spec); err != nil {
		return nil, err
	}
	if profile == nil {
		return nil, fmt.Errorf("evaluation: metrics profile is required")
	}
	if err := validateMetricsProfile(profile); err != nil {
		return nil, err
	}
	if profile.ID != spec.Profile || profile.Component != spec.Component {
		return nil, fmt.Errorf("evaluation: metrics profile %q does not match replay dataset", profile.ID)
	}

	sources := make([]O11YSourceSpec, 0, len(profile.Sources)+len(spec.ReferenceSources))
	usedEndpoints := make(map[O11YMetricSource]struct{}, len(profile.Sources))
	for i := range profile.Sources {
		profileSource := &profile.Sources[i]
		baseURL, ok := spec.Endpoints[profileSource.Type]
		if !ok || baseURL == "" {
			return nil, fmt.Errorf("evaluation: O11Y endpoint %q is required", profileSource.Type)
		}
		usedEndpoints[profileSource.Type] = struct{}{}

		queries := make([]O11YQuerySpec, len(profileSource.Queries))
		for j := range profileSource.Queries {
			if profileSource.Queries[j].Type != O11YQueryTypeRange {
				return nil, fmt.Errorf("evaluation: replay dataset query %q must be a range query", profileSource.Queries[j].Name)
			}
			queries[j] = profileSource.Queries[j]
			if profileSource.Queries[j].KeepLabels != nil {
				queries[j].KeepLabels = append([]string{}, profileSource.Queries[j].KeepLabels...)
			}
			queries[j].PromQL = strings.ReplaceAll(queries[j].PromQL, metricsProfileClusterIDPlaceholder, spec.Target.ClusterID)
		}
		sources = append(sources, O11YSourceSpec{Type: profileSource.Type, BaseURL: baseURL, Queries: queries})
	}
	if len(usedEndpoints) != len(spec.Endpoints) {
		return nil, fmt.Errorf("evaluation: O11Y endpoints must match metrics profile sources")
	}
	for i := range spec.ReferenceSources {
		sources = append(sources, cloneO11YSource(spec.ReferenceSources[i]))
	}

	exportSpec := &O11YDatasetExportSpec{
		Version:   ReplayVersionV1Alpha1,
		ID:        spec.ID,
		Component: spec.Component,
		Target:    spec.Target,
		Window:    spec.Window,
		Sources:   sources,
	}
	if err := validateO11YDatasetExportSpec(exportSpec); err != nil {
		return nil, err
	}
	return exportSpec, nil
}

func FetchO11YDataset(ctx context.Context, spec *O11YDatasetExportSpec, options O11YDatasetExportOptions) (*O11YDatasetArchive, error) {
	if err := validateO11YDatasetExportSpec(spec); err != nil {
		return nil, err
	}
	if options.ArchivePath == "" {
		return nil, fmt.Errorf("evaluation: O11Y dataset archive path is required")
	}
	if options.Client == nil {
		options.Client = &http.Client{Timeout: defaultO11YHTTPTimeout}
	}
	if err := makePrivateDir(filepath.Dir(options.ArchivePath)); err != nil {
		return nil, fmt.Errorf("evaluation: prepare O11Y dataset archive directory: %w", err)
	}
	if options.CacheDir != "" {
		if err := makePrivateDir(options.CacheDir); err != nil {
			return nil, fmt.Errorf("evaluation: prepare O11Y cache directory: %w", err)
		}
	}

	archive := &O11YDatasetArchive{
		Version:   ReplayVersionV1Alpha1,
		DatasetID: spec.ID,
		Component: spec.Component,
		Target:    spec.Target,
		Window:    spec.Window,
		Queries:   make([]O11YArchivedQuery, 0, queryCount(spec.Sources)),
	}
	for i := range spec.Sources {
		source := &spec.Sources[i]
		for j := range source.Queries {
			query := &source.Queries[j]
			requestURL, err := buildO11YQueryURL(source, query, &spec.Window, spec.Window.End)
			if err != nil {
				return nil, fmt.Errorf("evaluation: build O11Y dataset query %q: %w", query.Name, err)
			}
			record := O11YArchivedQuery{
				Name:          query.Name,
				Source:        source.Type,
				Type:          query.Type,
				PromQL:        query.PromQL,
				RequestDigest: o11yRequestDigestWithHeaders(requestURL, source.Headers),
			}

			body, cacheHit, err := readO11YCache(options.CacheDir, record.RequestDigest)
			if err != nil {
				return nil, fmt.Errorf("evaluation: read O11Y dataset query %q cache: %w", query.Name, err)
			}
			record.CacheHit = cacheHit
			if !cacheHit {
				if options.RequireBearerToken && options.BearerToken == "" {
					return nil, fmt.Errorf("evaluation: O11Y dataset query %q cache miss requires bearer token", query.Name)
				}
				body, err = fetchO11YQueryWithHeaders(ctx, options.Client, options.BearerToken, requestURL, source.Headers)
				if err != nil {
					return nil, fmt.Errorf("evaluation: fetch O11Y dataset query %q: %w", query.Name, err)
				}
			}

			response, err := decodePrometheusResponse(body)
			if err != nil {
				return nil, fmt.Errorf("evaluation: O11Y dataset query %q: %w", query.Name, err)
			}
			if response.IsPartial {
				return nil, fmt.Errorf("evaluation: O11Y dataset query %q returned a partial response", query.Name)
			}
			if !cacheHit && options.CacheDir != "" {
				if err := writeO11YCache(options.CacheDir, record.RequestDigest, body); err != nil {
					return nil, fmt.Errorf("evaluation: cache O11Y dataset query %q: %w", query.Name, err)
				}
			}

			series, sampleCount, err := normalizePrometheusResponse(response, query, &spec.Window, spec.Window.End)
			if err != nil {
				return nil, fmt.Errorf("evaluation: normalize O11Y dataset query %q: %w", query.Name, err)
			}
			record.SeriesCount = len(series)
			record.SampleCount = sampleCount
			record.Response = append(json.RawMessage(nil), body...)
			archive.Queries = append(archive.Queries, record)
		}
	}

	data, err := json.Marshal(archive)
	if err != nil {
		return nil, fmt.Errorf("evaluation: encode O11Y dataset archive: %w", err)
	}
	if err := writePrivateFile(options.ArchivePath, append(data, '\n')); err != nil {
		return nil, fmt.Errorf("evaluation: write O11Y dataset archive: %w", err)
	}
	return archive, nil
}

func DecodeO11YDatasetArchive(r io.Reader) (*O11YDatasetArchive, error) {
	decoder := json.NewDecoder(r)
	decoder.DisallowUnknownFields()

	var archive O11YDatasetArchive
	if err := decoder.Decode(&archive); err != nil {
		return nil, fmt.Errorf("evaluation: decode O11Y dataset archive: %w", err)
	}
	var extra any
	if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) {
		if err == nil {
			return nil, fmt.Errorf("evaluation: decode O11Y dataset archive: multiple JSON values")
		}
		return nil, fmt.Errorf("evaluation: decode O11Y dataset archive trailing data: %w", err)
	}
	return &archive, nil
}

func ReplayO11YDatasetArchive(spec *O11YDatasetExportSpec, archive *O11YDatasetArchive, archiveDigest string) (*ReplayDataset, error) {
	if err := validateO11YDatasetExportSpec(spec); err != nil {
		return nil, err
	}
	if err := validateO11YDatasetArchiveIdentity(spec, archive); err != nil {
		return nil, err
	}

	metrics := make([]MetricSeries, 0)
	maxAgeByMetric := make(map[string]time.Duration)
	queryIndex := 0
	for i := range spec.Sources {
		source := &spec.Sources[i]
		for j := range source.Queries {
			query := &source.Queries[j]
			record := &archive.Queries[queryIndex]
			queryIndex++
			if err := validateO11YArchivedQuery(source, query, &spec.Window, spec.Window.End, record); err != nil {
				return nil, err
			}
			response, err := decodePrometheusResponse(record.Response)
			if err != nil {
				return nil, fmt.Errorf("evaluation: replay O11Y dataset query %q: %w", query.Name, err)
			}
			if response.IsPartial {
				return nil, fmt.Errorf("evaluation: replay O11Y dataset query %q is partial", query.Name)
			}
			series, _, err := normalizePrometheusResponse(response, query, &spec.Window, spec.Window.End)
			if err != nil {
				return nil, fmt.Errorf("evaluation: replay O11Y dataset query %q: %w", query.Name, err)
			}
			metrics = append(metrics, series...)
			maxAgeByMetric[query.Name] = time.Duration(query.MaxGapSeconds) * time.Second
		}
	}
	sortMetricSeries(metrics)
	return &ReplayDataset{
		Version:        ReplayVersionV1Alpha1,
		ID:             spec.ID,
		Component:      spec.Component,
		Target:         spec.Target,
		Window:         spec.Window,
		ArchiveDigest:  archiveDigest,
		Metrics:        metrics,
		MaxAgeByMetric: maxAgeByMetric,
	}, nil
}

func LoadReplayDataset(datasetPath, casesDir string) (*ReplayDataset, error) {
	file, err := os.Open(datasetPath)
	if err != nil {
		return nil, fmt.Errorf("evaluation: open replay dataset spec %q: %w", datasetPath, err)
	}
	spec, decodeErr := DecodeReplayDatasetSpec(file)
	closeErr := file.Close()
	if decodeErr != nil {
		return nil, decodeErr
	}
	if closeErr != nil {
		return nil, fmt.Errorf("evaluation: close replay dataset spec %q: %w", datasetPath, closeErr)
	}

	profile, err := LoadMetricsProfile(casesDir, spec.Profile)
	if err != nil {
		return nil, err
	}
	exportSpec, err := BuildO11YDatasetExportSpec(spec, profile)
	if err != nil {
		return nil, err
	}
	archivePath := strings.TrimSuffix(datasetPath, filepath.Ext(datasetPath)) + ".json"
	data, err := os.ReadFile(archivePath)
	if err != nil {
		return nil, fmt.Errorf("evaluation: read replay dataset archive %q: %w", archivePath, err)
	}
	archive, err := DecodeO11YDatasetArchive(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	return ReplayO11YDatasetArchive(exportSpec, archive, digestBytes(data))
}

func validateReplayDatasetSpec(spec *ReplayDatasetSpec) error {
	if spec == nil {
		return fmt.Errorf("evaluation: replay dataset spec is required")
	}
	if spec.Version != ReplayVersionV1Alpha1 || spec.ID == "" || spec.Component == "" {
		return fmt.Errorf("evaluation: replay dataset version, id, and component are required")
	}
	if !isMetricsProfileReference(spec.Profile) {
		return fmt.Errorf("evaluation: replay dataset profile %q is invalid", spec.Profile)
	}
	if spec.Target.ClusterID == "" || spec.Target.Vendor == "" || spec.Target.Region == "" {
		return fmt.Errorf("evaluation: replay dataset target clusterId, vendor, and region are required")
	}
	if spec.Window.Start.IsZero() || spec.Window.End.IsZero() || !spec.Window.Start.Before(spec.Window.End) {
		return fmt.Errorf("evaluation: replay dataset window start must be before end")
	}
	if spec.Window.StepSeconds < 1 {
		return fmt.Errorf("evaluation: replay dataset stepSeconds must be greater than zero")
	}
	if len(spec.Endpoints) == 0 {
		return fmt.Errorf("evaluation: replay dataset endpoints are required")
	}
	return nil
}

func cloneO11YSource(source O11YSourceSpec) O11YSourceSpec {
	clone := O11YSourceSpec{
		Type:    source.Type,
		BaseURL: source.BaseURL,
		Queries: make([]O11YQuerySpec, len(source.Queries)),
	}
	for name, value := range source.Headers {
		if clone.Headers == nil {
			clone.Headers = make(map[string]string, len(source.Headers))
		}
		clone.Headers[name] = value
	}
	copy(clone.Queries, source.Queries)
	for i := range clone.Queries {
		if source.Queries[i].KeepLabels != nil {
			clone.Queries[i].KeepLabels = append([]string{}, source.Queries[i].KeepLabels...)
		}
	}
	return clone
}

func validateO11YDatasetExportSpec(spec *O11YDatasetExportSpec) error {
	if spec == nil {
		return fmt.Errorf("evaluation: O11Y dataset export spec is required")
	}
	if spec.Version != ReplayVersionV1Alpha1 || spec.ID == "" || spec.Component == "" {
		return fmt.Errorf("evaluation: O11Y dataset version, id, and component are required")
	}
	if spec.Target.ClusterID == "" || spec.Target.Vendor == "" || spec.Target.Region == "" {
		return fmt.Errorf("evaluation: O11Y dataset target is required")
	}
	if spec.Window.Start.IsZero() || spec.Window.End.IsZero() || !spec.Window.Start.Before(spec.Window.End) || spec.Window.StepSeconds < 1 {
		return fmt.Errorf("evaluation: O11Y dataset window is invalid")
	}
	if len(spec.Sources) == 0 {
		return fmt.Errorf("evaluation: O11Y dataset sources are required")
	}

	queryNames := make(map[string]struct{})
	for i := range spec.Sources {
		source := &spec.Sources[i]
		if source.Type != O11YMetricSourceApplication && source.Type != O11YMetricSourceInfrastructure {
			return fmt.Errorf("evaluation: O11Y dataset source %q is unsupported", source.Type)
		}
		if source.BaseURL == "" {
			return fmt.Errorf("evaluation: O11Y dataset source %q baseUrl is required", source.Type)
		}
		for name, value := range source.Headers {
			if name == "" || value == "" {
				return fmt.Errorf("evaluation: O11Y dataset source %q headers must have names and values", source.Type)
			}
		}
		if err := validateO11YQueries(i, source.Queries, queryNames); err != nil {
			return err
		}
		for j := range source.Queries {
			if source.Queries[j].Type != O11YQueryTypeRange {
				return fmt.Errorf("evaluation: O11Y dataset query %q must be a range query", source.Queries[j].Name)
			}
		}
	}
	return nil
}

func validateO11YDatasetArchiveIdentity(spec *O11YDatasetExportSpec, archive *O11YDatasetArchive) error {
	if archive == nil {
		return fmt.Errorf("evaluation: O11Y dataset archive is required")
	}
	if archive.Version != ReplayVersionV1Alpha1 || archive.DatasetID != spec.ID || archive.Component != spec.Component {
		return fmt.Errorf("evaluation: O11Y dataset archive identity does not match %q", spec.ID)
	}
	if archive.Target != spec.Target || !archive.Window.Start.Equal(spec.Window.Start) || !archive.Window.End.Equal(spec.Window.End) || archive.Window.StepSeconds != spec.Window.StepSeconds {
		return fmt.Errorf("evaluation: O11Y dataset archive target or window does not match %q", spec.ID)
	}
	if len(archive.Queries) != queryCount(spec.Sources) {
		return fmt.Errorf("evaluation: O11Y dataset archive queries = %d, want %d", len(archive.Queries), queryCount(spec.Sources))
	}
	return nil
}

func digestBytes(data []byte) string {
	digest := sha256.Sum256(data)
	return fmt.Sprintf("%x", digest)
}