package main

import (
	"context"

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

type cloudQuotaSchemaHandler struct {
	client *ClinicClient
}

func (h *cloudQuotaSchemaHandler) ListTables(
	ctx context.Context,
) (_r []*trino.TrinoThriftTableMetadata, _err error) {
	return []*trino.TrinoThriftTableMetadata{
		&trino.TrinoThriftTableMetadata{
			SchemaTableName: &trino.TrinoThriftSchemaTableName{
				SchemaName: "cloud_account_quota",
				TableName:  "aws",
			},
			Columns: []*trino.TrinoThriftColumnMetadata{
				&trino.TrinoThriftColumnMetadata{
					Name:    "aws_account",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the aws account number the resource belong to"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "region",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the region of the resource"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "resource",
					Type:    "varchar(128)",
					Comment: lo.ToPtr("the resource name"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "limit",
					Type:    "double",
					Comment: lo.ToPtr("usage limit of the resource"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "used",
					Type:    "double",
					Comment: lo.ToPtr("used of the resource"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "usage",
					Type:    "double",
					Comment: lo.ToPtr("the ratio usage (used/limit) of the resource"),
				},
			},
			Comment: lo.ToPtr("the limit and usage of resources in aws"),
		},
		&trino.TrinoThriftTableMetadata{
			SchemaTableName: &trino.TrinoThriftSchemaTableName{
				SchemaName: "cloud_account_quota",
				TableName:  "gcp",
			},
			Columns: []*trino.TrinoThriftColumnMetadata{
				&trino.TrinoThriftColumnMetadata{
					Name:    "project",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the gcp project the resource belong to"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "region",
					Type:    "varchar(64)",
					Comment: lo.ToPtr("the region of the resource"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "resource",
					Type:    "varchar(128)",
					Comment: lo.ToPtr("the resource name"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "limit",
					Type:    "double",
					Comment: lo.ToPtr("usage limit of the resource"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "used",
					Type:    "double",
					Comment: lo.ToPtr("used of the resource"),
				},
				&trino.TrinoThriftColumnMetadata{
					Name:    "usage",
					Type:    "double",
					Comment: lo.ToPtr("the ratio usage (used/limit) of the resource"),
				},
			},
			Comment: lo.ToPtr("the limit and usage of resources in gcp"),
		},
	}, nil
}

func (h *cloudQuotaSchemaHandler) GetSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	outputConstraint *trino.TrinoThriftTupleDomain,
) (_r []*schema.SplitInfo, _err error) {
	return []*schema.SplitInfo{
		&schema.SplitInfo{
			SchemaTableName:  schemaTableName,
			OutputConstraint: outputConstraint,
		},
	}, nil
}

func (h *cloudQuotaSchemaHandler) GetIndexSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	outputConstraint *trino.TrinoThriftTupleDomain,
	indexColumnNames []string,
	keys *trino.TrinoThriftPageResult_,
) (_r []*schema.SplitInfo, _err error) {
	return []*schema.SplitInfo{
		&schema.SplitInfo{
			SchemaTableName:  schemaTableName,
			OutputConstraint: outputConstraint,
		},
	}, nil
}

func (h *cloudQuotaSchemaHandler) GetRows(
	ctx context.Context,
	split *schema.SplitInfo,
	columns []string,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	if split.SchemaTableName.TableName == "aws" {
		return h.loadAWSQuota(ctx, columns)
	} else if split.SchemaTableName.TableName == "gcp" {
		return h.loadGCPQuota(ctx, columns)
	}

	panic("unreachable")
}

func (h *cloudQuotaSchemaHandler) loadAWSQuota(
	ctx context.Context, columns []string,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	blocks := map[string]*trino.TrinoThriftBlock{
		"aws_account": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"region": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"resource": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"limit": &trino.TrinoThriftBlock{
			DoubleData: &trino.TrinoThriftDouble{Nulls: []bool{}, Doubles: []float64{}},
		},
		"used": &trino.TrinoThriftBlock{
			DoubleData: &trino.TrinoThriftDouble{Nulls: []bool{}, Doubles: []float64{}},
		},
		"usage": &trino.TrinoThriftBlock{
			DoubleData: &trino.TrinoThriftDouble{Nulls: []bool{}, Doubles: []float64{}},
		},
	}

	xs, err := h.client.AWSAccountQuota()
	if err != nil {
		return nil, err
	}

	for _, quota := range xs {
		account_data := blocks["aws_account"].VarcharData
		account_data.Nulls = append(account_data.Nulls, false)
		account_data.Sizes = append(account_data.Sizes, int32(len(quota.AccountOrProject)))
		account_data.Bytes = append(account_data.Bytes, []byte(quota.AccountOrProject)...)

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

		resource_data := blocks["resource"].VarcharData
		resource_data.Nulls = append(resource_data.Nulls, false)
		resource_data.Sizes = append(resource_data.Sizes, int32(len(quota.Resource)))
		resource_data.Bytes = append(resource_data.Bytes, []byte(quota.Resource)...)

		limit_data := blocks["limit"].DoubleData
		limit_data.Nulls = append(limit_data.Nulls, false)
		limit_data.Doubles = append(limit_data.Doubles, quota.Limit)

		used_data := blocks["used"].DoubleData
		used_data.Nulls = append(used_data.Nulls, false)
		used_data.Doubles = append(used_data.Doubles, quota.Used)

		usage_data := blocks["usage"].DoubleData
		usage_data.Nulls = append(usage_data.Nulls, false)
		usage_data.Doubles = append(usage_data.Doubles, quota.Usage)
	}

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

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

func (h *cloudQuotaSchemaHandler) loadGCPQuota(
	ctx context.Context, columns []string,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	blocks := map[string]*trino.TrinoThriftBlock{
		"project": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"region": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"resource": &trino.TrinoThriftBlock{
			VarcharData: &trino.TrinoThriftVarchar{Nulls: []bool{}, Sizes: []int32{}, Bytes: []byte{}},
		},
		"limit": &trino.TrinoThriftBlock{
			DoubleData: &trino.TrinoThriftDouble{Nulls: []bool{}, Doubles: []float64{}},
		},
		"used": &trino.TrinoThriftBlock{
			DoubleData: &trino.TrinoThriftDouble{Nulls: []bool{}, Doubles: []float64{}},
		},
		"usage": &trino.TrinoThriftBlock{
			DoubleData: &trino.TrinoThriftDouble{Nulls: []bool{}, Doubles: []float64{}},
		},
	}

	xs, err := h.client.GCPAccountQuota()
	if err != nil {
		return nil, err
	}

	for _, quota := range xs {
		project_data := blocks["project"].VarcharData
		project_data.Nulls = append(project_data.Nulls, false)
		project_data.Sizes = append(project_data.Sizes, int32(len(quota.AccountOrProject)))
		project_data.Bytes = append(project_data.Bytes, []byte(quota.AccountOrProject)...)

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

		resource_data := blocks["resource"].VarcharData
		resource_data.Nulls = append(resource_data.Nulls, false)
		resource_data.Sizes = append(resource_data.Sizes, int32(len(quota.Resource)))
		resource_data.Bytes = append(resource_data.Bytes, []byte(quota.Resource)...)

		limit_data := blocks["limit"].DoubleData
		limit_data.Nulls = append(limit_data.Nulls, false)
		limit_data.Doubles = append(limit_data.Doubles, quota.Limit)

		used_data := blocks["used"].DoubleData
		used_data.Nulls = append(used_data.Nulls, false)
		used_data.Doubles = append(used_data.Doubles, quota.Used)

		usage_data := blocks["usage"].DoubleData
		usage_data.Nulls = append(usage_data.Nulls, false)
		usage_data.Doubles = append(usage_data.Doubles, quota.Usage)
	}

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

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