1package testdatasource
2
3import (
4	"context"
5	"encoding/json"
6	"fmt"
7	"math"
8	"time"
9
10	"github.com/grafana/grafana-plugin-sdk-go/backend"
11	"github.com/grafana/grafana-plugin-sdk-go/data"
12)
13
14var modeValueAsRow = "values-as-rows"
15var modeValueAsFields = "values-as-fields"
16var modeValueAsLabeledFields = "values-as-labeled-fields"
17var modeTimeseries = "timeseries"
18var modeTimeseriesWide = "timeseries-wide"
19
20var allStateCodes = []string{"AL", "AK", "AZ", "AR", "CA", "CO", "CT", "DE", "DC", "FL", "GA", "HI", "ID", "IL", "IN", "IA", "KS", "KY", "LA", "ME", "MT", "NE", "NV", "NH", "NJ", "NM", "NY", "NC", "ND", "OH", "OK", "OR", "MD", "MA", "MI", "MN", "MS", "MO", "PA", "RI", "SC", "SD", "TN", "TX", "UT", "VT", "VA", "WA", "WV", "WI", "WY"}
21
22type usaQueryWrapper struct {
23	USA usaQuery `json:"usa"`
24}
25
26type usaQuery struct {
27	Mode   string   `json:"mode"`
28	Period string   `json:"period"`
29	Fields []string `json:"fields"`
30	States []string `json:"states"`
31
32	// From the main query
33	maxDataPoints int64
34	timeRange     backend.TimeRange
35	interval      time.Duration
36	period        time.Duration
37}
38
39func (s *Service) handleUSAScenario(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
40	resp := backend.NewQueryDataResponse()
41
42	for _, q := range req.Queries {
43		wrapper := &usaQueryWrapper{}
44		err := json.Unmarshal(q.JSON, wrapper)
45		if err != nil {
46			return nil, fmt.Errorf("failed to parse query json: %v", err)
47		}
48
49		usa := wrapper.USA
50		periodString := usa.Period
51		if periodString == "" {
52			periodString = "30m"
53		}
54		period, err := time.ParseDuration(periodString)
55		if err != nil {
56			return nil, fmt.Errorf("failed to parse query json: %v", err)
57		}
58
59		if len(usa.Fields) == 0 {
60			usa.Fields = []string{"foo", "bar", "baz"}
61		}
62		if len(usa.States) == 0 {
63			usa.States = allStateCodes
64		}
65
66		usa.period = period
67		usa.maxDataPoints = q.MaxDataPoints * 2
68		usa.timeRange = q.TimeRange
69		usa.interval = q.Interval
70
71		resp.Responses[q.RefID] = doStateQuery(usa)
72	}
73
74	return resp, nil
75}
76
77func doStateQuery(query usaQuery) backend.DataResponse {
78	switch query.Mode {
79	default:
80	case modeTimeseries:
81		return getStateValueAsTimeseries(query, false)
82	case modeTimeseriesWide:
83		return getStateValueAsTimeseries(query, true)
84	case modeValueAsFields:
85		return getStateValueAsFrame(query.timeRange.To, query, false)
86	case modeValueAsLabeledFields:
87		return getStateValueAsFrame(query.timeRange.To, query, true)
88	}
89	return getStateValueAsRow(query.timeRange.To, query)
90}
91
92func getStateValueAsFrame(t time.Time, query usaQuery, asLabel bool) backend.DataResponse {
93	dr := backend.DataResponse{}
94
95	var labels data.Labels
96	for _, fname := range query.Fields {
97		frame := data.NewFrame(fname)
98		vals := getStateValues(t, fname, query)
99		for idx, state := range query.States {
100			name := state
101			if asLabel {
102				labels = data.Labels{"state": state}
103				name = ""
104			}
105			field := data.NewField(name, labels, []float64{vals[idx]})
106			frame.Fields = append(frame.Fields, field)
107		}
108		dr.Frames = append(dr.Frames, frame)
109	}
110	return dr
111}
112
113// One frame for each time+value
114func getStateValueAsTimeseries(query usaQuery, wide bool) backend.DataResponse {
115	dr := backend.DataResponse{}
116	tr := query.timeRange
117
118	var labels data.Labels
119	for _, fname := range query.Fields {
120		timeWalkerMs := tr.From.UnixNano() / int64(time.Millisecond)
121		to := tr.To.UnixNano() / int64(time.Millisecond)
122		stepMillis := query.interval.Milliseconds()
123
124		// Calculate all values
125		timeVals := make([]time.Time, 0)
126		stateVals := make([][]float64, 0)
127		for i := int64(0); i < query.maxDataPoints && timeWalkerMs < to; i++ {
128			t := time.Unix(timeWalkerMs/int64(1e+3), (timeWalkerMs%int64(1e+3))*int64(1e+6)).UTC()
129
130			vals := getStateValues(t, fname, query)
131			timeVals = append(timeVals, t)
132			stateVals = append(stateVals, vals)
133
134			timeWalkerMs += stepMillis
135		}
136
137		values := make([]float64, len(timeVals))
138		for idx, state := range query.States {
139			for i := 0; i < len(timeVals); i++ {
140				values[i] = stateVals[i][idx]
141			}
142
143			labels = data.Labels{"state": state}
144			frame := data.NewFrame(fname,
145				data.NewField(data.TimeSeriesTimeFieldName, nil, timeVals),
146				data.NewField(data.TimeSeriesValueFieldName, labels, values),
147			)
148			dr.Frames = append(dr.Frames, frame)
149		}
150	}
151
152	// Stick them next to eachother
153	if wide {
154		wideFrame := data.NewFrame("", dr.Frames[0].Fields[0])
155		for _, frame := range dr.Frames {
156			field := frame.Fields[1]
157			field.Name = frame.Name
158			wideFrame.Fields = append(wideFrame.Fields, field)
159		}
160		dr.Frames = data.Frames{wideFrame}
161	}
162
163	return dr
164}
165
166func getStateValueAsRow(t time.Time, query usaQuery) backend.DataResponse {
167	frame := data.NewFrame("", data.NewField("state", nil, query.States))
168	for _, f := range query.Fields {
169		frame.Fields = append(frame.Fields, data.NewField(f, nil, getStateValues(t, f, query)))
170	}
171
172	return backend.DataResponse{
173		Frames: data.Frames{frame},
174	}
175}
176
177func getStateValues(t time.Time, field string, query usaQuery) []float64 {
178	tv := float64(t.UnixNano())
179	pn := float64(query.period.Nanoseconds())
180	incr := pn / float64(len(query.States))
181
182	fn := math.Sin
183
184	// period offsets
185	switch field {
186	case "bar":
187		fn = math.Cos
188	case "baz":
189		fn = math.Tan
190	}
191
192	values := make([]float64, len(query.States))
193	for i := range query.States {
194		tv += incr
195		value := fn(float64(int64(tv) % int64(pn)))
196		// We shall truncate float64 to a certain precision, because otherwise we might
197		// get different cos/tan results across different CPU architectures and compilers.
198		values[i] = float64(int(value*1e10)) / 1e10
199	}
200	return values
201}
202