package query

import (
	"bytes"
	"context"
	"fmt"
	"sort"

	"github.com/pingcap/log"
	"github.com/tidbcloud/metrics-server/central/db/entity"
	"github.com/tidbcloud/metrics-server/central/metrics/prom"
	"github.com/tidbcloud/metrics-server/central/monitor_data_migration"
	"github.com/tidbcloud/metrics-server/central/utils"
	"github.com/tidbcloud/metrics-server/idl/pbgen/proto/core"
	"github.com/tidbcloud/metrics-server/internal/datasource"
	"github.com/tidbcloud/metrics-server/internal/grpc/grpcli"
	"github.com/tidbcloud/metrics-server/internal/otel_collector/receiver/metadb"
	"go.uber.org/zap"
)

type Manager struct {
}

func NewManager() *Manager {
	return &Manager{}
}

func (m *Manager) QueryMetrics(ctx context.Context, projectID uint64, cluster *core.Cluster, useClusterMetrics bool) ([]byte, error) {
	type result struct {
		value prom.Metric
		err   error
		index int
	}

	var allMetricQueries []datasource.MetricQuerierFunc
	if useClusterMetrics {
		allMetricQueries = registryClusterMetrics(cluster)
	} else {
		allMetricQueries = registryMetrics(cluster)
	}

	projectMigrationStatus := monitor_data_migration.GetProjectMigrationStatusFromContext(ctx)
	tidbClusterInfo := datasource.TidbClusterInfo{
		Vendor:    cluster.CloudProviderName,
		Region:    cluster.RegionName,
		TenantID:  entity.ID(cluster.TenantId),
		ProjectID: entity.ID(cluster.ProjectId),
		ClusterID: entity.ID(cluster.Id),
	}
	ctx, err := datasource.SetMetricDataAddressToContextWithTidbClusterInfoV2(ctx, projectMigrationStatus, tidbClusterInfo)
	if err != nil {
		return nil, fmt.Errorf("failed to call datasource.SetMetricDataAddressToContextWithTidbClusterInfoV2: %s", err.Error())
	}

	resultsCh := make(chan result, len(allMetricQueries))
	for idx, query := range allMetricQueries {
		go func(idx int, query datasource.MetricQuerierFunc) {
			m, err := query(ctx)
			resultsCh <- result{
				value: m,
				err:   err,
				index: idx,
			}
		}(idx, query)
	}
	results := make([]result, 0, len(allMetricQueries))
	for range allMetricQueries {
		r := <-resultsCh
		if r.err != nil {
			log.Warn("GetMeasurements API failed to query metrics",
				zap.Uint64("clusterID", cluster.Id),
				zap.Error(r.err))
		}
		results = append(results, r)
	}
	sort.Slice(results, func(i, j int) bool {
		return results[i].index < results[j].index
	})
	outputData := bytes.NewBuffer(make([]byte, 0, 10*1024))
	for _, r := range results {
		if r.err == nil {
			_ = r.value.EncodePrometheusFormat(outputData)
		}
	}

	return outputData.Bytes(), nil
}

func (m *Manager) QueryProjectMetrics(ctx context.Context, projectID uint64) ([]byte, error) {
	resp, err := grpcli.ListClusters(ctx, &core.ListClustersReq{
		Filter: &core.ClusterFilter{
			ProjectIds: []uint64{projectID},
			// Statuses:   []core.ClusterStatus{core.ClusterStatus_CLUSTER_STATUS_NORMAL, core.ClusterStatus_CLUSTER_STATUS_SCALING},
			Deleted:   utils.BoolPtr(false),
			IsDevTier: utils.BoolPtr(false),
		},
	})
	if err != nil {
		return nil, fmt.Errorf("failed to list clusters under project %d : %w", projectID, err)
	}
	clusters := make(map[uint64]*core.Cluster)
	for _, c := range resp.Clusters {
		if !c.IsDevTier && metadb.StatusExpected(c.Status) {
			log.Info(fmt.Sprintf("status %s found for cluster %s", c.Status.String(), c.Name))
			clusters[c.Id] = c
		}
	}

	type clusterMetricsResp struct {
		clusterID uint64
		data      []byte
		err       error
	}
	clusterMetricsCh := make(chan clusterMetricsResp, len(clusters))
	for _, cluster := range clusters {
		go func(cluster *core.Cluster) {
			data, err := m.QueryMetrics(ctx, projectID, cluster, true)
			clusterMetricsCh <- clusterMetricsResp{
				clusterID: cluster.Id,
				data:      data,
				err:       err,
			}
		}(cluster)
	}

	outputData := bytes.NewBuffer(make([]byte, 0, 100*1024))
	for i := 0; i < len(clusters); i++ {
		r := <-clusterMetricsCh
		if r.err != nil {
			log.Warn("failed to query cluster metrics", zap.Uint64("clusterID", r.clusterID), zap.Error(r.err))
			continue
		}
		outputData.Write(r.data)
	}

	return outputData.Bytes(), nil
}