package vmquery

import (
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
)

// QueryResponse is the top-level PromQL/VictoriaMetrics API envelope.
type QueryResponse struct {
	Status    string    `json:"status"`
	ErrorType string    `json:"errorType,omitempty"`
	Error     string    `json:"error,omitempty"`
	Data      QueryData `json:"data"`
}

// Err returns nil when Status == "success", otherwise an error carrying the
// server-supplied errorType/error message. Callers should consult this before
// reading Data.
func (r *QueryResponse) Err() error {
	if r == nil {
		return errors.New("vmquery: nil response")
	}
	if r.Status == "success" {
		return nil
	}
	if r.Error != "" {
		return errors.New(r.Error)
	}
	if r.ErrorType != "" {
		return errors.New(r.ErrorType)
	}
	return errors.New("vmquery: query failed")
}

// Count sums every sample value in the response. Returns 0 for empty results.
func (r *QueryResponse) Count() (float64, error) {
	values, err := r.Values()
	if err != nil {
		return 0, err
	}
	var sum float64
	for _, v := range values {
		sum += v
	}
	return sum, nil
}

// Values flattens every sample value in the response, regardless of
// resultType shape. matrix → values across all series; vector → value of
// each series; scalar/string → the single sample's value.
func (r *QueryResponse) Values() ([]float64, error) {
	if err := r.Err(); err != nil {
		return nil, err
	}
	return r.Data.Values()
}

// QueryData holds the result payload. Result is RawMessage because the shape
// varies by ResultType (matrix/vector vs scalar/string).
type QueryData struct {
	ResultType string          `json:"resultType"`
	Result     json.RawMessage `json:"result"`
}

// Values decodes Result according to ResultType and returns every sample
// value as a flat slice. Returns an empty slice for empty results.
func (d QueryData) Values() ([]float64, error) {
	switch d.ResultType {
	case "matrix":
		series, err := decodeMatrix(d.Result)
		if err != nil {
			return nil, err
		}
		out := []float64{}
		for _, s := range series {
			for _, sample := range s.Values {
				v, err := sample.Value()
				if err != nil {
					return nil, err
				}
				out = append(out, v)
			}
		}
		return out, nil
	case "vector":
		series, err := decodeMatrix(d.Result)
		if err != nil {
			return nil, err
		}
		out := make([]float64, 0, len(series))
		for _, s := range series {
			if s.Value == nil {
				continue
			}
			v, err := s.Value.Value()
			if err != nil {
				return nil, err
			}
			out = append(out, v)
		}
		return out, nil
	case "scalar", "string":
		sample, err := decodeScalar(d.Result)
		if err != nil {
			return nil, err
		}
		v, err := sample.Value()
		if err != nil {
			return nil, err
		}
		return []float64{v}, nil
	}
	// Unknown resultType: treat as empty so callers don't fail on a successful
	// query with an unfamiliar shape. Empty Result also falls through here.
	return []float64{}, nil
}

// Series is one row of a matrix or vector response.
type Series struct {
	Metric map[string]string `json:"metric,omitempty"`
	Values []Sample          `json:"values,omitempty"` // matrix
	Value  *Sample           `json:"value,omitempty"`  // vector
}

// Sample is a Prometheus [timestamp, value] pair. The value slot is a quoted
// string in PromQL responses.
type Sample [2]json.RawMessage

// Value parses the value slot of the sample as a float64. Handles both the
// PromQL quoted-string form ("7") and a bare numeric form (7).
func (s Sample) Value() (float64, error) {
	return parseSampleValue(s[1])
}

func parseSampleValue(b json.RawMessage) (float64, error) {
	// PromQL responses quote the value. Try string first.
	var s string
	if err := json.Unmarshal(b, &s); err == nil {
		f, perr := strconv.ParseFloat(s, 64)
		if perr != nil {
			return 0, fmt.Errorf("vmquery: parse sample value %q: %w", s, perr)
		}
		return f, nil
	}
	// Fall back to bare numeric JSON (defensive).
	var f float64
	if err := json.Unmarshal(b, &f); err != nil {
		return 0, fmt.Errorf("vmquery: invalid sample value %s: %w", b, err)
	}
	return f, nil
}

func decodeMatrix(b json.RawMessage) ([]Series, error) {
	if len(b) == 0 || string(b) == "null" {
		return nil, nil
	}
	var out []Series
	if err := json.Unmarshal(b, &out); err != nil {
		return nil, fmt.Errorf("vmquery: decode series: %w", err)
	}
	return out, nil
}

func decodeScalar(b json.RawMessage) (Sample, error) {
	var s Sample
	if err := json.Unmarshal(b, &s); err != nil {
		return Sample{}, fmt.Errorf("vmquery: decode scalar: %w", err)
	}
	return s, nil
}