package database

import (
	"context"
	"fmt"

	"github.com/glebarez/sqlite"
	log "github.com/sirupsen/logrus"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/constants"
	"github.com/tidbcloud/observability/o11ymetering/pkg/common/models"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/logger"
)

type GormRW struct {
	DB *Database
}

type Database struct {
	base *gorm.DB
}

func (s GormRW) GetTask(ctx context.Context, id string) (*models.StatisticTask, error) {
	task := &models.StatisticTask{}
	err := s.DB.base.First(task, "id = ?", id).Error
	if err != nil {
		return nil, err
	} else {
		return task, nil
	}
}

func (s GormRW) SaveStatisticRecordsForTask(ctx context.Context, taskID string, records []*models.StatisticRecord) ([]*models.StatisticRecord, error) {
	task := models.StatisticTask{
		ID: taskID,
	}
	err := s.DB.base.Transaction(func(tx *gorm.DB) error {
		if len(records) > 0 {
			err := tx.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(records, 900).Error
			if err != nil {
				return err
			}
		}
		err := tx.Model(task).Update("state", constants.StateCalculated).Error
		return err
	})
	return records, err
}

func (s GormRW) CopyRecords(ctx context.Context, source string, region string) error {
	sourceDB, err := gorm.Open(sqlite.Open(source+"?_busy_timeout=5000&_txlock=immediate"), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Error),
	})
	if err != nil {
		log.Errorf("open localSqliteFile failed, err != %s", err.Error())
		return err
	}

	defer func() {
		if sourceDB != nil {
			if d, e := sourceDB.DB(); e == nil && d != nil {
				d.Close()
			}
		}
	}()

	err = sourceDB.AutoMigrate(models.StatisticRecord{})
	if err != nil {
		log.Errorf(err.Error())
		return err
	}

	// get total count of records
	total := int64(0)
	err = sourceDB.Model(&models.StatisticRecord{}).
		Where("region = ?", region).
		Count(&total).Error
	if err != nil {
		log.Errorf("query total records failed, err = %s", err)
		return err
	}
	log.Infof("copy statistic records, total = %d", total)

	// paginate to avoid loading all data at once
	pageSize := 900
	return s.DB.base.Transaction(func(tx *gorm.DB) error {
		for offset := 0; offset < int(total); offset += pageSize {
			var pageRecords []*models.StatisticRecord

			// 	query page records
			err := sourceDB.Model(&models.StatisticRecord{}).
				Where("region = ?", region).
				Offset(offset).
				Limit(pageSize).
				Find(&pageRecords).Error

			if err != nil {
				log.Errorf("query page records failed, offset = %d, err = %s", offset, err)
				return err
			}

			if len(pageRecords) == 0 {
				break
			}

			// copy records to merged db
			err = tx.Clauses(clause.OnConflict{DoNothing: true}).
				CreateInBatches(pageRecords, len(pageRecords)).Error
			if err != nil {
				log.Errorf("copy records to merged db failed, err = %s", err)
				return err
			}

			log.Infof("copied %d records (offset: %d)", len(pageRecords), offset)
		}
		return nil
	})
}

func (s GormRW) LatestMergedRecordPeriod(ctx context.Context, region string) (string, error) {
	records := make([]*models.StatisticRecord, 0)

	err := s.DB.base.Model(&models.StatisticRecord{}).
		Where("region = ?", region).
		Order("period desc").
		Order("id asc").
		Limit(1).
		Find(&records).Error
	if err != nil {
		return "", err
	}

	if len(records) == 0 {
		return "", nil
	}
	return records[0].Period, nil
}

func (s GormRW) DeleteRecordsByPeriod(ctx context.Context, periods []string) ([]*models.StatisticRecord, error) {
	panic("implement me")
}

func (s GormRW) BatchCreateTask(ctx context.Context, tasks []*models.StatisticTask) ([]*models.StatisticTask, error) {
	if len(tasks) > 0 {
		err := s.DB.base.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(tasks, len(tasks)).Error
		if err != nil {
			return tasks, err
		}
	}
	return tasks, nil
}

func (s GormRW) PickEarliestPendingTasks(ctx context.Context, region string, limit int) ([]*models.StatisticTask, error) {
	tasks := make([]*models.StatisticTask, 0)

	err := s.DB.base.Model(&models.StatisticTask{}).
		Where("state = ?", "Pending").
		Where("region = ?", region).
		Order("period asc").
		Order("id asc").
		Limit(limit).
		Find(&tasks).Error
	if err != nil {
		return tasks, err
	}
	return tasks, err
}

func (s GormRW) UpdateTaskState(ctx context.Context, id string, originalState constants.StatisticState, state constants.StatisticState) error {
	task := models.StatisticTask{
		ID: id,
	}
	err := s.DB.base.Model(&task).Update("state", state).Error
	if err != nil {
		return err
	}
	return nil
}

func (s GormRW) GetLastedTaskPeriod(ctx context.Context, region string) (string, error) {
	tasks := make([]*models.StatisticTask, 0)

	err := s.DB.base.Model(&models.StatisticTask{}).
		Where("region = ?", region).
		Order("period desc").
		Limit(1).
		Find(&tasks).Error

	if err != nil || len(tasks) == 0 {
		return "", err
	}
	return tasks[0].Period, nil
}

func (s GormRW) DeleteTasksByPeriod(ctx context.Context, periods []string) ([]*models.StatisticTask, error) {
	panic("implement me")
}

func (s GormRW) IfReportExists(ctx context.Context, timestamp int32, region string) (bool, error) {
	var reports []*models.MonthlyStatisticReport
	err := s.DB.base.Model(models.MonthlyStatisticReport{}).
		Where("timestamp = ?", timestamp).
		Where("region = ?", region).
		Limit(1).
		Find(&reports).Error

	if err != nil {
		return false, err
	}
	if len(reports) > 0 {
		return true, nil
	} else {
		return false, nil
	}
}

func (s GormRW) IfCostReady(ctx context.Context, timestamp int32, region string) (bool, error) {
	var costs []*models.O11yCost
	err := s.DB.base.Table("o11y_cost").
		Where("timestamp = ?", timestamp).
		Where("region = ?", region).
		Limit(1).
		Find(&costs).Error

	if err != nil {
		return false, err
	}
	if len(costs) > 0 {
		return true, nil
	} else {
		return false, nil
	}
}

func (s GormRW) GetMonthlyCost(ctx context.Context, region string, timestamp int32, cloudService string) (float64, error) {
	var value float64
	err := s.DB.base.Table("o11y_cost").
		Where("timestamp = ?", timestamp).
		Where("region = ?", region).
		Where("cloud_service = ?", cloudService).
		Select("cost").Find(&value).Error

	return value, err
}

func (s GormRW) InsertReport(ctx context.Context, reports []*models.MonthlyStatisticReport) error {
	err := s.DB.base.Transaction(func(tx *gorm.DB) error {
		innerErr := tx.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(reports, len(reports)).Error
		return innerErr
	})
	return err
}

func (s GormRW) MonthlyStatisticSum(ctx context.Context, region string, startPeriod string, endPeriod string, statisticName constants.StatisticName, category constants.Category) (float64, error) {
	var result float64
	err := s.DB.base.Table("statistic_records").
		Where("region = ?", region).
		Where("statistic_name = ?", statisticName).
		Where("category = ?", category).
		Where("period >= ?", startPeriod).
		Where("period < ?", endPeriod).
		Select("COALESCE(sum(value), 0)").
		Find(&result).Error

	return result, err
}

func (s GormRW) MonthlyTenantStatistic(ctx context.Context, region string, startPeriod string, endPeriod string, statisticName constants.StatisticName) ([]*models.MonthlyTenantStatistic, error) {
	var result []*models.MonthlyTenantStatistic
	err := s.DB.base.Table("statistic_records").
		Where("region = ?", region).
		Where("statistic_name = ?", statisticName).
		Where("period >= ?", startPeriod).
		Where("period < ?", endPeriod).
		Select("region, statistic_name, category, sub_category, metering_group, env, tenant, sum(value) as sum_value").
		Group("category, sub_category, metering_group, env, tenant").Find(&result).Error

	return result, err
}

// CalculateMonthlyProportionSum performs monthly metering proportion sum calculation.
func (s GormRW) CalculateMonthlyProportionSum(ctx context.Context, targetTimestamp int64, targetMonth string) ([]models.MonthlyProportionSum, error) {
	log.Infof("Starting metering statistics calculation for %s (Timestamp: %d)\n", targetMonth, targetTimestamp)

	// --- Step 1: Aggregate statistical data by provider and metering_group ---
	log.Info("Step 1: Aggregating costs by provider and metering_group...")
	var reports []models.MonthlyStatisticReport
	if err := s.DB.base.Where("timestamp = ?", targetTimestamp).Find(&reports).Error; err != nil {
		return nil, fmt.Errorf("failed to query raw statistical reports: %w", err)
	}

	aggregatedData := make(map[string]map[constants.MeteringGroup]float64)

	for _, report := range reports {
		provider := report.Vendor

		if _, ok := aggregatedData[provider]; !ok {
			aggregatedData[provider] = make(map[constants.MeteringGroup]float64)
		}
		aggregatedData[provider][report.MeteringGroup] += report.Cost
	}

	log.Info("Step 1 completed. Aggregated results example:")
	for p, mgData := range aggregatedData {
		for mg, cost := range mgData {
			log.Infof("  Provider: %s, MeteringGroup: %s, TotalCost: %.2f\n", p, mg, cost)
		}
	}

	log.Info("Step 2: Calculating amplification ratios and total values...")

	providerFactors := make(map[string]struct {
		AmplificationFactor            float64
		ProviderTotalCostForProportion float64
	})

	for provider, mgData := range aggregatedData {
		var dedicatedDataplaneCost float64
		var dedicatedRegionalCost float64
		var providerTotalCost float64

		for mg, cost := range mgData {
			if mg == "dedicated_dataplane" {
				dedicatedDataplaneCost = cost
			} else if mg == "dedicated_regional" {
				dedicatedRegionalCost = cost
			}
			providerTotalCost += cost
		}

		var amplificationFactor float64
		if dedicatedDataplaneCost != 0 {
			amplificationFactor = (dedicatedDataplaneCost + dedicatedRegionalCost) / dedicatedDataplaneCost
		} else {
			amplificationFactor = 1.0
		}

		providerFactors[provider] = struct {
			AmplificationFactor            float64
			ProviderTotalCostForProportion float64
		}{
			AmplificationFactor:            amplificationFactor,
			ProviderTotalCostForProportion: providerTotalCost,
		}

		log.Infof("  Provider: %s\n", provider)
		log.Infof("    DedicatedDataplane Cost: %.2f\n", dedicatedDataplaneCost)
		log.Infof("    DedicatedRegional Cost: %.2f\n", dedicatedRegionalCost)
		log.Infof("    Amplification Factor: %.5f\n", amplificationFactor)
		log.Infof("    Total Value for Proportion Denominator (Excluding Regional): %.2f\n", providerTotalCost)
	}
	fmt.Println("Step 2 completed.")

	// --- Step 3: Aggregate data and insert into monthly_proportion_sum table ---
	fmt.Println("---")
	fmt.Println("Step 3: Aggregating data and inserting into monthly_proportion_sum table...")

	var rawReportsForProportion []models.MonthlyStatisticReport
	if err := s.DB.base.Where("timestamp = ? AND metering_group != ?", targetTimestamp, "dedicated_regional").Find(&rawReportsForProportion).Error; err != nil {
		return nil, fmt.Errorf("failed to query raw reports for proportion calculation: %w", err)
	}

	type FinalGroupKey struct {
		Timestamp     int64
		Month         string
		Provider      string
		Env           string
		MeteringGroup string
		Tenant        string
	}
	finalAggregatedCosts := make(map[FinalGroupKey]float64)

	for _, report := range rawReportsForProportion {
		provider := report.Vendor

		key := FinalGroupKey{
			Timestamp:     int64(report.Timestamp),
			Month:         targetMonth,
			Provider:      provider,
			Env:           string(report.Env),
			MeteringGroup: string(report.MeteringGroup),
			Tenant:        report.Tenant,
		}
		finalAggregatedCosts[key] += report.Cost
	}

	var recordsToInsert = make([]models.MonthlyProportionSum, 0, len(finalAggregatedCosts))
	for key, groupCost := range finalAggregatedCosts {
		factorData := providerFactors[key.Provider]
		var proportion float64

		switch {
		case key.MeteringGroup == "dedicated_dataplane":
			if factorData.ProviderTotalCostForProportion != 0 {
				proportion = groupCost * factorData.AmplificationFactor / factorData.ProviderTotalCostForProportion
			} else {
				proportion = 0
			}
		default:
			if factorData.ProviderTotalCostForProportion != 0 {
				proportion = groupCost / factorData.ProviderTotalCostForProportion
			} else {
				proportion = 0
			}
		}

		recordsToInsert = append(recordsToInsert, models.MonthlyProportionSum{
			Timestamp:     key.Timestamp,
			Month:         key.Month,
			Provider:      key.Provider,
			Env:           key.Env,
			MeteringGroup: key.MeteringGroup,
			Tenant:        key.Tenant,
			Proportion:    proportion,
		})
	}

	if len(recordsToInsert) > 0 {
		log.Infof("Preparing to insert %d records into monthly_proportion_sum table...\n", len(recordsToInsert))
		if err := s.DB.base.CreateInBatches(recordsToInsert, 100).Error; err != nil {
			return nil, fmt.Errorf("failed to insert monthly_proportion_sum data: %w", err)
		}
		log.Infof("Data successfully inserted into monthly_proportion_sum.")
	} else {
		log.Infof("No data to insert into monthly_proportion_sum.")
	}

	log.Infof("Metering statistics calculation completed.")
	return recordsToInsert, nil
}

// InsertRecordsToMergeDB inserts statistic records directly into the merge database
// This is used by the replay functionality to write records directly to merge_statistic table
func InsertRecordsToMergeDB(ctx context.Context, mergeRW interface{}, records []*models.StatisticRecord) error {
	if len(records) == 0 {
		return nil
	}

	// Cast to GormRW to access the underlying database
	gormRW, ok := mergeRW.(*GormRW)
	if !ok {
		return fmt.Errorf("merge reader/writer is not a GormRW instance")
	}

	// Insert records directly into merge database
	// Use DoNothing on conflict to handle duplicate records gracefully
	err := gormRW.DB.base.Transaction(func(tx *gorm.DB) error {
		return tx.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(records, 900).Error
	})

	return err
}