package main

import (
	"context"
	"fmt"
	"strings"

	"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 dedicatedClusterSchemaHandler struct {
	client *ClinicClient
}

func (h *dedicatedClusterSchemaHandler) ListTables(
	ctx context.Context,
) (_r []*trino.TrinoThriftTableMetadata, _err error) {
	return []*trino.TrinoThriftTableMetadata{
		&trino.TrinoThriftTableMetadata{
			SchemaTableName: &trino.TrinoThriftSchemaTableName{
				SchemaName: "dedicated_clusters",
				TableName:  "cluster_config",
			},
			Columns: []*trino.TrinoThriftColumnMetadata{
				&trino.TrinoThriftColumnMetadata{
					Name:    "cluster_id",
					Type:    "varchar(255)",
					Comment: lo.ToPtr("the id of the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "kind",
					Type:    "varchar(255)",
					Comment: lo.ToPtr("the kind (tidb_variable, tidb_config, pd_config, tiflash_config .etc.) of the configration"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "config_name",
					Type:    "varchar",
					Comment: lo.ToPtr("the name of the configration"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "config_value",
					Type:    "varchar",
					Comment: lo.ToPtr("the value of the configration"),
				},
			},
			Comment: lo.ToPtr("the configuration of each component in dedicated clusters"),
			IndexableKeys: [][]string{
				[]string{"cluster_id"},
			},
		},
		&trino.TrinoThriftTableMetadata{
			SchemaTableName: &trino.TrinoThriftSchemaTableName{
				SchemaName: "dedicated_clusters",
				TableName:  "cluster_status",
			},
			Columns: []*trino.TrinoThriftColumnMetadata{
				&trino.TrinoThriftColumnMetadata{
					Name:    "cluster_id",
					Type:    "varchar(255)",
					Comment: lo.ToPtr("the id of the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "cluster_proxy_availability",
					Type:    "varchar(255)",
					Comment: lo.ToPtr("if the cluster proxy component is available ('YES' or 'NO')"),
				},
			},
			Comment: lo.ToPtr("the status of the dedicated clusters"),
			IndexableKeys: [][]string{
				[]string{"cluster_id"},
			},
		},
		&trino.TrinoThriftTableMetadata{
			SchemaTableName: &trino.TrinoThriftSchemaTableName{
				SchemaName: "dedicated_clusters",
				TableName:  "cluster_list",
			},
			Columns: []*trino.TrinoThriftColumnMetadata{
				&trino.TrinoThriftColumnMetadata{
					Name:    "tenant_id",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the id of tenant who own the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "tenant_name",
					Type:    "varchar(255)",
					Comment: lo.ToPtr("the name of tenant who own the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "cluster_id",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the id of the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "cluster_name",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the name of the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "region",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the region where the cluster running"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "provider",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the cloud provider (aws or gcp)"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "version",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the version of the cluster"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "cluster_created_at",
					Type:    "timestamp",
					Comment: lo.ToPtr("the time when the cluster created at"),
				},
			},
			Comment: lo.ToPtr("the list of the dedicated clusters"),
			IndexableKeys: [][]string{
				[]string{"cluster_id"},
			},
		},
	}, nil
}

func (h *dedicatedClusterSchemaHandler) 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 *dedicatedClusterSchemaHandler) 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.Sprint(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 *dedicatedClusterSchemaHandler) GetRows(
	ctx context.Context,
	split *schema.SplitInfo,
	columns []string,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	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]
	})

	switch split.SchemaTableName.TableName {
	case "cluster_list":
		return h.loadClusterList(ctx, columns, clusters)
	case "cluster_status":
		return h.loadClusterStatus(ctx, columns, clusters)
	case "cluster_config":
		return h.loadClusterConfig(ctx, columns, clusters)
	default:
		panic("unreachable")
	}
}

func (h *dedicatedClusterSchemaHandler) loadClusterList(
	ctx context.Context, columns []string, clusters []Cluster,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	blocks := map[string]*trino.TrinoThriftBlock{
		"tenant_id": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"tenant_name": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"cluster_id": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"cluster_name": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"region": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"provider": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"version": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"cluster_created_at": &trino.TrinoThriftBlock{
			TimestampData: &trino.TrinoThriftTimestamp{Nulls: []bool{}, Timestamps: []int64{}},
		},
	}

	for _, cluster := range clusters {
		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(cluster.TenantID)))
		tenant_id_data.Bytes = append(tenant_id_data.Bytes, []byte(cluster.TenantID)...)

		tenant_name_data := blocks["tenant_name"].VarcharData
		tenant_name_data.Nulls = append(tenant_name_data.Nulls, false)
		tenant_name_data.Sizes = append(tenant_name_data.Sizes, int32(len(cluster.TenantName)))
		tenant_name_data.Bytes = append(tenant_name_data.Bytes, []byte(cluster.TenantName)...)

		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(cluster.ClusterID)))
		cluster_id_data.Bytes = append(cluster_id_data.Bytes, []byte(cluster.ClusterID)...)

		cluster_name_data := blocks["cluster_name"].VarcharData
		cluster_name_data.Nulls = append(cluster_name_data.Nulls, false)
		cluster_name_data.Sizes = append(cluster_name_data.Sizes, int32(len(cluster.ClusterName)))
		cluster_name_data.Bytes = append(cluster_name_data.Bytes, []byte(cluster.ClusterName)...)

		region_data := blocks["region"].VarcharData
		region_data.Nulls = append(region_data.Nulls, false)
		region_data.Sizes = append(region_data.Sizes, int32(len(cluster.ClusterRegionName)))
		region_data.Bytes = append(region_data.Bytes, []byte(cluster.ClusterRegionName)...)

		provider_data := blocks["provider"].VarcharData
		provider_data.Nulls = append(provider_data.Nulls, false)
		provider_data.Sizes = append(provider_data.Sizes, int32(len(cluster.ClusterProviderName)))
		provider_data.Bytes = append(provider_data.Bytes, []byte(cluster.ClusterProviderName)...)

		version_data := blocks["version"].VarcharData
		version_data.Nulls = append(version_data.Nulls, false)
		version_data.Sizes = append(version_data.Sizes, int32(len(cluster.ClusterVersion)))
		version_data.Bytes = append(version_data.Bytes, []byte(cluster.ClusterVersion)...)

		cluster_created_at_data := blocks["cluster_created_at"].TimestampData
		cluster_created_at_data.Nulls = append(cluster_created_at_data.Nulls, false)
		cluster_created_at_data.Timestamps = append(cluster_created_at_data.Timestamps, int64(cluster.ClusterCreatedAt))
	}

	column_blocks := []*trino.TrinoThriftBlock{}
	for _, column := range columns {
		column_blocks = append(column_blocks, blocks[column])
	}
	r := &trino.TrinoThriftPageResult_{
		ColumnBlocks: column_blocks,
		RowCount:     int32(len(clusters)),
		NextToken:    nil,
	}
	return r, nil
}

func (h *dedicatedClusterSchemaHandler) loadClusterStatus(
	ctx context.Context, columns []string, clusters []Cluster,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	blocks := map[string]*trino.TrinoThriftBlock{
		"cluster_id": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"cluster_proxy_availability": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
	}

	var status_list []bool
	var errs []error
	if lo.Contains(columns, "cluster_proxy_availability") {
		status_list, errs = h.client.ClusterProxyAvailability(ctx, clusters)
	}

	for idx, cluster := range clusters {
		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(cluster.ClusterID)))
		cluster_id_data.Bytes = append(cluster_id_data.Bytes, []byte(cluster.ClusterID)...)

		if !lo.Contains(columns, "cluster_proxy_availability") {
			continue
		}
		availability_data := blocks["cluster_proxy_availability"].VarcharData
		if errs[idx] != nil {
			availability_data.Nulls = append(availability_data.Nulls, true)
			availability_data.Sizes = append(availability_data.Sizes, 0)
		} else {
			availability_data.Nulls = append(availability_data.Nulls, false)
			if status_list[idx] {
				availability_data.Sizes = append(availability_data.Sizes, 3)
				availability_data.Bytes = append(availability_data.Bytes, []byte("YES")...)
			} else {
				availability_data.Sizes = append(availability_data.Sizes, 2)
				availability_data.Bytes = append(availability_data.Bytes, []byte("NO")...)
			}
		}
	}

	column_blocks := []*trino.TrinoThriftBlock{}
	for _, column := range columns {
		column_blocks = append(column_blocks, blocks[column])
	}
	r := &trino.TrinoThriftPageResult_{
		ColumnBlocks: column_blocks,
		RowCount:     int32(len(clusters)),
		NextToken:    nil,
	}
	return r, nil
}

func (h *dedicatedClusterSchemaHandler) loadClusterConfig(
	ctx context.Context, columns []string, clusters []Cluster,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	blocks := map[string]*trino.TrinoThriftBlock{
		"cluster_id": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"kind": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"config_name": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"config_value": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
	}

	row_count := 0
	cluster_configs, errs := h.client.ClusterConfig(ctx, clusters)

	for idx, cluster := range clusters {
		if errs[idx] != nil {
			fmt.Println(_err)
			continue
		}
		configs := cluster_configs[idx]

		for _, items := range configs {
			// ignore some useless content
			if (items[0] == "tidb_variable" && items[1] == "tidb_config") ||
				strings.Contains(items[1], "-addr") ||
				items[1] == "labels.host" {
				continue
			}
			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(cluster.ClusterID)))
			cluster_id_data.Bytes = append(cluster_id_data.Bytes, []byte(cluster.ClusterID)...)

			kind_data := blocks["kind"].VarcharData
			kind_data.Nulls = append(kind_data.Nulls, false)
			kind_data.Sizes = append(kind_data.Sizes, int32(len(items[0])))
			kind_data.Bytes = append(kind_data.Bytes, []byte(items[0])...)

			config_name_data := blocks["config_name"].VarcharData
			config_name_data.Nulls = append(config_name_data.Nulls, false)
			config_name_data.Sizes = append(config_name_data.Sizes, int32(len(items[1])))
			config_name_data.Bytes = append(config_name_data.Bytes, []byte(items[1])...)

			config_value_data := blocks["config_value"].VarcharData
			config_value_data.Nulls = append(config_value_data.Nulls, false)
			config_value_data.Sizes = append(config_value_data.Sizes, int32(len(items[2])))
			config_value_data.Bytes = append(config_value_data.Bytes, []byte(items[2])...)

			row_count += 1
		}
	}

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

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