package billing

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"time"

	"github.com/mem9-ai/drive9/pkg/meta"
	"github.com/mem9-ai/drive9/pkg/metrics"
)

// MetaStore is the subset of *meta.Store the billing worker uses. The
// interface exists so collector/sender tests can run without a database.
type MetaStore interface {
	ListTenantStorageTierUsage(ctx context.Context) ([]meta.TenantStorageTierUsage, error)
	InsertBillingStorageSnapshots(ctx context.Context, snapshots []meta.BillingStorageSnapshot) error
	LatestBillingStorageSnapshotAtOrBefore(ctx context.Context, tenantID string, ts time.Time) (meta.BillingStorageSnapshot, error)
	ListBillingSnapshotTenants(ctx context.Context, before time.Time) ([]meta.BillingSnapshotTenant, error)
	DeleteBillingStorageSnapshotsBefore(ctx context.Context, before time.Time) (int64, error)
	SumBillingOpsDeltas(ctx context.Context, from, to time.Time) ([]meta.BillingOpsSum, error)
	DeleteBillingOpsDeltasBefore(ctx context.Context, before time.Time) (int64, error)
	BillingCollectCheckpoint(ctx context.Context) (time.Time, bool, error)
	AdvanceBillingCollectCheckpoint(ctx context.Context, lastHour time.Time) error
	ListIncompleteBillingStorageHours(ctx context.Context, since time.Time) ([]meta.BillingIncompleteHour, error)
	WriteBillingUsageWithOutbox(ctx context.Context, ledger []meta.BillingUsageLedgerRow, outbox []meta.BillingOutboxRow) error
	// WriteBillingHourWithOutbox closes one hour atomically: ledger rows,
	// outbox rows, and the checkpoint advance commit in one transaction.
	WriteBillingHourWithOutbox(ctx context.Context, ledger []meta.BillingUsageLedgerRow, outbox []meta.BillingOutboxRow, checkpointHour time.Time) error
	ClaimBillingOutboxBatch(ctx context.Context, owner string, lease time.Duration, limit int) ([]meta.BillingOutboxRow, error)
	MarkBillingOutboxSent(ctx context.Context, costID, owner string) error
	MarkBillingOutboxFailed(ctx context.Context, costID, owner string, nextAttemptAt time.Time, sendErr error) error
	DeleteTerminalBillingOutboxBefore(ctx context.Context, before time.Time) (int64, error)
	PruneBillingLedgerRows(ctx context.Context, before time.Time) (int64, error)
	BillingOutboxStats(ctx context.Context) (meta.BillingOutboxStats, error)
}

// Billing reads ops usage from the billing_ops_delta table: every
// drive9-server pod flushes its in-process ops counter deltas there, so the
// billing data path never depends on an external metrics (o11y) stack.
// traffic_path is a billing dimension: db9 (through the drive9 HTTP data
// plane) and s3 (direct-to-S3 presign/multipart) ops bill as separate
// messages per tenant per hour.

// org label values produced by the server's local/dev mode; they are not
// billable TiDB Cloud organizations and must never reach the ledger.
var unbillableOrgIDs = map[string]bool{"": true, "guest": true, "unknown": true}

// tenant label values produced when the server cannot attribute a request;
// they are not billable drive9 tenants.
var unbillableTenantIDs = map[string]bool{"": true, "unknown": true}

// opsSKUForPath maps direction × traffic_path to the ops SKU. db9 (through
// the drive9 HTTP data plane) is the non-pooled SKU; s3 (direct
// presign/multipart) is the pooled one.
var opsSKUForPath = map[[2]string]string{
	{"read", "db9"}:  meta.BillingUnitReadOps,
	{"write", "db9"}: meta.BillingUnitWriteOps,
	{"read", "s3"}:   meta.BillingUnitPooledReadOps,
	{"write", "s3"}:  meta.BillingUnitPooledWriteOps,
}

// BillingMeasure is the SQS message contract from the billing design doc.
// Ops messages are per drive9 tenant per unit per traffic path per hour;
// storage messages are per tenant per tier (no path dimension). TenantID is
// the numeric TiDB Cloud organization ID required by the consumer, and the
// drive9 tenant travels in labels as resource_id/resource_type=fs.
type BillingMeasure struct {
	TenantID      uint64            `json:"tenant_id"` // numeric TiDB Cloud organization ID
	CostID        string            `json:"cost_id"`
	Cost          int64             `json:"cost"`
	CostTimestamp int64             `json:"cost_time_stamp"`
	Labels        map[string]string `json:"labels"`
	Unit          string            `json:"unit"`
}

// unitFor composes the emitted BillingMeasure.unit from a SKU stem:
// fs-{sku}:{provider}:{region} (e.g. fs-read-ops:aws:us-west-2).
// Provider/region also remain in labels; the composite unit lets the consumer
// route on the unit string alone.
func (c *Collector) unitFor(sku string) string {
	return "fs-" + sku + ":" + c.cfg.Provider + ":" + c.cfg.Region
}

// Collector closes complete UTC hours into the usage ledger and SQS outbox.
type Collector struct {
	cfg  Config
	meta MetaStore
	now  func() time.Time
}

func newCollector(cfg Config, store MetaStore, now func() time.Time) *Collector {
	if now == nil {
		now = time.Now
	}
	return &Collector{cfg: cfg.withDefaults(), meta: store, now: now}
}

// SnapshotOnce records one storage snapshot per tenant. Leader-only.
func (c *Collector) SnapshotOnce(ctx context.Context) error {
	usages, err := c.meta.ListTenantStorageTierUsage(ctx)
	if err != nil {
		return fmt.Errorf("list tenant storage tier usage: %w", err)
	}
	capturedAt := c.now().UTC()
	snapshots := make([]meta.BillingStorageSnapshot, 0, len(usages))
	for _, usage := range usages {
		snapshots = append(snapshots, meta.BillingStorageSnapshot{
			TenantID:   usage.TenantID,
			OrgID:      usage.TiDBCloudOrgID,
			DBBytes:    usage.DBBytes,
			S3Bytes:    usage.S3Bytes,
			CapturedAt: capturedAt,
		})
	}
	if err := c.meta.InsertBillingStorageSnapshots(ctx, snapshots); err != nil {
		return fmt.Errorf("insert storage snapshots: %w", err)
	}
	metrics.RecordTenantOperationCount("", "billing", "storage_snapshot", "ok")
	return nil
}

// PruneOnce prunes the billing tables. Snapshot and ops-delta rows are cut at
// the collection checkpoint minus the incomplete-retry window (so a lagging
// collector never loses unprocessed data); terminal outbox rows and
// zero/stale ledger rows are cut at now - 2× the window. The 2× factor is the
// load-bearing gap: retry only looks at hours >= now - window, so anything
// older than now - 2×window can never be resurrected.
func (c *Collector) PruneOnce(ctx context.Context) error {
	last, ok, err := c.meta.BillingCollectCheckpoint(ctx)
	if err != nil {
		return fmt.Errorf("read collect checkpoint for prune: %w", err)
	}
	if !ok {
		return nil
	}
	cutoff := last.Add(-c.cfg.IncompleteRetryWindow)
	deleted, err := c.meta.DeleteBillingStorageSnapshotsBefore(ctx, cutoff)
	if err != nil {
		return fmt.Errorf("prune storage snapshots: %w", err)
	}
	if deleted > 0 {
		metrics.RecordGauge("billing", "storage_snapshots_pruned", float64(deleted))
	}
	deleted, err = c.meta.DeleteBillingOpsDeltasBefore(ctx, cutoff)
	if err != nil {
		return fmt.Errorf("prune ops deltas: %w", err)
	}
	if deleted > 0 {
		metrics.RecordGauge("billing", "ops_deltas_pruned", float64(deleted))
	}
	terminalCutoff := c.now().UTC().Add(-2 * c.cfg.IncompleteRetryWindow)
	deleted, err = c.meta.DeleteTerminalBillingOutboxBefore(ctx, terminalCutoff)
	if err != nil {
		return fmt.Errorf("prune terminal outbox rows: %w", err)
	}
	if deleted > 0 {
		metrics.RecordGauge("billing", "outbox_terminal_pruned", float64(deleted))
	}
	// Zero-quantity ledger rows (written so retried hours converge) carry no
	// billing information; prune them once they are beyond 2x the retry
	// window. Stale incomplete rows past the retry window go with them.
	deleted, err = c.meta.PruneBillingLedgerRows(ctx, terminalCutoff)
	if err != nil {
		return fmt.Errorf("prune ledger rows: %w", err)
	}
	if deleted > 0 {
		metrics.RecordGauge("billing", "ledger_rows_pruned", float64(deleted))
	}
	return nil
}

// CollectDueHours closes every complete UTC hour that is older than
// HourCloseDelay and newer than the checkpoint. Hours that fail to collect
// (e.g. VM unreachable) stop the sweep and are retried on the next tick;
// the checkpoint only advances past fully processed hours. On first run the
// checkpoint is initialized so the sweep starts at the newest closable hour
// and never backfills earlier history.
func (c *Collector) CollectDueHours(ctx context.Context) error {
	last, ok, err := c.meta.BillingCollectCheckpoint(ctx)
	if err != nil {
		return fmt.Errorf("read collect checkpoint: %w", err)
	}
	closableEnd := c.now().UTC().Add(-c.cfg.HourCloseDelay).Truncate(time.Hour)
	if !ok {
		// The newest closable hour starts at closableEnd-1h; initialize the
		// checkpoint one hour earlier so the sweep starts exactly there.
		if err := c.meta.AdvanceBillingCollectCheckpoint(ctx, closableEnd.Add(-2*time.Hour)); err != nil {
			return fmt.Errorf("init collect checkpoint: %w", err)
		}
		return nil
	}
	for start := last.Add(time.Hour); !start.After(closableEnd.Add(-time.Hour)); start = start.Add(time.Hour) {
		if err := c.collectHour(ctx, start); err != nil {
			return err
		}
	}
	return nil
}

// collectHour closes one UTC hour. Ops come from the billing_ops_delta table
// (flushed by every server pod); storage capacity (MiB) comes from the
// freshest snapshot at each hour's end. The hour's ledger rows, outbox rows,
// and the checkpoint advance commit in one meta transaction.
func (c *Collector) collectHour(ctx context.Context, start time.Time) error {
	start = start.UTC().Truncate(time.Hour)
	end := start.Add(time.Hour)

	sums, err := c.meta.SumBillingOpsDeltas(ctx, start, end)
	if err != nil {
		metrics.RecordTenantOperationCount("", "billing", "collect_hour", "meta_error")
		return fmt.Errorf("sum ops for hour %s: %w", start.Format(time.RFC3339), err)
	}

	ledger, outbox, unmapped, err := c.buildOpsRecords(start, sums)
	if err != nil {
		return err
	}
	if unmapped > 0 {
		metrics.RecordGauge("billing", "collect_unmapped_tenants", float64(unmapped))
	}

	storageLedger, storageOutbox, err := c.buildStorageRecords(ctx, start, end)
	if err != nil {
		return err
	}
	ledger = append(ledger, storageLedger...)
	outbox = append(outbox, storageOutbox...)

	if err := c.meta.WriteBillingHourWithOutbox(ctx, ledger, outbox, start); err != nil {
		metrics.RecordTenantOperationCount("", "billing", "collect_hour", "meta_error")
		return fmt.Errorf("write hour %s records: %w", start.Format(time.RFC3339), err)
	}
	metrics.RecordTenantOperationCount("", "billing", "collect_hour", "ok")
	return nil
}

// buildOpsRecords maps aggregated ops deltas into per-tenant ledger rows and
// SQS payloads. Rows without a billable tenant/organization identity are
// counted as unmapped and skipped.
func (c *Collector) buildOpsRecords(start time.Time, sums []meta.BillingOpsSum) (ledger []meta.BillingUsageLedgerRow, outbox []meta.BillingOutboxRow, unmapped int, err error) {
	for _, s := range sums {
		orgNum, parseErr := strconv.ParseUint(s.OrgID, 10, 64)
		if unbillableOrgIDs[s.OrgID] || unbillableTenantIDs[s.TenantID] || parseErr != nil || orgNum == 0 {
			unmapped++
			continue
		}
		// direction × traffic_path selects the SKU: db9 (through the drive9
		// HTTP data plane) is the non-pooled ops SKU, s3 (direct
		// presign/multipart) is the pooled one. Unknown combos are skipped
		// with a metric — billing is the last boundary before charging, so
		// never guess a SKU.
		sku, ok := opsSKUForPath[[2]string{s.Direction, s.TrafficPath}]
		if !ok {
			// Split the result label by which dimension is off so a firing
			// counter is self-explaining.
			result := "unknown_traffic_path_skipped"
			if s.Direction != "read" && s.Direction != "write" {
				result = "unknown_direction_skipped"
			}
			metrics.RecordTenantOperationCount("", "billing", "collect_hour", result)
			continue
		}
		unit := c.unitFor(sku)
		if s.Ops <= 0 {
			continue
		}
		labels := map[string]string{
			"service":       "drive9",
			"provider":      c.cfg.Provider,
			"region":        c.cfg.Region,
			"resource_id":   s.TenantID,
			"resource_type": "fs",
			"traffic_path":  s.TrafficPath,
		}
		costID := CostID(orgNum, unit, start, c.cfg.Provider, c.cfg.Region, s.TenantID, s.TrafficPath)
		row := meta.BillingUsageLedgerRow{
			CostID:      costID,
			TenantID:    s.TenantID,
			OrgID:       s.OrgID,
			Unit:        unit,
			PeriodStart: start,
			Quantity:    s.Ops,
			Labels:      labels,
			Status:      meta.BillingLedgerStatusComplete,
		}
		ledger = append(ledger, row)
		payload, merr := json.Marshal(billingMeasureFor(row, orgNum))
		if merr != nil {
			return nil, nil, 0, merr
		}
		outbox = append(outbox, meta.BillingOutboxRow{CostID: costID, Payload: payload})
	}
	return ledger, outbox, unmapped, nil
}

// buildStorageRecords closes one hour of storage per tenant. Hours without
// sufficient snapshot coverage are written as incomplete ledger rows and
// produce no outbox message — they must not be billed as zero.
func (c *Collector) buildStorageRecords(ctx context.Context, start, end time.Time) (ledger []meta.BillingUsageLedgerRow, outbox []meta.BillingOutboxRow, err error) {
	tenants, err := c.meta.ListBillingSnapshotTenants(ctx, end)
	if err != nil {
		return nil, nil, fmt.Errorf("list snapshot tenants: %w", err)
	}
	for _, tenant := range tenants {
		tenantLedger, tenantOutbox, err := c.buildStorageRecordsForTenant(ctx, tenant, start, end)
		if err != nil {
			return nil, nil, err
		}
		ledger = append(ledger, tenantLedger...)
		outbox = append(outbox, tenantOutbox...)
	}
	return ledger, outbox, nil
}

// buildStorageRecordsForTenant closes one hour of storage for one tenant,
// using the tenant's newest snapshot at or before the hour end (carried by
// ListBillingSnapshotTenants — no per-tenant re-query). The cost quantity is
// the point-in-time capacity in MiB; the billing side prices it per MiB per
// hour. Both storage SKUs always get a ledger row (zero capacity is a
// complete zero, so a retried hour converges its sibling SKU instead of
// leaving a stale incomplete marker); outbox messages are only created for
// positive quantities. Hours whose freshest snapshot is missing or stale are
// incomplete and never billed as zero.
func (c *Collector) buildStorageRecordsForTenant(ctx context.Context, tenant meta.BillingSnapshotTenant, start, end time.Time) (ledger []meta.BillingUsageLedgerRow, outbox []meta.BillingOutboxRow, err error) {
	orgNum, parseErr := strconv.ParseUint(tenant.OrgID, 10, 64)
	if unbillableOrgIDs[tenant.OrgID] || unbillableTenantIDs[tenant.TenantID] || parseErr != nil || orgNum == 0 {
		return nil, nil, nil
	}
	if tenant.CapturedAt.IsZero() || end.Sub(tenant.CapturedAt) > c.cfg.SnapshotStaleTolerance {
		return c.incompleteStorageRows(tenant, start), nil, nil
	}
	labels := map[string]string{
		"service":       "drive9",
		"provider":      c.cfg.Provider,
		"region":        c.cfg.Region,
		"resource_id":   tenant.TenantID,
		"resource_type": "fs",
	}
	for _, row := range []struct {
		sku      string
		quantity int64
	}{
		// Capacity in MiB from the freshest snapshot, rounded UP: a tenant with
		// any bytes at all bills at least 1 MiB for the hour.
		{meta.BillingUnitStoragePerf, mibCeil(tenant.DBBytes)},
		{meta.BillingUnitStoragePooled, mibCeil(tenant.S3Bytes)},
	} {
		unit := c.unitFor(row.sku)
		costID := CostID(orgNum, unit, start, c.cfg.Provider, c.cfg.Region, tenant.TenantID)
		ledgerRow := meta.BillingUsageLedgerRow{
			CostID:      costID,
			TenantID:    tenant.TenantID,
			OrgID:       tenant.OrgID,
			Unit:        unit,
			PeriodStart: start,
			Quantity:    row.quantity,
			Labels:      labels,
			Status:      meta.BillingLedgerStatusComplete,
		}
		ledger = append(ledger, ledgerRow)
		if row.quantity <= 0 {
			continue
		}
		payload, merr := json.Marshal(billingMeasureFor(ledgerRow, orgNum))
		if merr != nil {
			return nil, nil, merr
		}
		outbox = append(outbox, meta.BillingOutboxRow{CostID: costID, Payload: payload})
	}
	return ledger, outbox, nil
}

// incompleteStorageRows records the hour as not billable for both storage
// units so the retry pass (RetryIncompleteHours) and audits can find it; no
// outbox message is created.
func (c *Collector) incompleteStorageRows(tenant meta.BillingSnapshotTenant, start time.Time) []meta.BillingUsageLedgerRow {
	orgNum, err := strconv.ParseUint(tenant.OrgID, 10, 64)
	if err != nil {
		return nil
	}
	labels := map[string]string{
		"service":       "drive9",
		"provider":      c.cfg.Provider,
		"region":        c.cfg.Region,
		"resource_id":   tenant.TenantID,
		"resource_type": "fs",
	}
	var rows []meta.BillingUsageLedgerRow
	for _, sku := range []string{meta.BillingUnitStoragePerf, meta.BillingUnitStoragePooled} {
		rows = append(rows, meta.BillingUsageLedgerRow{
			CostID:      CostID(orgNum, c.unitFor(sku), start, c.cfg.Provider, c.cfg.Region, tenant.TenantID),
			TenantID:    tenant.TenantID,
			OrgID:       tenant.OrgID,
			Unit:        c.unitFor(sku),
			PeriodStart: start,
			Quantity:    0,
			Labels:      labels,
			Status:      meta.BillingLedgerStatusIncomplete,
		})
	}
	metrics.RecordTenantOperationCount("", "billing", "incomplete_storage_hour", "ok")
	return rows
}

// RetryIncompleteHours re-evaluates incomplete storage hours inside a bounded
// recent window (IncompleteRetryWindow, default 24h). This covers the one
// recoverable gap shape: snapshots that were captured inside the hour but
// committed after the hour closed (commit/boundary races). Hours whose gap
// comes from worker downtime can never gain retroactive coverage — snapshots
// are point-in-time — so they stay incomplete and alerted, and are never
// billed as zero.
func (c *Collector) RetryIncompleteHours(ctx context.Context) error {
	since := c.now().UTC().Add(-c.cfg.IncompleteRetryWindow)
	hours, err := c.meta.ListIncompleteBillingStorageHours(ctx, since)
	if err != nil {
		return fmt.Errorf("list incomplete storage hours: %w", err)
	}
	for _, hour := range hours {
		if ctx.Err() != nil {
			return nil
		}
		start := hour.PeriodStart.UTC().Truncate(time.Hour)
		tenant := meta.BillingSnapshotTenant{TenantID: hour.TenantID, OrgID: hour.OrgID}
		// Incomplete hours re-check coverage with the tenant's latest
		// snapshot at the hour end; a missing one keeps the hour incomplete.
		snap, err := c.meta.LatestBillingStorageSnapshotAtOrBefore(ctx, hour.TenantID, start.Add(time.Hour))
		switch {
		case err == nil:
			tenant.DBBytes, tenant.S3Bytes, tenant.CapturedAt = snap.DBBytes, snap.S3Bytes, snap.CapturedAt
		case errors.Is(err, meta.ErrNotFound):
			// no snapshot yet
		default:
			return fmt.Errorf("load retry snapshot for tenant %s: %w", hour.TenantID, err)
		}
		ledger, outbox, err := c.buildStorageRecordsForTenant(ctx, tenant, start, start.Add(time.Hour))
		if err != nil {
			return err
		}
		complete := false
		for _, row := range ledger {
			if row.Status == meta.BillingLedgerStatusComplete {
				complete = true
				break
			}
		}
		if !complete {
			continue
		}
		// The ledger upsert and outbox insert are idempotent per cost_id, so
		// completing an hour through the retry path converges with a normal
		// close and never duplicates messages.
		if err := c.meta.WriteBillingUsageWithOutbox(ctx, ledger, outbox); err != nil {
			return fmt.Errorf("rewrite completed hour %s tenant %s: %w", start.Format(time.RFC3339), hour.TenantID, err)
		}
		metrics.RecordTenantOperationCount("", "billing", "retry_incomplete_hour", "completed")
	}
	return nil
}

// mibBytes is the storage capacity unit: snapshot bytes are reported to the
// billing side as capacity in MiB, rounded up.
const mibBytes = int64(1) << 20

// mibCeil rounds byte counts up to whole MiB (zero stays zero):
// a tenant with any non-zero capacity bills at least 1 MiB.
func mibCeil(bytes int64) int64 {
	if bytes <= 0 {
		return 0
	}
	return (bytes-1)/mibBytes + 1
}

// billingMeasureFor renders the SQS payload for a complete ledger row.
// TenantID carries the numeric TiDB Cloud organization ID required by the
// consumer; the drive9 tenant travels in labels.
func billingMeasureFor(row meta.BillingUsageLedgerRow, orgNum uint64) BillingMeasure {
	return BillingMeasure{
		TenantID:      orgNum,
		CostID:        row.CostID,
		Cost:          row.Quantity,
		CostTimestamp: row.PeriodStart.Unix(),
		Labels:        row.Labels,
		Unit:          row.Unit,
	}
}