package main

import (
	"context"
	"fmt"
	"time"

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

type dedicatedEventSchemaHandler struct {
	client *ClinicClient
}

func getTimeRange(outputConstraint *trino.TrinoThriftTupleDomain, defaultDuration time.Duration) (int64, int64) {
	end_ts := time.Now().Unix()
	begin_ts := end_ts - int64(defaultDuration.Seconds())
	domain := outputConstraint.Domains["time"]
	if domain == nil || len(domain.ValueSet.RangeValueSet.Ranges) != 1 {
		return begin_ts, end_ts
	}

	r := domain.ValueSet.RangeValueSet.Ranges[0]

	if r.Low.Value != nil {
		begin_ts = r.Low.Value.TimestampData.Timestamps[0] / 1000
	}
	if r.High.Value != nil {
		end_ts = r.High.Value.TimestampData.Timestamps[0] / 1000
	}

	return begin_ts, end_ts
}

func (h *dedicatedEventSchemaHandler) ListTables(
	ctx context.Context,
) (_r []*trino.TrinoThriftTableMetadata, _err error) {
	return []*trino.TrinoThriftTableMetadata{
		&trino.TrinoThriftTableMetadata{
			SchemaTableName: &trino.TrinoThriftSchemaTableName{
				SchemaName: "dedicated_events",
				TableName:  "events_in_recent_7_days",
			},
			Columns: []*trino.TrinoThriftColumnMetadata{
				&trino.TrinoThriftColumnMetadata{
					Name:    "tenant_id",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the owner of the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "cluster_id",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the id of the cluster where the event is collected from"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "name",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the name of the event"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "status",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the status of the event"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "time",
					Type:    "timestamp",
					Comment: lo.ToPtr("the start time of the event (UTC+0)"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "summary",
					Type:    "varchar(1024)",
					Comment: lo.ToPtr("a brief description of what happened"),
				},
			},
			Comment: lo.ToPtr("events happened on clusters in the recent 7 days"),
			IndexableKeys: [][]string{
				[]string{"cluster_id"},
			},
		},
	}, nil
}

func (h *dedicatedEventSchemaHandler) GetSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	outputConstraint *trino.TrinoThriftTupleDomain,
) (_r []*schema.SplitInfo, _err error) {
	max_d := time.Hour * 24 * 7
	begin_ts, end_ts := getTimeRange(outputConstraint, max_d)
	if end_ts-begin_ts > int64(max_d.Seconds()) {
		begin_ts = end_ts - int64(max_d.Seconds())
	}
	splits := []*schema.SplitInfo{}
	page := 0

	for {
		page += 1
		clusters, _, err := h.client.ListClusters(map[string][]string{
			"limit":        []string{fmt.Sprint(PAGE_SIZE)},
			"page":         []string{fmt.Sprint(page)},
			"sort":         []string{"clusterCreatedAt"},
			"order":        []string{"ascend"},
			"deploy_type":  []string{"dedicated"},
			"show_deleted": []string{"false"},
		})
		if err != nil {
			return nil, err
		}

		cs := []Cluster{}
		for _, cluster := range clusters {
			if int64(cluster.ClusterCreatedAt) > end_ts*1000 {
				continue
			}
			cs = append(cs, cluster)
		}
		if len(cs) > 0 {
			splits = append(splits, &schema.SplitInfo{
				SchemaTableName:  schemaTableName,
				OutputConstraint: outputConstraint,
				PartInfo:         cs,
			})
		}
		if len(cs) != PAGE_SIZE {
			break
		}
	}

	page = 0
	for {
		page += 1
		clusters, _, err := h.client.ListClusters(map[string][]string{
			"limit":          []string{fmt.Sprint(PAGE_SIZE)},
			"page":           []string{fmt.Sprint(page)},
			"sort":           []string{"clusterDeletedAt"},
			"order":          []string{"descend"},
			"deploy_type":    []string{"dedicated"},
			"show_deleted":   []string{"true"},
			"cluster_status": []string{"cleared", "deleted", "deletecomplete"},
		})
		if err != nil {
			return nil, err
		}

		cs := []Cluster{}
		for _, cluster := range clusters {
			if int64(cluster.ClusterDeletedAt) < begin_ts*1000 {
				continue
			}
			cs = append(cs, cluster)
		}
		if len(cs) > 0 {
			splits = append(splits, &schema.SplitInfo{
				SchemaTableName:  schemaTableName,
				OutputConstraint: outputConstraint,
				PartInfo:         cs,
			})
		}
		if len(cs) != PAGE_SIZE {
			break
		}
	}

	return splits, nil
}

func (h *dedicatedEventSchemaHandler) GetIndexSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	outputConstraint *trino.TrinoThriftTupleDomain,
	indexColumnNames []string,
	keys *trino.TrinoThriftPageResult_,
) (_r []*schema.SplitInfo, _err error) {
	splits, err := h.GetSplits(ctx, schemaTableName, outputConstraint)
	if err != nil {
		return nil, err
	}

	for idx := range splits {
		splits[idx].IndexColumnNames = indexColumnNames
		splits[idx].Keys = keys
	}

	return splits, nil
}

func (h *dedicatedEventSchemaHandler) GetRows(
	ctx context.Context,
	split *schema.SplitInfo,
	columns []string,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	clusters := split.PartInfo.([]Cluster)
	df := dataframe.LoadStructs(clusters)
	indexes := split.Filter(df)

	// filter clusters by indexes
	clusters = lo.Map(indexes, func(idx int, _ int) Cluster {
		return clusters[idx]
	})

	switch split.SchemaTableName.TableName {
	case "events_in_recent_7_days":
		max_d := time.Hour * 24 * 7
		begin_ts, end_ts := getTimeRange(split.OutputConstraint, max_d)
		if end_ts-begin_ts > int64(max_d.Seconds()) {
			begin_ts = end_ts - int64(max_d.Seconds())
		}
		return h.load7DaysClusterEvents(ctx, columns, clusters, begin_ts, end_ts)
	default:
		panic("unreachable")
	}

	return nil, nil
}

func (h *dedicatedEventSchemaHandler) load7DaysClusterEvents(
	ctx context.Context, columns []string, clusters []Cluster, begin_ts, end_ts int64,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	blocks := map[string]*trino.TrinoThriftBlock{
		"tenant_id": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"cluster_id": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"name": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"status": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"time": &trino.TrinoThriftBlock{
			TimestampData: &trino.TrinoThriftTimestamp{Nulls: []bool{}, Timestamps: []int64{}},
		},
		"summary": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
	}

	events, errs := h.client.ClusterEvents(ctx, clusters, begin_ts, end_ts)
	for _, err := range errs {
		if err != nil {
			fmt.Println(err)
		}
	}

	for _, event := range events {
		tenant_id_data := blocks["tenant_id"].VarcharData
		tenant_id_data.Nulls = append(tenant_id_data.Nulls, false)
		tenant_id_data.Sizes = append(tenant_id_data.Sizes, int32(len(event.TenantID)))
		tenant_id_data.Bytes = append(tenant_id_data.Bytes, []byte(event.TenantID)...)

		cluster_id_data := blocks["cluster_id"].VarcharData
		cluster_id_data.Nulls = append(cluster_id_data.Nulls, false)
		cluster_id_data.Sizes = append(cluster_id_data.Sizes, int32(len(event.ClusterID)))
		cluster_id_data.Bytes = append(cluster_id_data.Bytes, []byte(event.ClusterID)...)

		name_data := blocks["name"].VarcharData
		name_data.Nulls = append(name_data.Nulls, false)
		name_data.Sizes = append(name_data.Sizes, int32(len(event.Name)))
		name_data.Bytes = append(name_data.Bytes, []byte(event.Name)...)

		status_data := blocks["status"].VarcharData
		status_data.Nulls = append(status_data.Nulls, false)
		status_data.Sizes = append(status_data.Sizes, int32(len(event.Status)))
		status_data.Bytes = append(status_data.Bytes, []byte(event.Status)...)

		time_data := blocks["time"].TimestampData
		time_data.Nulls = append(time_data.Nulls, false)
		time_data.Timestamps = append(time_data.Timestamps, int64(event.Time))

		summary_data := blocks["summary"].VarcharData
		summary_data.Nulls = append(summary_data.Nulls, false)
		summary_data.Sizes = append(summary_data.Sizes, int32(len(event.Summary)))
		summary_data.Bytes = append(summary_data.Bytes, []byte(event.Summary)...)
	}

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

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