package main

import (
	"bytes"
	"context"
	"encoding/gob"
	"fmt"

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

const PAGE_SIZE = 10

type ClinicService struct {
	client   *ClinicClient
	handlers []schema.SchemaHandler
}

func NewClinicService(client *ClinicClient) (*ClinicService, error) {
	gob.Register([]Cluster{})

	handlers := []schema.SchemaHandler{
		&cloudQuotaSchemaHandler{client},
		&dedicatedClusterSchemaHandler{client},
		&dedicatedEventSchemaHandler{client},
	}

	h, err := newMetricsSchemaHandler(client)
	if err != nil {
		return nil, err
	}
	handlers = append(handlers, h)

	return &ClinicService{
		client:   client,
		handlers: handlers,
	}, nil
}

func (s *ClinicService) TrinoListSchemaNames(ctx context.Context) (_r []string, _err error) {
	tables := []*trino.TrinoThriftTableMetadata{}
	for _, h := range s.handlers {
		tbs, err := h.ListTables(ctx)
		if err != nil {
			return nil, err
		}
		tables = append(tables, tbs...)
	}
	schemas := lo.Map(tables, func(t *trino.TrinoThriftTableMetadata, i int) string {
		return t.SchemaTableName.SchemaName
	})
	return lo.Uniq(schemas), nil
}

func (s *ClinicService) TrinoListTables(
	ctx context.Context,
	schemaNameOrNull *trino.TrinoThriftNullableSchemaName,
) (_r []*trino.TrinoThriftSchemaTableName, _err error) {
	tables := []*trino.TrinoThriftTableMetadata{}
	for _, h := range s.handlers {
		tbs, err := h.ListTables(ctx)
		if err != nil {
			return nil, err
		}
		tables = append(tables, tbs...)
	}
	names := lo.Map(
		tables,
		func(t *trino.TrinoThriftTableMetadata, i int) *trino.TrinoThriftSchemaTableName {
			return t.SchemaTableName
		},
	)
	return names, nil
}

func (s *ClinicService) TrinoGetTableMetadata(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
) (_r *trino.TrinoThriftNullableTableMetadata, _err error) {
	for _, h := range s.handlers {
		tables, err := h.ListTables(ctx)
		if err != nil {
			return nil, err
		}
		for _, table := range tables {
			if table.SchemaTableName.SchemaName == schemaTableName.SchemaName &&
				table.SchemaTableName.TableName == schemaTableName.TableName {

				return &trino.TrinoThriftNullableTableMetadata{
					TableMetadata: table,
				}, nil
			}
		}
	}
	return nil, fmt.Errorf("no such table: %s.%s", schemaTableName.SchemaName, schemaTableName.TableName)
}

func (s *ClinicService) encode(splits []*schema.SplitInfo) []*trino.TrinoThriftSplit {
	xs := []*trino.TrinoThriftSplit{}
	for _, split := range splits {
		var buf bytes.Buffer
		enc := gob.NewEncoder(&buf)
		if err := enc.Encode(split); err != nil {
			panic(err)
		}

		xs = append(xs, &trino.TrinoThriftSplit{
			SplitId: &trino.TrinoThriftId{
				ID: buf.Bytes(),
			},
			Hosts: []*trino.TrinoThriftHostAddress{},
		})
	}
	return xs
}

func (s *ClinicService) TrinoGetSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	desiredColumns *trino.TrinoThriftNullableColumnSet,
	outputConstraint *trino.TrinoThriftTupleDomain,
	maxSplitCount int32,
	nextToken *trino.TrinoThriftNullableToken,
) (_r *trino.TrinoThriftSplitBatch, _err error) {
	sh, err := s.selectSchemaHandler(ctx, schemaTableName.SchemaName)
	if err != nil {
		return nil, err
	}
	splits, err := sh.GetSplits(ctx, schemaTableName, outputConstraint)
	if err != nil {
		return nil, err
	}

	return &trino.TrinoThriftSplitBatch{
		Splits:    s.encode(splits),
		NextToken: nil,
	}, nil
}

func (s *ClinicService) TrinoGetIndexSplits(
	ctx context.Context,
	schemaTableName *trino.TrinoThriftSchemaTableName,
	indexColumnNames []string,
	outputColumnNames []string,
	keys *trino.TrinoThriftPageResult_,
	outputConstraint *trino.TrinoThriftTupleDomain,
	maxSplitCount int32,
	nextToken *trino.TrinoThriftNullableToken,
) (_r *trino.TrinoThriftSplitBatch, _err error) {
	sh, err := s.selectSchemaHandler(ctx, schemaTableName.SchemaName)
	if err != nil {
		return nil, err
	}
	splits, err := sh.GetIndexSplits(ctx, schemaTableName, outputConstraint, indexColumnNames, keys)
	if err != nil {
		return nil, err
	}

	return &trino.TrinoThriftSplitBatch{
		Splits:    s.encode(splits),
		NextToken: nil,
	}, nil
}

func (s *ClinicService) TrinoGetRows(
	ctx context.Context,
	splitId *trino.TrinoThriftId,
	columns []string,
	maxBytes int64,
	nextToken *trino.TrinoThriftNullableToken,
) (_r *trino.TrinoThriftPageResult_, _err error) {
	var buf bytes.Buffer
	buf.Write(splitId.ID)
	dec := gob.NewDecoder(&buf)

	var split schema.SplitInfo
	if err := dec.Decode(&split); err != nil {
		return nil, err
	}

	sh, err := s.selectSchemaHandler(ctx, split.SchemaTableName.SchemaName)
	if err != nil {
		return nil, err
	}
	return sh.GetRows(ctx, &split, columns)
}