package edge

import (
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

	"github.com/carlosframework/platform/internal/registry"
	"github.com/carlosframework/platform/internal/release"
)

func TestUsageMeterRecordAndDrain(t *testing.T) {
	m := NewUsageMeter()
	rt := registry.Route{Host: "a.example.com", App: "hello", Account: "bab"}
	m.Record(rt, 1, 100, false)
	m.Record(rt, 1, 50, false)
	m.Record(registry.Route{Host: "b.example.com", App: "site"}, 1, 10, false) // Account "" -> release.OpsAccount

	got := m.Drain()
	a := got["a.example.com"]
	if a.Requests != 2 || a.BytesOut != 150 || a.Account != "bab" || a.App != "hello" {
		t.Fatalf("a row = %+v", a)
	}
	if got["b.example.com"].Account != release.OpsAccount {
		t.Fatalf("b row = %+v", got["b.example.com"])
	}
	if again := m.Drain(); len(again) != 0 {
		t.Fatalf("drain did not reset: %+v", again)
	}
}

// A hopped exchange is recorded ONCE — into the totals, as any exchange is —
// and additionally annotated into the hop counters, which are therefore a
// subset of the totals and never an addition to them. A consumer that sums
// BytesOut has already counted every backhauled byte; one that added
// HopBytesOut on top would bill the tenant twice for the same wire bytes.
func TestRecordHopAnnotatesSubset(t *testing.T) {
	m := NewUsageMeter()
	rt := registry.Route{Host: "a.example.com", App: "hello", Account: "bab"}
	m.Record(rt, 1, 100, true)
	m.Record(rt, 1, 50, false)

	row := m.Drain()["a.example.com"]
	if row.Requests != 2 || row.BytesOut != 150 {
		t.Fatalf("row = %+v, want both exchanges counted once in the totals", row)
	}
	if row.HopRequests != 1 || row.HopBytesOut != 100 {
		t.Fatalf("row = %+v, want only the hopped exchange annotated", row)
	}
	if row.HopRequests > row.Requests || row.HopBytesOut > row.BytesOut {
		t.Fatalf("row = %+v — the hop counters must stay inside the totals", row)
	}
	if row.Account != "bab" || row.App != "hello" {
		t.Fatalf("row = %+v, want the hop attributed to the route's own account and app", row)
	}
}

func TestUsageMeterConcurrent(t *testing.T) {
	m := NewUsageMeter()
	rt := registry.Route{Host: "a.example.com", App: "hello", Account: "bab"}
	var wg sync.WaitGroup
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < 100; j++ {
				m.Record(rt, 1, 10, false)
			}
		}()
	}
	wg.Wait()
	row := m.Drain()["a.example.com"]
	if row.Requests != 5000 || row.BytesOut != 50000 {
		t.Fatalf("row = %+v", row)
	}
}

func TestCountingWriterCountsBytes(t *testing.T) {
	var n int64
	rec := httptest.NewRecorder()
	w := &countingWriter{ResponseWriter: rec, n: &n}
	w.WriteHeader(http.StatusTeapot)
	w.Write([]byte("hello "))
	w.Write([]byte("world"))
	if n != 11 {
		t.Fatalf("counted %d bytes", n)
	}
	if rec.Code != http.StatusTeapot || rec.Body.String() != "hello world" {
		t.Fatalf("forwarding broken: %d %q", rec.Code, rec.Body.String())
	}
}

func TestCountingWriterObservesStatusAndType(t *testing.T) {
	var n int64
	rec := httptest.NewRecorder()
	w := &countingWriter{ResponseWriter: rec, n: &n}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(103) // informational: forwarded, not latched
	w.WriteHeader(http.StatusBadGateway)
	w.WriteHeader(http.StatusOK) // superfluous second final status: ignored
	if w.status != http.StatusBadGateway || w.ctype != "text/html; charset=utf-8" {
		t.Fatalf("latched %d %q", w.status, w.ctype)
	}
	// Implicit 200 at first Write.
	rec2 := httptest.NewRecorder()
	w2 := &countingWriter{ResponseWriter: rec2, n: &n}
	w2.Write([]byte("x"))
	if w2.status != http.StatusOK {
		t.Fatalf("implicit status = %d", w2.status)
	}
	// Nothing written at all: status stays 0 (the caller treats it as a
	// bodiless 200 that is never a pageview and never an error).
	w3 := &countingWriter{ResponseWriter: httptest.NewRecorder(), n: &n}
	if w3.status != 0 {
		t.Fatal("untouched writer must report 0")
	}
}