package nextgenmetering

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	sdkcommon "github.com/pingcap/metering_sdk/common"
	deltalake "github.com/tidbcloud/metering/delta-rs-golang"
	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/constants"
	"github.com/tidbcloud/metering/pkg/common/metrics"
	"github.com/tidbcloud/metering/pkg/common/models"
	"github.com/tidbcloud/metering/pkg/common/utils"
	"github.com/tidbcloud/metering/pkg/common/utils/deltautils"
	"go.uber.org/zap"
)

// NextGenMetering handles the next generation metering service
type NextGenMetering struct {
	conf            *config.ServiceConfig
	loader          *MeteringLoader
	tableManager    *deltautils.TableManager
	log             *zap.Logger
	taskName        string
	skuGenerator    *SKUGenerator
	handlerRegistry *SKUHandlerRegistry
}

// NewNextGenMetering creates a new NextGenMetering instance
func NewNextGenMetering(conf *config.ServiceConfig, log *zap.Logger) (*NextGenMetering, error) {
	if conf.NextGenMeteringConfig == nil {
		return nil, fmt.Errorf("NextGenMeteringConfig is required")
	}

	loader := NewMeteringLoader(conf, log)
	if err := loader.Init(); err != nil {
		return nil, fmt.Errorf("failed to initialize metering loader: %w", err)
	}

	// Get or initialize global table manager
	tableManager := deltautils.GetGlobalTableManager()
	if tableManager == nil {
		var err error
		tableManager, err = deltautils.InitGlobalTableManager(conf, log)
		if err != nil {
			return nil, fmt.Errorf("failed to initialize global table manager: %w", err)
		}
	}

	logger := log.With(zap.String("module", "nextgenmetering"))
	skuGen := NewSKUGenerator(conf, log)
	return &NextGenMetering{
		conf:            conf,
		loader:          loader,
		tableManager:    tableManager,
		log:             logger,
		taskName:        "nextgen-metering",
		skuGenerator:    skuGen,
		handlerRegistry: NewSKUHandlerRegistry(skuGen, log),
	}, nil
}

// createDeltaTable creates a new delta table
func (n *NextGenMetering) CreateDeltaTable() error {
	return n.tableManager.CreateTable(context.Background())
}

// openDeltaTable opens an existing delta table
func (n *NextGenMetering) openDeltaTable() error {
	return n.tableManager.OpenTable(context.Background())
}

// Close releases resources associated with the NextGenMetering instance
// Note: The global table manager is NOT closed here, as it may be shared by other tasks
func (n *NextGenMetering) Close() error {
	var errors []error

	if n.loader != nil {
		n.log.Info("Closing loader")
		if err := n.loader.Close(); err != nil {
			n.log.Error("Failed to close loader", zap.Error(err))
			errors = append(errors, fmt.Errorf("failed to close loader: %w", err))
		}
		n.loader = nil
	}

	if len(errors) > 0 {
		n.log.Error("Errors occurred during close", zap.Int("error_count", len(errors)))
		return fmt.Errorf("errors during close: %v", errors)
	}

	n.log.Info("NextGenMetering resources closed successfully")
	return nil
}

// processMinute processes metering data for a single minute and returns the generated SKU records
func (n *NextGenMetering) ProcessMinute(ctx context.Context, timestamp int64) ([]models.BillingReportV2Deltalake, error) {
	n.log.Debug("Processing minute", zap.Int64("timestamp", timestamp))

	// Read metering data for this minute
	meteringDataList, err := n.loader.LoadMeteringDataByTimestamp(ctx, timestamp, "")
	if err != nil {
		return nil, fmt.Errorf("failed to read metering data: %w", err)
	}

	if len(meteringDataList) == 0 {
		n.log.Debug("No metering data found for timestamp", zap.Int64("timestamp", timestamp))
		return []models.BillingReportV2Deltalake{}, nil
	}

	// Generate SKU records
	skuRecords, err := n.generateBillingRecords(ctx, meteringDataList, timestamp)
	if err != nil {
		return nil, fmt.Errorf("failed to generate billing records: %w", err)
	}

	n.log.Debug("Successfully processed minute",
		zap.Int64("timestamp", timestamp),
		zap.Int("records", len(skuRecords)))

	return skuRecords, nil
}

// generateBillingRecords converts metering data to BillingReportV2Deltalake records
func (n *NextGenMetering) generateBillingRecords(ctx context.Context, meteringDataList []*sdkcommon.MeteringData, timestamp int64) ([]models.BillingReportV2Deltalake, error) {
	var billingRecords []models.BillingReportV2Deltalake

	for _, meteringData := range meteringDataList {
		if n.conf.NextGenMeteringConfig.Plan == models.PlanBYOC && meteringData.Category != models.NextGenCategoryWorkerNode {
			n.log.Debug("BYOC mode: skipping non-worker-node category",
				zap.String("category", meteringData.Category),
				zap.Int64("timestamp", timestamp))
			continue
		}
		if meteringData.Category == models.NextGenCategoryReservedNode {
			aggregatedDataPoints := n.aggregateReservedNodeDataPoints(meteringData.Data)
			for _, item := range aggregatedDataPoints {
				clusterRecords, err := n.generateBillingRecordsForCluster(ctx, item.clusterID, meteringData.SharedPoolID, timestamp, []map[string]interface{}{item.dataPoint}, meteringData.Category)
				if err != nil {
					n.log.Error("Failed to generate billing records for reserved-node aggregate",
						zap.String("cluster_id", item.clusterID),
						zap.String("self_id", meteringData.SelfID),
						zap.String("category", meteringData.Category),
						zap.Int64("timestamp", timestamp),
						zap.Error(err))
					return nil, err
				}

				billingRecords = append(billingRecords, clusterRecords...)
			}
			continue
		}

		// Process each data point in the metering data
		// Each data point may have its own cluster_id
		for _, dataPoint := range meteringData.Data {
			// Extract cluster_id from the data point
			clusterID, ok := dataPoint[models.NextGenMeteringLogicalClusterID].(string)
			if !ok && meteringData.Category != models.NextGenCategoryWorkerNode {
				n.log.Warn("Data point missing cluster_id, use shared pool processing",
					zap.String("self_id", meteringData.SelfID),
					zap.String("category", meteringData.Category),
					zap.Int64("timestamp", timestamp))
			}

			// Generate records for this specific cluster and data point
			clusterRecords, err := n.generateBillingRecordsForCluster(ctx, clusterID, meteringData.SharedPoolID, timestamp, []map[string]interface{}{dataPoint}, meteringData.Category)
			if err != nil {
				n.log.Error("Failed to generate billing records for cluster",
					zap.String("cluster_id", clusterID),
					zap.String("self_id", meteringData.SelfID),
					zap.String("category", meteringData.Category),
					zap.Int64("timestamp", timestamp),
					zap.Error(err))
				// if one cluster fails, we log and return error
				return nil, err
			}

			billingRecords = append(billingRecords, clusterRecords...)
		}
	}

	n.log.Info("Generated billing records",
		zap.Int("total_records", len(billingRecords)),
		zap.Int64("timestamp", timestamp))

	return billingRecords, nil
}

type reservedNodeAggregate struct {
	clusterID string
	dataPoint map[string]interface{}
}

func (n *NextGenMetering) aggregateReservedNodeDataPoints(dataPoints []map[string]interface{}) []reservedNodeAggregate {
	type aggregateState struct {
		total uint64
		unit  string
	}

	states := make(map[string]*aggregateState)
	order := make([]string, 0, len(dataPoints))

	for _, dataPoint := range dataPoints {
		clusterID, _ := dataPoint[models.NextGenMeteringLogicalClusterID].(string)
		state, exists := states[clusterID]
		if !exists {
			state = &aggregateState{}
			states[clusterID] = state
			order = append(order, clusterID)
		}

		usageData, ok := n.skuGenerator.ConvertToMeteringValue(dataPoint["usage"])
		if !ok || usageData == nil {
			continue
		}

		state.total += usageData.Value
		if state.unit == "" {
			state.unit = usageData.Unit
		}
	}

	result := make([]reservedNodeAggregate, 0, len(order))
	for _, clusterID := range order {
		state := states[clusterID]
		if state == nil || state.total == 0 {
			continue
		}
		unit := state.unit
		if unit == "" {
			unit = "vCPU"
		}

		result = append(result, reservedNodeAggregate{
			clusterID: clusterID,
			dataPoint: map[string]interface{}{
				models.NextGenMeteringNodeType: "all",
				"usage": map[string]interface{}{
					"value": state.total,
					"unit":  unit,
				},
			},
		})
	}

	return result
}

// generateBillingRecordsForCluster generates billing records for a specific cluster
// This method fetches metadata for the cluster and enriches the metering data
func (n *NextGenMetering) generateBillingRecordsForCluster(ctx context.Context, clusterID string, sharedPoolID string, timestamp int64, dataPoints []map[string]interface{}, category string) ([]models.BillingReportV2Deltalake, error) {
	var billingRecords []models.BillingReportV2Deltalake

	// Use the new dealSKU method for each data point
	for _, dataPoint := range dataPoints {
		records, err := n.dealSKU(ctx, clusterID, sharedPoolID, timestamp, dataPoint, category)
		if err != nil {
			return nil, err
		}
		billingRecords = append(billingRecords, records...)
	}

	return billingRecords, nil
}

// dealSKU creates billing records for specific cluster based on SKU rules
func (n *NextGenMetering) dealSKU(ctx context.Context, clusterID string, sharedPoolID string, timestamp int64, dataPoint map[string]interface{}, category string) ([]models.BillingReportV2Deltalake, error) {
	now := time.Now()
	var sharedMetaData, metaData, backupMetaData *sdkcommon.MetaData
	var err error

	// Load metadata based on cluster type
	if clusterID == "" || clusterID == models.NextGenMeteringSystemID {
		sharedMetaData, err = n.loader.LoadMetaDataByTypeWithCategory(ctx, sharedPoolID, sdkcommon.MetaTypeSharedpool, "", timestamp)
		if err != nil {
			n.log.Error("Failed to load shared meta data",
				zap.String("shared_pool_id", sharedPoolID),
				zap.Int64("timestamp", timestamp),
				zap.Error(err))
			return nil, err
		}
		if sharedMetaData == nil {
			metrics.MetricsNextGenMeteringSkipCount.WithLabelValues(n.conf.Vendor, category, constants.MetricsNextGenReasonMetadataMissing).Inc()
			n.log.Warn("Shared metadata not found, skipping processing",
				zap.String("shared_pool_id", sharedPoolID),
				zap.Int64("timestamp", timestamp))
			return []models.BillingReportV2Deltalake{}, nil
		}
	} else {
		metaData, err = n.loader.LoadMetaDataByTypeWithCategory(ctx, clusterID, sdkcommon.MetaTypeLogic, "", timestamp)
		if err != nil {
			n.log.Error("Failed to load meta data",
				zap.String("cluster_id", clusterID),
				zap.Int64("timestamp", timestamp),
				zap.Error(err))
			return nil, err
		}
		if metaData == nil {
			metrics.MetricsNextGenMeteringSkipCount.WithLabelValues(n.conf.Vendor, category, constants.MetricsNextGenReasonMetadataMissing).Inc()
			n.log.Warn("Metadata not found, skipping processing",
				zap.String("cluster_id", clusterID),
				zap.Int64("timestamp", timestamp))
			return []models.BillingReportV2Deltalake{}, nil
		}

		// Load backup category metadata for logic cluster
		backupMetaData, err = n.loader.LoadMetaDataByTypeWithCategory(ctx, clusterID, sdkcommon.MetaTypeLogic, models.NextGenCategoryBackup, timestamp)
		if err != nil {
			n.log.Error("Failed to load backup meta data",
				zap.String("cluster_id", clusterID),
				zap.Int64("timestamp", timestamp),
				zap.Error(err))
			// Don't return error, just log warning and continue without backup metadata
		}
		if backupMetaData == nil {
			n.log.Debug("Backup metadata not found, backup-storage-v2 SKU will be skipped",
				zap.String("cluster_id", clusterID),
				zap.Int64("timestamp", timestamp))
		}
	}

	// Extract metadata information
	var tenantID uint64
	var enrichmentData map[string]interface{}
	var backupEnrichmentData map[string]interface{}
	var cloudProvider, region, plan string

	if metaData != nil && metaData.Metadata != nil {
		enrichmentData = metaData.Metadata
		if tenant, ok := metaData.Metadata[models.NextGenMetaTenantID].(string); ok {
			if tid, parseErr := strconv.ParseUint(tenant, 10, 64); parseErr == nil {
				tenantID = tid
			}
		}
		if regionVal, ok := metaData.Metadata[models.NextGenMeteringRegion].(string); ok {
			region = regionVal
		}
		cloudProvider = n.conf.Vendor
		plan = resolvePlan(enrichmentData, n.conf.NextGenMeteringConfig.Plan)
	} else if sharedMetaData != nil && sharedMetaData.Metadata != nil {
		enrichmentData = sharedMetaData.Metadata
		if tenant, ok := sharedMetaData.Metadata[models.NextGenMetaTenantID].(string); ok {
			if tid, parseErr := strconv.ParseUint(tenant, 10, 64); parseErr == nil {
				tenantID = tid
			}
		}
		if regionVal, ok := sharedMetaData.Metadata[models.NextGenMeteringRegion].(string); ok {
			region = regionVal
		}
		cloudProvider = n.conf.Vendor
		plan = resolvePlan(enrichmentData, n.conf.NextGenMeteringConfig.Plan)
	} else {
		// if no metadata found, we log and skip processing
		metrics.MetricsNextGenMeteringSkipCount.WithLabelValues(n.conf.Vendor, category, constants.MetricsNextGenReasonMetadataMissing).Inc()
		n.log.Warn("No metadata found for cluster or shared pool, skipping processing",
			zap.String("cluster_id", clusterID),
			zap.String("shared_pool_id", sharedPoolID),
			zap.String("category", category),
			zap.Int64("timestamp", timestamp))
		return []models.BillingReportV2Deltalake{}, nil
	}
	if backupMetaData != nil && backupMetaData.Metadata != nil {
		backupEnrichmentData = backupMetaData.Metadata
	}
	plan = normalizeSKUPlan(plan, enrichmentData)

	gate := n.conf.NextGenMeteringConfig.GetFeatureGateFilter(constants.FeatureGateNextGenSkipDeletedMetadata)
	if gate.IsFullSwitch(timestamp) {
		//if meta has deleted flag, skip processing
		if val, ok := enrichmentData[models.NextGenMetaDeleted]; ok {
			if deleted, ok := val.(bool); ok && deleted {
				metrics.MetricsNextGenMeteringSkipCount.WithLabelValues(n.conf.Vendor, category, constants.MetricsNextGenReasonMetadataDeleted).Inc()
				n.log.Info("Metadata marked as deleted, skipping processing",
					zap.String("cluster_id", clusterID),
					zap.String("shared_pool_id", sharedPoolID),
					zap.String("category", category),
					zap.Int64("timestamp", timestamp))
				return []models.BillingReportV2Deltalake{}, nil
			}
		}
	}

	// Generate SKU records based on data point content and category
	return n.generateSKURecords(dataPoint, category, clusterID, sharedPoolID, tenantID, timestamp,
		cloudProvider, region, plan, enrichmentData, backupEnrichmentData, now)
}

func isPremiumReservedMode(plan string, enrichmentData map[string]interface{}) bool {
	if strings.EqualFold(plan, "premium_reserved") || strings.EqualFold(plan, "reserved") {
		return true
	}
	if capacityMode, ok := enrichmentData[models.NextGenMetaCapacityMode].(string); ok {
		return strings.EqualFold(capacityMode, "TOPOLOGY_DRIVEN")
	}
	return false
}

func normalizeSKUPlan(plan string, enrichmentData map[string]interface{}) string {
	if isPremiumReservedMode(plan, enrichmentData) {
		return models.PlanReserved
	}
	return plan
}

// generateSKURecords generates billing records based on data point and metadata
func (n *NextGenMetering) generateSKURecords(dataPoint map[string]interface{}, category, clusterID, sharedPoolID string, tenantID uint64, timestamp int64, cloudProvider, region, plan string, enrichmentData map[string]interface{}, backupMetaData map[string]interface{}, now time.Time) ([]models.BillingReportV2Deltalake, error) {
	plan = normalizeSKUPlan(plan, enrichmentData)

	if category == models.NextGenCategoryReservedNode {
		records, err := n.skuGenerator.GenerateReservedNodeSKUs(dataPoint, clusterID, sharedPoolID, tenantID, timestamp, cloudProvider, region, plan, enrichmentData, now)
		if err != nil {
			return nil, err
		}
		return records, nil
	}

	handler, ok := n.handlerRegistry.GetHandler(plan, category)
	if !ok {
		sourceName := "unknown"
		if sourceNameVal, ok := dataPoint[models.NextGenMeteringSourceName].(string); ok {
			sourceName = sourceNameVal
		}
		metrics.MetricsNextGenMeteringSkipCount.WithLabelValues(n.conf.Vendor, category, constants.MetricsNextGenReasonSKURuleNotMatch).Inc()
		n.log.Warn("No handler for plan+category, skipping data point",
			zap.String("plan", plan),
			zap.String("category", category),
			zap.String("source_name", sourceName),
			zap.String("cluster_id", clusterID),
			zap.Int64("timestamp", timestamp))
		return []models.BillingReportV2Deltalake{}, nil
	}

	params := &GenerateSKUParams{
		ClusterID:      clusterID,
		SharedPoolID:   sharedPoolID,
		TenantID:       tenantID,
		Timestamp:      timestamp,
		CloudProvider:  cloudProvider,
		Region:         region,
		Plan:           plan,
		EnrichmentData: enrichmentData,
		BackupMetaData: backupMetaData,
		Now:            now,
	}

	records, err := handler.GenerateSKUs(dataPoint, params)
	if err != nil {
		return nil, err
	}
	return records, nil
}

// resolvePlan determines the effective plan: config-level for byoc, metadata-level otherwise
func resolvePlan(enrichmentData map[string]interface{}, configPlan string) string {
	// byoc only relies on config plan, ignore metadata plan
	if configPlan == models.PlanBYOC {
		return models.PlanBYOC
	}
	// others use metadata plan if exists, otherwise use config plan
	if planVal, ok := enrichmentData[models.NextGenMetaPlan].(string); ok && planVal != "" {
		return planVal
	}
	if configPlan != "" {
		return configPlan
	}
	return models.PlanPremium
}

// WriteToDeltaTable writes billing records to delta table
func (n *NextGenMetering) WriteToDeltaTable(records []models.BillingReportV2Deltalake, saveMode string) error {
	err := n.openDeltaTable()
	if err != nil {
		return err
	}
	return n.tableManager.WriteToDeltaTable(context.Background(), records, saveMode)
}

// GetTableVersion returns the current version of the Delta Lake table
// This method reuses the already opened table if available
func (n *NextGenMetering) GetTableVersion(ctx context.Context) (int64, error) {
	err := n.openDeltaTable()
	if err != nil {
		return 0, err
	}
	return n.tableManager.GetTableVersion(ctx)
}

// QueryMaxTimestampWithPartitionFilter returns the maximum cost_time_stamp since the provided partition time using delta-rs query
func (n *NextGenMetering) QueryMaxTimestampWithPartitionFilter(ctx context.Context, partitionTime int64) (int64, error) {
	if n.conf.NextGenMeteringConfig == nil {
		return 0, fmt.Errorf("NextGenMeteringConfig is required")
	}

	if n.conf.NextGenMeteringConfig.DeltaTablePath == "" {
		return 0, fmt.Errorf("delta table path not configured")
	}

	// Generate hour dimension for partition pruning - use 30 days before partitionTime to ensure we capture data
	// even when there's no metering data at the exact partitionTime
	hourDimension := utils.GenerateHourDimensionWithOffset(partitionTime, -30)
	query := fmt.Sprintf(`
		SELECT MAX(cost_time_stamp) AS max_timestamp
		FROM metering_table
		WHERE hour_dimension >= '%s'
	`, hourDimension)

	queryOptions := &deltalake.QueryOptions{TableAlias: "metering_table"}

	rows, err := n.tableManager.Query(ctx, query, queryOptions)
	if err != nil {
		n.log.Error("failed to execute delta-rs max timestamp query",
			zap.String("table_path", n.conf.NextGenMeteringConfig.DeltaTablePath),
			zap.String("hour_dimension", hourDimension),
			zap.Error(err))
		return 0, fmt.Errorf("failed to query max timestamp: %w", err)
	}

	if len(rows) == 0 {
		n.log.Info("no rows returned from delta-rs max timestamp query",
			zap.String("table_path", n.conf.NextGenMeteringConfig.DeltaTablePath),
			zap.String("hour_dimension", hourDimension))
		return 0, nil
	}

	value, ok := rows[0]["max_timestamp"]
	if !ok || value == nil {
		n.log.Info("max timestamp not found in delta-rs query result",
			zap.String("table_path", n.conf.NextGenMeteringConfig.DeltaTablePath),
			zap.String("hour_dimension", hourDimension))
		return 0, nil
	}

	maxTimestamp, err := convertToInt64(value)
	if err != nil {
		n.log.Error("failed to convert delta-rs max timestamp value",
			zap.String("table_path", n.conf.NextGenMeteringConfig.DeltaTablePath),
			zap.String("hour_dimension", hourDimension),
			zap.Any("value", value),
			zap.Error(err))
		return 0, fmt.Errorf("failed to convert max timestamp: %w", err)
	}

	n.log.Info("successfully retrieved max timestamp using delta-rs",
		zap.String("table_path", n.conf.NextGenMeteringConfig.DeltaTablePath),
		zap.String("hour_dimension", hourDimension),
		zap.Int64("max_timestamp", maxTimestamp))

	return maxTimestamp, nil
}

func convertToInt64(value interface{}) (int64, error) {
	switch v := value.(type) {
	case int64:
		return v, nil
	case int32:
		return int64(v), nil
	case int:
		return int64(v), nil
	case float64:
		return int64(v), nil
	case float32:
		return int64(v), nil
	case uint64:
		return int64(v), nil
	case uint32:
		return int64(v), nil
	case string:
		parsed, err := strconv.ParseInt(v, 10, 64)
		if err != nil {
			return 0, err
		}
		return parsed, nil
	default:
		return 0, fmt.Errorf("unsupported type %T for max timestamp", value)
	}
}

// TableExists checks if the Delta Lake table exists at the configured path
func (n *NextGenMetering) TableExists(ctx context.Context) (bool, error) {
	return n.tableManager.TableExists(ctx)
}