package costexplorer

import (
	"errors"
	"fmt"
	"sort"
	"time"

	"github.com/pingcap/dbaas/billing2/billing_errors"
	"github.com/pingcap/dbaas/billing2/constant"
	"github.com/pingcap/dbaas/billing2/db"
	"github.com/pingcap/dbaas/billing2/db/entity"
	"github.com/pingcap/dbaas/billing2/forms"
	"github.com/pingcap/dbaas/billing2/model/billing/snapshot"
	"github.com/pingcap/dbaas/billing2/utils/timeutils"
	"gorm.io/gorm"
)

// snapshotOverlaySegment is a disjoint interval covered by one applied rerun.
// Rerun windows are day-aligned today, but keeping the interval explicit makes
// the source boundary visible and prevents summing overlapping runs.
type snapshotOverlaySegment struct {
	RunID     uint
	InvoiceID uint
	Start     time.Time
	End       time.Time
}

type snapshotOverlayCandidate struct {
	ID          uint      `gorm:"column:id"`
	RunID       uint      `gorm:"column:run_id"`
	InvoiceID   uint      `gorm:"column:invoice_id"`
	WindowStart time.Time `gorm:"column:window_start"`
	WindowEnd   time.Time `gorm:"column:window_end"`
	CreatedAt   time.Time `gorm:"column:link_created_at"`
}

type snapshotOverlayPlan struct {
	Segments []snapshotOverlaySegment
}

func loadSnapshotOverlayPlan(tenantID, billingConfigID uint, start, end time.Time, invoiceID uint) (*snapshotOverlayPlan, error) {
	return loadSnapshotOverlayPlanWithDB(db.GetReadOnlyDB(), tenantID, billingConfigID, start, end, invoiceID)
}

func loadSnapshotOverlayPlanWithDB(queryDB *gorm.DB, tenantID, billingConfigID uint, start, end time.Time, invoiceID uint) (*snapshotOverlayPlan, error) {
	if tenantID == 0 || billingConfigID == 0 || !end.After(start) || !snapshot.IsSnapshotEnabledForTenant(tenantID) {
		return &snapshotOverlayPlan{}, nil
	}
	if queryDB == nil {
		queryDB = db.GetReadOnlyDB()
	}

	query := queryDB.Model(&entity.BillingSnapshotInvoiceLink{}).
		Select(`
			billing_snapshot_invoice_link.id,
			billing_snapshot_invoice_link.run_id,
			billing_snapshot_invoice_link.invoice_id,
			billing_snapshot_run.window_start,
			billing_snapshot_run.window_end,
			billing_snapshot_invoice_link.created_at as link_created_at`).
		Joins("JOIN billing_snapshot_run ON billing_snapshot_run.run_id = billing_snapshot_invoice_link.run_id").
		Joins("JOIN billing_snapshot_input_bundle ON billing_snapshot_input_bundle.bundle_id = billing_snapshot_run.input_bundle_id").
		Where("billing_snapshot_invoice_link.tenant_id = ?", tenantID).
		Where("billing_snapshot_invoice_link.billing_config_id = ?", billingConfigID).
		Where("billing_snapshot_invoice_link.disposition = ?", entity.BillingSnapshotInvoiceDispositionNormal).
		Where("billing_snapshot_run.status IN ?", []string{
			entity.BillingSnapshotRunStatusSuccess,
			entity.BillingSnapshotRunStatusFrozen,
		}).
		// A link only exists for an authoritative rewrite: either a published
		// rerun (raw_user_unit_costs / rerun_*) or an applied replace on a
		// primary run (legacy_user_unit_costs / base_run). The latter is what
		// ApplyReplaceUnissuedInvoice publishes; excluding it here made the
		// overlay silently fall back to legacy user_unit_costs, so the
		// corrected amount never reached GetInvoiceCosts/invoice views.
		Where("billing_snapshot_run.source_mode IN ?", []string{"legacy_user_unit_costs", "raw_user_unit_costs"}).
		Where("billing_snapshot_input_bundle.bundle_type IN ?", []string{
			entity.BillingSnapshotInputBundleTypeBaseRun,
			entity.BillingSnapshotInputBundleTypeRerunCandidate,
			entity.BillingSnapshotInputBundleTypeRerunPublished,
		}).
		Where("billing_snapshot_run.window_start < ? AND billing_snapshot_run.window_end > ?", end, start)
	if invoiceID != 0 {
		query = query.Where("billing_snapshot_invoice_link.invoice_id = ?", invoiceID)
	}

	var candidates []snapshotOverlayCandidate
	if err := query.Find(&candidates).Error; err != nil {
		return nil, err
	}
	if len(candidates) == 0 {
		return &snapshotOverlayPlan{}, nil
	}

	// The same invoice may now carry several links (one per applied rerun
	// layer), and different invoice types can also overlap. Resolve overlaps
	// deterministically in favour of the most recently applied link (highest
	// link id) so a small rerun applied later overlays a larger one beneath it.
	boundaries := []time.Time{start, end}
	for _, candidate := range candidates {
		candidateStart := maxTime(start, candidate.WindowStart)
		candidateEnd := minTime(end, candidate.WindowEnd)
		if candidateEnd.After(candidateStart) {
			boundaries = append(boundaries, candidateStart, candidateEnd)
		}
	}
	sort.Slice(boundaries, func(i, j int) bool { return boundaries[i].Before(boundaries[j]) })
	uniqueBoundaries := boundaries[:0]
	for _, boundary := range boundaries {
		if len(uniqueBoundaries) == 0 || !boundary.Equal(uniqueBoundaries[len(uniqueBoundaries)-1]) {
			uniqueBoundaries = append(uniqueBoundaries, boundary)
		}
	}

	segments := make([]snapshotOverlaySegment, 0, len(uniqueBoundaries))
	for i := 0; i+1 < len(uniqueBoundaries); i++ {
		segmentStart, segmentEnd := uniqueBoundaries[i], uniqueBoundaries[i+1]
		if !segmentEnd.After(segmentStart) {
			continue
		}
		var selected *snapshotOverlayCandidate
		for j := range candidates {
			candidate := &candidates[j]
			if !candidate.WindowStart.After(segmentStart) && !candidate.WindowEnd.Before(segmentEnd) {
				if selected == nil || candidate.CreatedAt.After(selected.CreatedAt) ||
					(candidate.CreatedAt.Equal(selected.CreatedAt) && candidate.ID > selected.ID) {
					selected = candidate
				}
			}
		}
		if selected == nil {
			continue
		}
		segment := snapshotOverlaySegment{
			RunID:     selected.RunID,
			InvoiceID: selected.InvoiceID,
			Start:     segmentStart,
			End:       segmentEnd,
		}
		if len(segments) > 0 {
			last := &segments[len(segments)-1]
			if last.RunID == segment.RunID && last.InvoiceID == segment.InvoiceID && last.End.Equal(segment.Start) {
				last.End = segment.End
				continue
			}
		}
		segments = append(segments, segment)
	}
	return &snapshotOverlayPlan{Segments: segments}, nil
}

func (p *snapshotOverlayPlan) legacySegments(start, end time.Time) []snapshotOverlaySegment {
	if p == nil || !end.After(start) {
		return nil
	}
	segments := make([]snapshotOverlaySegment, 0, len(p.Segments)+1)
	cursor := start
	for _, snapshotSegment := range p.Segments {
		if snapshotSegment.Start.After(cursor) {
			segments = append(segments, snapshotOverlaySegment{Start: cursor, End: snapshotSegment.Start})
		}
		if snapshotSegment.End.After(cursor) {
			cursor = snapshotSegment.End
		}
	}
	if cursor.Before(end) {
		segments = append(segments, snapshotOverlaySegment{Start: cursor, End: end})
	}
	return segments
}

func maxTime(left, right time.Time) time.Time {
	if left.After(right) {
		return left
	}
	return right
}

func minTime(left, right time.Time) time.Time {
	if left.Before(right) {
		return left
	}
	return right
}

func snapshotOverlayInvoiceID(costQuery *forms.CostExplorerQueryV2) uint {
	if costQuery == nil {
		return 0
	}
	if costQuery.InvoiceID != 0 {
		return costQuery.InvoiceID
	}
	if len(costQuery.InvoiceIDs) == 1 && costQuery.InvoiceIDs[0] > 0 {
		return uint(costQuery.InvoiceIDs[0])
	}
	return 0
}

func snapshotOverlayQuerySupported(costQuery *forms.CostExplorerQueryV2) bool {
	if costQuery == nil {
		return false
	}
	if costQuery.StartDate.IsZero() || costQuery.EndDate.IsZero() {
		return snapshotOverlayInvoiceID(costQuery) != 0
	}
	if !costQuery.EndDate.After(costQuery.StartDate) {
		return false
	}
	// Settlement facts intentionally do not retain binded_package_id. Falling
	// back to legacy for this uncommon filter is safer than returning a partial
	// result.
	if len(costQuery.BindedPackageIDs) != 0 || len(costQuery.InvoiceIDs) > 1 {
		return false
	}
	return true
}

func snapshotFactMatchesCostQuery(fact entity.BillingSnapshotSettlementFact, costQuery *forms.CostExplorerQueryV2) bool {
	if costQuery == nil {
		return true
	}
	labels := normalizeSnapshotFactLabels(fact.Labels)
	for key, value := range costQuery.Labels {
		if labels[key] != value {
			return false
		}
	}
	if costQuery.IgnoreBeta && labels["beta"] == "true" {
		return false
	}
	if len(costQuery.Units) != 0 {
		found := false
		for _, unit := range costQuery.Units {
			if unit == fact.Unit {
				found = true
				break
			}
		}
		if !found {
			return false
		}
	}
	// Cluster filters are represented in labels for settlement facts. The
	// legacy query uses a denormalized column, so this is the equivalent check.
	if clusterID, ok := costQuery.Labels[constant.ClusterLabelKey]; ok && labels[constant.ClusterLabelKey] != clusterID {
		return false
	}
	return true
}

type snapshotOverlayCostKey struct {
	TenantID       uint
	Unit           string
	LabelsKey      string
	ConsumedSource string
	InvoiceID      int
	Internal       bool
}

type snapshotOverlayCostMerger struct {
	ignoreInvoiceID bool
	items           map[snapshotOverlayCostKey]entity.UnitCostSummary
}

func newSnapshotOverlayCostMerger(ignoreInvoiceID bool) *snapshotOverlayCostMerger {
	return &snapshotOverlayCostMerger{
		ignoreInvoiceID: ignoreInvoiceID,
		items:           map[snapshotOverlayCostKey]entity.UnitCostSummary{},
	}
}

func (m *snapshotOverlayCostMerger) add(cost entity.UnitCostSummary) {
	if m == nil {
		return
	}
	if m.ignoreInvoiceID {
		cost.InvoiceID = 0
	}
	key := snapshotOverlayCostKey{
		TenantID:       cost.TenantID,
		Unit:           cost.Unit,
		LabelsKey:      labelsStableKey(cost.Labels),
		ConsumedSource: cost.ConsumedSource,
		InvoiceID:      cost.InvoiceID,
		Internal:       cost.Internal,
	}
	item, ok := m.items[key]
	if !ok {
		m.items[key] = cost
		return
	}
	item.RawPrice = item.RawPrice.Add(cost.RawPrice)
	item.CostPrice = item.CostPrice.Add(cost.CostPrice)
	item.Cost += cost.Cost
	m.items[key] = item
}

func (m *snapshotOverlayCostMerger) costs() []entity.UnitCostSummary {
	keys := make([]snapshotOverlayCostKey, 0, len(m.items))
	for key := range m.items {
		keys = append(keys, key)
	}
	sort.Slice(keys, func(i, j int) bool {
		left, right := keys[i], keys[j]
		if left.TenantID != right.TenantID {
			return left.TenantID < right.TenantID
		}
		if left.Unit != right.Unit {
			return left.Unit < right.Unit
		}
		if left.LabelsKey != right.LabelsKey {
			return left.LabelsKey < right.LabelsKey
		}
		if left.ConsumedSource != right.ConsumedSource {
			return left.ConsumedSource < right.ConsumedSource
		}
		if left.InvoiceID != right.InvoiceID {
			return left.InvoiceID < right.InvoiceID
		}
		return !left.Internal && right.Internal
	})
	result := make([]entity.UnitCostSummary, 0, len(keys))
	for _, key := range keys {
		result = append(result, m.items[key])
	}
	return result
}

func forEachSnapshotOverlayCost(plan *snapshotOverlayPlan, tenantID, billingConfigID uint, costQuery *forms.CostExplorerQueryV2, handle func(entity.UnitCostSummary) error, queryDB *gorm.DB) error {
	if plan == nil {
		return nil
	}
	for _, segment := range plan.Segments {
		if !isDayAlignedSnapshotWindow(segment.Start, segment.End) {
			return billing_errors.NewInvalidInputError(errors.New("snapshot overlay windows must be UTC day-aligned"))
		}
		rows, err := queryDB.Model(&entity.BillingSnapshotSettlementFact{}).
			Where("run_id = ? AND tenant_id = ? AND billing_config_id = ?", segment.RunID, tenantID, billingConfigID).
			Where("billed_date >= ? AND billed_date < ?", timeutils.GetDayDate(segment.Start), timeutils.GetDayDate(segment.End)).
			Order("id ASC").Rows()
		if err != nil {
			return err
		}
		for rows.Next() {
			var fact entity.BillingSnapshotSettlementFact
			if err := queryDB.ScanRows(rows, &fact); err != nil {
				rows.Close()
				return err
			}
			if !snapshotFactMatchesCostQuery(fact, costQuery) {
				continue
			}
			cost := settlementFactToUnitCostSummary(fact)
			cost.Labels = normalizeSnapshotFactLabels(cost.Labels)
			if segment.InvoiceID != 0 && fact.ConsumedSource == entity.ConsumedSource_Invoice {
				cost.InvoiceID = int(segment.InvoiceID)
			}
			if err := handle(cost); err != nil {
				rows.Close()
				return err
			}
		}
		if err := rows.Err(); err != nil {
			rows.Close()
			return err
		}
		rows.Close()
	}
	return nil
}

func forEachSnapshotOverlayDetail(plan *snapshotOverlayPlan, tenantID, billingConfigID uint, invoiceID uint, matches func(entity.BillingSnapshotSettlementFact) bool, handle func(entity.CostsDetailsQuery) error) error {
	if plan == nil {
		return nil
	}
	queryDB := db.GetReadOnlyDB()
	for _, segment := range plan.Segments {
		if !isDayAlignedSnapshotWindow(segment.Start, segment.End) {
			return billing_errors.NewInvalidInputError(errors.New("snapshot overlay windows must be UTC day-aligned"))
		}
		rows, err := queryDB.Model(&entity.BillingSnapshotSettlementFact{}).
			Where("run_id = ? AND tenant_id = ? AND billing_config_id = ? AND is_internal = ?", segment.RunID, tenantID, billingConfigID, false).
			Where("billed_date >= ? AND billed_date < ?", timeutils.GetDayDate(segment.Start), timeutils.GetDayDate(segment.End)).
			Order("id ASC").Rows()
		if err != nil {
			return err
		}
		for rows.Next() {
			var fact entity.BillingSnapshotSettlementFact
			if err := queryDB.ScanRows(rows, &fact); err != nil {
				rows.Close()
				return err
			}
			if matches != nil && !matches(fact) {
				continue
			}
			detail := settlementFactToCostsDetailsQuery(fact)
			detail.Labels = normalizeSnapshotFactLabels(detail.Labels)
			if invoiceID != 0 {
				detail.InvoiceID = invoiceID
			}
			if err := handle(detail); err != nil {
				rows.Close()
				return err
			}
		}
		if err := rows.Err(); err != nil {
			rows.Close()
			return err
		}
		rows.Close()
	}
	return nil
}

func forEachDetailWithSnapshotOverlay(
	tenantID, billingConfigID uint,
	start, end time.Time,
	invoiceID uint,
	matches func(entity.BillingSnapshotSettlementFact) bool,
	legacy func(time.Time, time.Time, func(entity.CostsDetailsQuery) error) error,
	handle func(entity.CostsDetailsQuery) error,
) error {
	if !end.After(start) {
		return nil
	}
	plan, err := loadSnapshotOverlayPlan(tenantID, billingConfigID, start, end, invoiceID)
	if err != nil {
		return err
	}
	if plan == nil || len(plan.Segments) == 0 {
		return legacy(start, end, handle)
	}

	cursor := start
	for _, segment := range plan.Segments {
		if segment.Start.After(cursor) {
			if err := legacy(cursor, segment.Start, handle); err != nil {
				return err
			}
		}
		if err := forEachSnapshotOverlayDetail(&snapshotOverlayPlan{Segments: []snapshotOverlaySegment{segment}}, tenantID, billingConfigID, invoiceID, matches, handle); err != nil {
			return err
		}
		if segment.End.After(cursor) {
			cursor = segment.End
		}
	}
	if cursor.Before(end) {
		return legacy(cursor, end, handle)
	}
	return nil
}

func isDayAlignedSnapshotWindow(start, end time.Time) bool {
	start = start.UTC()
	end = end.UTC()
	return start.Hour() == 0 && start.Minute() == 0 && start.Second() == 0 && start.Nanosecond() == 0 &&
		end.Hour() == 0 && end.Minute() == 0 && end.Second() == 0 && end.Nanosecond() == 0
}

func hasSnapshotOverlay(tenantID, billingConfigID uint, start, end time.Time) bool {
	plan, err := loadSnapshotOverlayPlan(tenantID, billingConfigID, start, end, 0)
	return err == nil && plan != nil && len(plan.Segments) != 0
}

func invoiceWindowForSnapshotOverlay(tenantID, billingConfigID, invoiceID uint) (time.Time, time.Time, bool, error) {
	_, found, err := loadSnapshotInvoiceLink(tenantID, billingConfigID, invoiceID)
	if err != nil || !found {
		return time.Time{}, time.Time{}, found, err
	}
	var invoice entity.Invoice
	if err := db.GetDB().Where("id = ? AND tenant_id = ? AND binded_billing_config_id = ?", invoiceID, tenantID, billingConfigID).First(&invoice).Error; err != nil {
		return time.Time{}, time.Time{}, true, err
	}
	window, err := snapshot.BuildInvoiceWindow(invoice.Type, invoice.InvoiceDate)
	if err != nil {
		return time.Time{}, time.Time{}, true, fmt.Errorf("build invoice window for snapshot overlay: %w", err)
	}
	return window.Start, window.End, true, nil
}