package server

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

// TestFlushBillingOpsPersistsDeltas drives a real write through the HTTP API
// and verifies the billing ops delta lands in the central meta DB.
func TestFlushBillingOpsPersistsDeltas(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())

	s := newTestServer(t)
	s.meta = metaStore
	// Drain any tracker state left by earlier tests so this flush only carries
	// the write performed below.
	metrics.MarkFileOpsFlushed(metrics.SnapshotFileOpsDeltas())
	t.Cleanup(func() { metrics.DeleteTenantCounters("local") })

	ts := httptest.NewServer(s)
	defer ts.Close()
	req, _ := http.NewRequest(http.MethodPut, ts.URL+"/v1/fs/billing-flush.txt", strings.NewReader("hello"))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	_ = resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("write: %d", resp.StatusCode)
	}

	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("flush: %v", err)
	}
	sums, err := metaStore.SumBillingOpsDeltas(context.Background(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatalf("sum deltas: %v", err)
	}
	var writeOps int64
	for _, sum := range sums {
		if sum.TenantID == "local" && sum.Direction == "write" {
			writeOps += sum.Ops
		}
	}
	if writeOps != 1 {
		t.Fatalf("write ops delta = %d, want 1", writeOps)
	}

	// A second flush with no new activity writes nothing (baseline advanced).
	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("second flush: %v", err)
	}
	sums, err = metaStore.SumBillingOpsDeltas(context.Background(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	writeOps = 0
	for _, sum := range sums {
		if sum.TenantID == "local" && sum.Direction == "write" {
			writeOps += sum.Ops
		}
	}
	if writeOps != 1 {
		t.Fatalf("write ops after second flush = %d, want still 1", writeOps)
	}
}

// TestFlushBillingOpsSurvivesTenantEviction pins the idempotency-key shape:
// a tenant evicted and reactivated between flushes restarts its in-memory
// counter from zero, and the post-reactivation flush must still persist (the
// unique key includes the tenant generation).
func TestFlushBillingOpsSurvivesTenantEviction(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())

	s := newTestServer(t)
	s.meta = metaStore
	metrics.MarkFileOpsFlushed(metrics.SnapshotFileOpsDeltas())
	t.Cleanup(func() { metrics.DeleteTenantCounters("local") })

	write := func() {
		t.Helper()
		srv := httptest.NewServer(s)
		defer srv.Close()
		req, _ := http.NewRequest(http.MethodPut, srv.URL+"/v1/fs/evict.txt", strings.NewReader("x"))
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		_ = resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			t.Fatalf("write: %d", resp.StatusCode)
		}
	}

	write()
	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("first flush: %v", err)
	}
	// Evict the tenant's series (as the 15-minute TTL sweep would), then write
	// again: the restarted counter must still be billed.
	metrics.DeleteTenantCounters("local")
	write()
	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("second flush: %v", err)
	}

	sums, err := metaStore.SumBillingOpsDeltas(context.Background(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	var writeOps int64
	for _, sum := range sums {
		if sum.TenantID == "local" && sum.Direction == "write" {
			writeOps += sum.Ops
		}
	}
	if writeOps != 2 {
		t.Fatalf("write ops across eviction = %d, want 2", writeOps)
	}
}

// TestFlushBillingOpsRetriesSameBatchAfterFailure pins that a failed flush
// retries the identical snapshot (same cumulative points), so an ambiguous
// commit is deduped by the idempotency key instead of double-counting.
func TestFlushBillingOpsRetriesSameBatchAfterFailure(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())

	s := newTestServer(t)
	s.meta = metaStore
	metrics.MarkFileOpsFlushed(metrics.SnapshotFileOpsDeltas())
	t.Cleanup(func() { metrics.DeleteTenantCounters("local") })

	ts := httptest.NewServer(s)
	defer ts.Close()
	req, _ := http.NewRequest(http.MethodPut, ts.URL+"/v1/fs/retry.txt", strings.NewReader("x"))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	_ = resp.Body.Close()

	// First flush fails on a canceled context; the batch stays pending.
	canceledCtx, cancel := context.WithCancel(context.Background())
	cancel()
	if err := s.FlushBillingOps(canceledCtx, "pod-test"); err == nil {
		t.Fatal("expected canceled-context flush to fail")
	}
	// More ops land before the retry; they must not alter the pending batch.
	req, _ = http.NewRequest(http.MethodPut, ts.URL+"/v1/fs/retry2.txt", strings.NewReader("y"))
	resp2, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	_ = resp2.Body.Close()

	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("retry flush: %v", err)
	}
	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("final flush: %v", err)
	}

	sums, err := metaStore.SumBillingOpsDeltas(context.Background(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	var writeOps int64
	for _, sum := range sums {
		if sum.TenantID == "local" && sum.Direction == "write" {
			writeOps += sum.Ops
		}
	}
	if writeOps != 2 {
		t.Fatalf("write ops after retry = %d, want exactly 2", writeOps)
	}
}

// TestFlushBillingOpsDropsWedgedBatch pins the escape valve: a batch that
// fails maxBillingFlushFailures times with non-context errors is dropped
// (baselines advance past it), the metric is recorded, and flushing resumes
// afterwards — post-snapshot increments are billed normally.
func TestFlushBillingOpsDropsWedgedBatch(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	testtidb.ResetMetaDB(t, metaStore.DB())

	s := newTestServer(t)
	s.meta = metaStore
	metrics.MarkFileOpsFlushed(metrics.SnapshotFileOpsDeltas())
	t.Cleanup(func() { metrics.DeleteTenantCounters("local") })

	ts := httptest.NewServer(s)
	defer ts.Close()
	write := func(path string) {
		t.Helper()
		req, _ := http.NewRequest(http.MethodPut, ts.URL+path, strings.NewReader("x"))
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			t.Fatal(err)
		}
		_ = resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			t.Fatalf("write %s: %d", path, resp.StatusCode)
		}
	}
	write("/v1/fs/wedged.txt")

	// Kill the DB: inserts fail with sql.ErrDBClosed (a non-context error).
	if err := metaStore.Close(); err != nil {
		t.Fatal(err)
	}
	for i := 1; i <= maxBillingFlushFailures; i++ {
		err := s.FlushBillingOps(context.Background(), "pod-test")
		if err == nil {
			t.Fatalf("flush %d unexpectedly succeeded", i)
		}
		if i < maxBillingFlushFailures && strings.Contains(err.Error(), "dropped") {
			t.Fatalf("flush %d dropped early: %v", i, err)
		}
	}
	if s.billingFlushPending != nil || s.billingFlushFailures != 0 {
		t.Fatalf("after drop: pending=%v failures=%d, want cleared", s.billingFlushPending, s.billingFlushFailures)
	}

	// New increments during the outage must survive the drop: they were
	// recorded after the dropped snapshot.
	write("/v1/fs/after-drop.txt")

	metaStore2, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore2.Close() })
	testtidb.ResetMetaDB(t, metaStore2.DB())
	s.meta = metaStore2

	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("flush after recovery: %v", err)
	}
	sums, err := metaStore2.SumBillingOpsDeltas(context.Background(), time.Now().Add(-time.Hour), time.Now().Add(time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	var writeOps int64
	for _, sum := range sums {
		if sum.TenantID == "local" && sum.Direction == "write" {
			writeOps += sum.Ops
		}
	}
	if writeOps != 1 {
		t.Fatalf("write ops after drop = %d, want exactly 1 (wedged batch dropped, later op billed)", writeOps)
	}
	// Recovery complete: the next flush is a no-op.
	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatal(err)
	}
}

// TestFlushBillingOpsCanceledDoesNotCountTowardDropCap pins that a
// context-canceled flush (shutdown, slow DB) never counts toward the
// pending-batch drop cap.
func TestFlushBillingOpsCanceledDoesNotCountTowardDropCap(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = metaStore.Close() })
	testtidb.ResetMetaDB(t, metaStore.DB())

	s := newTestServer(t)
	s.meta = metaStore
	metrics.MarkFileOpsFlushed(metrics.SnapshotFileOpsDeltas())
	t.Cleanup(func() { metrics.DeleteTenantCounters("local") })

	ts := httptest.NewServer(s)
	defer ts.Close()
	req, _ := http.NewRequest(http.MethodPut, ts.URL+"/v1/fs/cancel.txt", strings.NewReader("x"))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	_ = resp.Body.Close()

	for i := 0; i < maxBillingFlushFailures+5; i++ {
		canceledCtx, cancel := context.WithCancel(context.Background())
		cancel()
		if err := s.FlushBillingOps(canceledCtx, "pod-test"); err == nil {
			t.Fatal("expected canceled flush to fail")
		}
	}
	if s.billingFlushFailures != 0 {
		t.Fatalf("canceled flushes counted as failures: %d", s.billingFlushFailures)
	}
	if err := s.FlushBillingOps(context.Background(), "pod-test"); err != nil {
		t.Fatalf("flush after cancellations: %v", err)
	}
}

// TestBillingOpsFlushLoopRecordsErrorMetric drives the periodic loop against a
// closed meta store and asserts the failure lands on the alerting metric.
func TestBillingOpsFlushLoopRecordsErrorMetric(t *testing.T) {
	metaStore, err := meta.Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	testtidb.ResetMetaDB(t, metaStore.DB())

	s := newTestServer(t)
	s.meta = metaStore
	metrics.MarkFileOpsFlushed(metrics.SnapshotFileOpsDeltas())
	t.Cleanup(func() { metrics.DeleteTenantCounters("local") })

	ts := httptest.NewServer(s)
	defer ts.Close()
	req, _ := http.NewRequest(http.MethodPut, ts.URL+"/v1/fs/loop-metric.txt", strings.NewReader("x"))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	_ = resp.Body.Close()

	// Kill the DB so every loop flush fails, then run the loop briefly.
	if err := metaStore.Close(); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	s.StartBillingOpsFlush(ctx, 10*time.Millisecond)
	time.Sleep(200 * time.Millisecond)
	cancel()

	rec := httptest.NewRecorder()
	metrics.WritePrometheus(rec)
	want := `drive9_service_operations_total{component="server",operation="billing_ops_flush",result="error"}`
	if !strings.Contains(rec.Body.String(), want) {
		t.Fatalf("missing flush error metric %q:\n%s", want, rec.Body.String())
	}
}
