package main

import (
	"bytes"
	"context"
	_ "embed"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"maps"
	"net/http"
	"regexp"
	"strings"
	"text/template"
	"time"

	"github.com/go-gota/gota/dataframe"
	"github.com/go-gota/gota/series"
	"github.com/pingcap-inc/trino-thrift-connectors/gen-go/trino"
	"github.com/pingcap-inc/trino-thrift-connectors/schema"
	"github.com/samber/lo"
)

type metricsSchemaHandler struct {
	client *ClinicClient

	// to use trino.TrinoThriftTableMetadata here, we have added a `expr` field in
	// trino.TrinoThriftColumnMetadata struct, if you regenerate the trino.go file,
	// please add that field again.
	tables []*trino.TrinoThriftTableMetadata
}

//go:embed config/metrics.json
var metrics_cfg []byte

func newMetricsSchemaHandler(client *ClinicClient) (*metricsSchemaHandler, error) {
	tables := []*trino.TrinoThriftTableMetadata{}
	if err := json.Unmarshal(metrics_cfg, &tables); err != nil {
		return nil, err
	}

	for _, t := range tables {
		t.IndexableKeys = [][]string{[]string{"cluster_id"}}
	}
	return &metricsSchemaHandler{client, tables}, nil
}

func (h *metricsSchemaHandler) ListTables(
	ctx context.Context,
) (_r []*trino.TrinoThriftTableMetadata, _err error) {
	return h.tables, nil
}

func (h *metricsSchemaHandler) GetSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	outputConstraint *trino.TrinoThriftTupleDomain,
) (_r []*schema.SplitInfo, _err error) {
	_, count, err := h.client.ListClusters(map[string][]string{
		"limit":          []string{fmt.Sprintf("%d", PAGE_SIZE)},
		"page":           []string{"1"},
		"sort":           []string{"clusterCreatedAt"},
		"order":          []string{"ascend"},
		"deploy_type":    []string{"dedicated"},
		"cluster_status": []string{"normal", "upgrading", "scaling"},
	})
	if err != nil {
		return nil, err
	}

	splits := []*schema.SplitInfo{}
	max_page := (int(count)-1)/int(PAGE_SIZE) + 1
	for idx := 0; idx < max_page; idx++ {
		splits = append(splits, &schema.SplitInfo{
			SchemaTableName:  schemaTableName,
			OutputConstraint: outputConstraint,
			PartInfo:         idx + 1,
		})
	}

	return splits, nil
}

func (h *metricsSchemaHandler) GetIndexSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	outputConstraint *trino.TrinoThriftTupleDomain,
	indexColumnNames []string,
	keys *trino.TrinoThriftPageResult_,
) (_r []*schema.SplitInfo, _err error) {
	_, count, err := h.client.ListClusters(map[string][]string{
		"limit":          []string{fmt.Sprintf("%d", PAGE_SIZE)},
		"page":           []string{"1"},
		"sort":           []string{"clusterCreatedAt"},
		"order":          []string{"ascend"},
		"deploy_type":    []string{"dedicated"},
		"cluster_status": []string{"normal", "upgrading", "scaling"},
	})
	if err != nil {
		return nil, err
	}

	splits := []*schema.SplitInfo{}
	max_page := (int(count)-1)/int(PAGE_SIZE) + 1
	for idx := 0; idx < max_page; idx++ {
		splits = append(splits, &schema.SplitInfo{
			SchemaTableName:  schemaTableName,
			OutputConstraint: outputConstraint,
			IndexColumnNames: indexColumnNames,
			Keys:             keys,
			PartInfo:         idx + 1,
		})
	}

	return splits, nil
}

func (h *metricsSchemaHandler) getTableMetadata(
	tbinfo *trino.TrinoThriftSchemaTableName,
) *trino.TrinoThriftTableMetadata {
	for _, tb := range h.tables {
		if tbinfo.SchemaName == tb.SchemaTableName.SchemaName &&
			tbinfo.TableName == tb.SchemaTableName.TableName {
			return tb
		}
	}
	return nil
}

func (h *metricsSchemaHandler) render(t any, cluster Cluster) (string, error) {
	tstr, ok := t.(string)
	if !ok {
		return "", fmt.Errorf("string required, got %v", t)
	}
	tmpl, err := template.New("expr").Parse(tstr)
	if err != nil {
		return "", err
	}

	var data bytes.Buffer
	if err := tmpl.Execute(&data, cluster); err != nil {
		return "", err
	}
	return data.String(), nil
}

func (h *metricsSchemaHandler) groups(expr string) []string {
	re := regexp.MustCompile(`by \(([A-Za-z0-9_, ]+)*\)`)
	xs := re.FindAllStringSubmatch(expr, -1)
	if len(xs) == 0 {
		return nil
	}
	return lo.Map(strings.Split(xs[len(xs)-1][1], ","), func(g string, _ int) string {
		return strings.Trim(g, " ")
	})
}

func (h *metricsSchemaHandler) loadColumnData(
	ctx context.Context,
	column *trino.TrinoThriftColumnMetadata,
	clusters []Cluster,
	begin_ts int64,
	end_ts int64,
) (*dataframe.DataFrame, error) {
	url := CLINIC_ENDPOINT + "/grafana/api/ds/query"
	meta, ok := column.Meta.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("meta field should be a map on column %s", column.Name)
	}
	if meta["expr"], ok = meta["expr"].(string); !ok {
		return nil, fmt.Errorf("meta.expr field should be a string on column %s", column.Name)
	}
	meta["instant"] = true
	meta["refId"] = column.Name
	meta["legendFormat"] = column.Name

	gs := h.groups(meta["expr"].(string))
	if len(gs) > 0 {
		meta["legendFormat"] = strings.Join(lo.Map(gs, func(g string, _ int) string {
			return "{{" + g + "}}"
		}), "/")
	}

	rs := []*http.Request{}
	for _, cluster := range clusters {
		var err error
		var referer string
		m := maps.Clone(meta)
		if m["expr"], err = h.render(m["expr"], cluster); err != nil {
			return nil, err
		}
		if r, ok := m["referer"]; ok {
			if referer, err = h.render(r, cluster); err != nil {
				return nil, err
			}
		}
		payload, err := json.Marshal(map[string]any{
			"queries": []map[string]any{m},
			"from":    fmt.Sprint(begin_ts),
			"to":      fmt.Sprint(end_ts),
		})
		if err != nil {
			return nil, err
		}
		r, _ := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(payload))
		if referer != "" {
			r.Header.Set("referer", referer)
		}
		rs = append(rs, r)
	}

	xs, errs := h.client.BatchDoWithCredential(rs)
	metrics := make([]PrometheusMetric, len(xs))

	for idx, res := range xs {
		// if this request failed
		if errs[idx] != nil {
			continue
		}
		defer res.Body.Close()

		var body []byte
		if body, errs[idx] = io.ReadAll(res.Body); errs[idx] != nil {
			continue
		}
		if res.StatusCode != 200 {
			errs[idx] = errors.New(string(body))
			continue
		}
		if errs[idx] = json.Unmarshal(body, &metrics[idx]); errs[idx] != nil {
			continue
		}
	}

	ss := []series.Series{}
	ids := []string{}
	times := []int{}
	values := []float64{}
	for idx, metric := range metrics {
		if errs[idx] != nil {
			continue
		}
		for _, frame := range metric.Results[column.Name].Frames {
			ids = append(ids, clusters[idx].ClusterID)
			times = append(times, int(begin_ts))
			values = append(values, frame.Data.Values[1][0])
		}
	}
	ss = append(ss, series.New(ids, series.String, "cluster_id"))
	ss = append(ss, series.New(times, series.Int, "time"))
	ss = append(ss, series.New(values, series.Float, column.Name))

	if len(gs) > 0 {
		for i, g := range gs {
			if g == "cluster_id" {
				continue
			}
			vs := []string{}
			for idx, metric := range metrics {
				if errs[idx] != nil {
					continue
				}
				for _, frame := range metric.Results[column.Name].Frames {
					xs := strings.Split(frame.Schema.Name, "/")
					vs = append(vs, xs[i])
				}
			}
			ss = append(ss, series.New(vs, series.String, g))
		}
	}

	// logging the erorrs
	for idx, err := range errs {
		if err != nil {
			fmt.Printf("error on fetch cluster %s metrics: %s\n", clusters[idx].ClusterID, err)
		}
	}

	df := dataframe.New(ss...)
	return &df, nil
}

func (h *metricsSchemaHandler) buildThriftBlocks(
	table *trino.TrinoThriftTableMetadata, df dataframe.DataFrame,
) (map[string]*trino.TrinoThriftBlock, error) {
	blocks := map[string]*trino.TrinoThriftBlock{}

	for _, column := range table.Columns {
		s := df.Col(column.Name)
		if s.Error() != nil {
			return nil, fmt.Errorf("error on processing column `%s`: %s", column.Name, s.Error())
		}

		if strings.HasPrefix(column.Type, "varchar") {
			values := s.Records()
			blocks[column.Name] = &trino.TrinoThriftBlock{
				VarcharData: &trino.TrinoThriftVarchar{
					Nulls: lo.Map(values, func(v string, _ int) bool {
						return false
					}), Sizes: lo.Map(values, func(v string, _ int) int32 {
						return int32(len(v))
					}), Bytes: []byte(strings.Join(values, "")),
				},
			}
		} else if column.Type == "double" {
			values := s.Float()
			blocks[column.Name] = &trino.TrinoThriftBlock{
				DoubleData: &trino.TrinoThriftDouble{
					Nulls: lo.Map(values, func(v float64, _ int) bool {
						return false
					}),
					Doubles: values,
				},
			}
		} else if column.Type == "timestamp" {
			values, err := s.Int()
			if err != nil {
				return nil, err
			}
			blocks[column.Name] = &trino.TrinoThriftBlock{
				TimestampData: &trino.TrinoThriftTimestamp{
					Nulls: lo.Map(values, func(v int, _ int) bool {
						return false
					}),
					Timestamps: lo.Map(values, func(v int, _ int) int64 {
						return int64(v)
					}),
				},
			}
		} else {
			return nil, fmt.Errorf("unsupported column type: %s", column.Type)
		}
	}

	return blocks, nil
}

func (h *metricsSchemaHandler) GetRows(
	ctx context.Context,
	split *schema.SplitInfo,
	columns []string,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	begin_ts, end_ts := getTimeRange(split.OutputConstraint, time.Minute*5)
	begin_ts, end_ts = begin_ts*1000, end_ts*1000
	clusters, _, err := h.client.ListClusters(map[string][]string{
		"limit":          []string{fmt.Sprintf("%d", PAGE_SIZE)},
		"page":           []string{fmt.Sprintf("%d", split.PartInfo)},
		"sort":           []string{"clusterCreatedAt"},
		"order":          []string{"ascend"},
		"deploy_type":    []string{"dedicated"},
		"cluster_status": []string{"normal", "upgrading", "scaling"},
	})
	if err != nil {
		return nil, err
	}

	df := dataframe.LoadStructs(clusters)
	indexes := split.Filter(df)

	// filter clusters by indexes
	clusters = lo.Map(indexes, func(idx int, _ int) Cluster {
		return clusters[idx]
	})
	if len(clusters) == 0 {
		return &trino.TrinoThriftPageResult_{
			ColumnBlocks: nil,
			RowCount:     0,
			NextToken:    nil,
		}, nil
	}
	df = dataframe.LoadStructs(clusters)

	table := h.getTableMetadata(split.SchemaTableName)
	if table == nil {
		return nil, fmt.Errorf(
			"table %s.%s not found",
			split.SchemaTableName.SchemaName,
			split.SchemaTableName.TableName,
		)
	}

	dfs := []dataframe.DataFrame{}
	for _, ci := range table.Columns {
		if ci.Meta == nil {
			continue
		}
		f, err := h.loadColumnData(ctx, ci, clusters, begin_ts, end_ts)
		if err != nil {
			return nil, err
		}
		dfs = append(dfs, *f)
	}

	for _, f := range dfs {
		joinOn := lo.Intersect(df.Names(), f.Names())
		df = df.OuterJoin(f, joinOn...)
	}

	blocks, err := h.buildThriftBlocks(table, df)
	if err != nil {
		return nil, err
	}

	column_blocks := []*trino.TrinoThriftBlock{}
	for _, column := range columns {
		column_blocks = append(column_blocks, blocks[column])
	}

	r := &trino.TrinoThriftPageResult_{
		ColumnBlocks: column_blocks,
		RowCount:     int32(df.Nrow()),
		NextToken:    nil,
	}
	return r, nil
}