package main

import (
	"encoding/json"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"
)

func TestStorageBytesForDB(t *testing.T) {
	dir := t.TempDir()
	db := filepath.Join(dir, "space.11ch.at.db")
	write := func(path string, n int) {
		t.Helper()
		if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(path, make([]byte, n), 0o600); err != nil {
			t.Fatal(err)
		}
	}
	write(db, 1000)
	write(db+"-wal", 200)
	write(filepath.Join(dir, "space.11ch.at-blobs", "aa"), 5000)
	write(filepath.Join(dir, "space.11ch.at-blobs", "nested", "bb"), 300)
	// A sibling instance's files must not leak into the count.
	write(filepath.Join(dir, "other.11ch.at.db"), 9999)

	if got := storageBytesForDB(db); got != 6500 {
		t.Errorf("storageBytesForDB = %d, want 6500 (db+wal+blobs)", got)
	}
	// Missing everything → 0, no error.
	if got := storageBytesForDB(filepath.Join(dir, "absent.db")); got != 0 {
		t.Errorf("absent instance = %d, want 0", got)
	}
}

func TestMeterOnceRecordsReadings(t *testing.T) {
	db := openTestCloudDB(t)
	idA := insertSpaceRow(t, db, "a.11ch.at", 1000)
	idB := insertSpaceRow(t, db, "b.11ch.at", 1000)
	// A pending create (NULL host) is never swept.
	if _, err := db.sql.Exec(`INSERT INTO spaces(host, caller_id, requester, created_at) VALUES(NULL,'c','r',1)`); err != nil {
		t.Fatal(err)
	}

	bytesByHost := map[string]int64{"a.11ch.at": 12345, "b.11ch.at": 0}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /instances/{host}/stats", func(w http.ResponseWriter, r *http.Request) {
		h := r.PathValue("host")
		if h == "b.11ch.at" {
			w.WriteHeader(http.StatusInternalServerError) // stats down for b
			return
		}
		writeJSON(w, http.StatusOK, map[string]any{"storage_bytes": bytesByHost[h], "setup_complete": true, "members": 3})
	})
	sock := fakeControl(t, mux.ServeHTTP)
	s := &cloudServer{db: db, control: controlClient(sock), controlToken: "t"}

	if metered := s.meterOnce(); metered != 1 {
		t.Fatalf("metered = %d, want 1 (a recorded, b skipped)", metered)
	}
	spaces, err := db.listSpaces("c", "r")
	if err != nil {
		t.Fatal(err)
	}
	byHost := map[string]Space{}
	for _, sp := range spaces {
		byHost[sp.Host] = sp
	}
	if got := byHost["a.11ch.at"]; got.StorageBytes != 12345 || got.MeteredAt == 0 || !got.SetupComplete || got.Members != 3 {
		t.Errorf("a = %+v, want 12345 bytes, claimed-ness recorded, timestamp", got)
	}
	if got := byHost["b.11ch.at"]; got.StorageBytes != 0 || got.MeteredAt != 0 || got.SetupComplete {
		t.Errorf("b = %+v, want untouched (stats error is a skip)", got)
	}
	_ = idA
	_ = idB

	// A later sweep refreshes the reading, never accumulates.
	bytesByHost["a.11ch.at"] = 200
	before := time.Now().Unix()
	s.meterOnce()
	spaces, _ = db.listSpaces("c", "r")
	for _, sp := range spaces {
		if sp.Host == "a.11ch.at" && (sp.StorageBytes != 200 || sp.MeteredAt < before) {
			t.Errorf("refresh: %+v, want 200 bytes, fresh timestamp", sp)
		}
	}
}

func TestParseByteSizeAndCaps(t *testing.T) {
	for in, want := range map[string]int64{
		"25GB": 25e9, "2TB": 2e12, "500MB": 5e8, "1024": 1024, "10 GB": 10e9, "1kb": 1000,
	} {
		if got, err := parseByteSize(in); err != nil || got != want {
			t.Errorf("parseByteSize(%q) = %d, %v; want %d", in, got, err, want)
		}
	}
	for _, bad := range []string{"", "-5GB", "0", "GB", "ten GB"} {
		if _, err := parseByteSize(bad); err == nil {
			t.Errorf("parseByteSize(%q) accepted, want error", bad)
		}
	}

	caps, _, err := parseStorageCaps("free=10GB, family=25GB,business=2TB")
	if err != nil || caps["free"] != 10e9 || caps["family"] != 25e9 || caps["business"] != 2e12 {
		t.Fatalf("caps = %v, %v", caps, err)
	}
	if c, sc, err := parseStorageCaps(""); c != nil || sc != nil || err != nil {
		t.Errorf("empty caps = %v, %v; want nil, nil (enforcement off)", c, err)
	}
	if _, _, err := parseStorageCaps("free=banana"); err == nil {
		t.Error("garbage cap accepted")
	}
}

func TestSetInstanceQuota(t *testing.T) {
	path := filepath.Join(t.TempDir(), "space.env")
	if err := os.WriteFile(path, []byte("LCHAT_ADMIN_TOKEN=secret\nLCHAT_DB=/data/x.db\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := setInstanceQuota(path, 25e9, "free"); err != nil {
		t.Fatal(err)
	}
	env, err := readEnvFile(path)
	if err != nil {
		t.Fatal(err)
	}
	// The cap and entitlement are set and every other var survived.
	if env["LCHAT_BLOB_QUOTA_BYTES"] != "25000000000" || env["LCHAT_ENTITLEMENT"] != "free" ||
		env["LCHAT_ADMIN_TOKEN"] != "secret" || env["LCHAT_DB"] != "/data/x.db" {
		t.Fatalf("env after set = %v", env)
	}
	// Idempotent update; an empty entitlement leaves the flag alone.
	if err := setInstanceQuota(path, 10e9, ""); err != nil {
		t.Fatal(err)
	}
	env, _ = readEnvFile(path)
	if env["LCHAT_BLOB_QUOTA_BYTES"] != "10000000000" || env["LCHAT_ENTITLEMENT"] != "free" {
		t.Fatalf("env after update = %v", env)
	}
	// A billing flip rides the same write; removal at 0 spares the flag.
	if err := setInstanceQuota(path, 0, "paid"); err != nil {
		t.Fatal(err)
	}
	env, _ = readEnvFile(path)
	if _, ok := env["LCHAT_BLOB_QUOTA_BYTES"]; ok {
		t.Fatalf("cap not removed: %v", env)
	}
	if env["LCHAT_ENTITLEMENT"] != "paid" {
		t.Fatalf("entitlement not flipped: %v", env)
	}
	// A missing instance surfaces not-exist (the handler's 404).
	if err := setInstanceQuota(filepath.Join(t.TempDir(), "absent.env"), 1, ""); !os.IsNotExist(err) {
		t.Errorf("absent env: err = %v, want not-exist", err)
	}
}

// quotaControl fakes controld's stats + quota endpoints, recording pushes.
func quotaControl(t *testing.T, bytesByHost map[string]int64, pushed map[string]int64) string {
	t.Helper()
	mux := http.NewServeMux()
	mux.HandleFunc("GET /instances/{host}/stats", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, map[string]any{"storage_bytes": bytesByHost[r.PathValue("host")]})
	})
	mux.HandleFunc("POST /instances/{host}/quota", func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			Bytes int64 `json:"bytes"`
		}
		json.NewDecoder(r.Body).Decode(&req)
		pushed[r.PathValue("host")] = req.Bytes
		writeJSON(w, http.StatusOK, map[string]any{"ok": true})
	})
	return fakeControl(t, mux.ServeHTTP)
}

// TestQuotaPushOnUpgradeAndSweep pins the enforcement loop: a webhook upgrade
// pushes the tier's cap immediately; the meter sweep reconciles free spaces
// and downgrades.
func TestQuotaPushOnUpgradeAndSweep(t *testing.T) {
	db := openTestCloudDB(t)
	paidID := insertSpaceRow(t, db, "paid.11ch.at", 1000)
	insertSpaceRow(t, db, "free.11ch.at", 1000)

	pushed := map[string]int64{}
	sock := quotaControl(t, map[string]int64{"paid.11ch.at": 1, "free.11ch.at": 1}, pushed)
	s := &cloudServer{
		db: db, control: controlClient(sock), controlToken: "t",
		storageCaps: map[string]int64{"free": 10e9, "family": 25e9},
	}

	// The upgrade path (what the checkout webhook calls).
	if err := db.markSpacePaid(paidID, "cus", "sub", "family", ""); err != nil {
		t.Fatal(err)
	}
	s.pushQuotaByID(paidID)
	if pushed["paid.11ch.at"] != 25e9 {
		t.Fatalf("upgrade push = %d, want 25GB", pushed["paid.11ch.at"])
	}

	// The sweep reconciles the free space (never pushed) — and doesn't
	// re-push the paid one (already at its cap).
	delete(pushed, "paid.11ch.at")
	s.meterOnce()
	if pushed["free.11ch.at"] != 10e9 {
		t.Fatalf("sweep push = %d, want 10GB", pushed["free.11ch.at"])
	}
	if _, again := pushed["paid.11ch.at"]; again {
		t.Error("sweep re-pushed an already-correct quota")
	}

	// A lapse (subscription deleted → free) re-applies the free cap.
	if err := db.setEntitlement(paidID, entFree, 0); err != nil {
		t.Fatal(err)
	}
	s.pushQuotaByID(paidID)
	if pushed["paid.11ch.at"] != 10e9 {
		t.Fatalf("lapse push = %d, want back to 10GB", pushed["paid.11ch.at"])
	}

	// No caps configured → nothing pushes anywhere.
	pushed2 := map[string]int64{}
	sock2 := quotaControl(t, map[string]int64{"paid.11ch.at": 1, "free.11ch.at": 1}, pushed2)
	s2 := &cloudServer{db: db, control: controlClient(sock2), controlToken: "t"}
	s2.meterOnce()
	s2.pushQuotaByID(paidID)
	if len(pushed2) != 0 {
		t.Errorf("unconfigured caps pushed: %v", pushed2)
	}
}

// The quota and meter pushes write DIFFERENT keys to the SAME env file from
// separate request goroutines; updateInstanceEnv serializes them so neither
// side's read-before-write snapshot can drop the other's keys. Hammered
// concurrently, every key family must survive.
func TestInstanceEnvWritesSerialized(t *testing.T) {
	envPath := filepath.Join(t.TempDir(), "space.env")
	if err := writeEnvFile(envPath, map[string]string{"LCHAT_ADMIN_TOKEN": "keepme"}); err != nil {
		t.Fatal(err)
	}
	var wg sync.WaitGroup
	for i := 0; i < 50; i++ {
		wg.Add(2)
		go func() {
			defer wg.Done()
			if err := setInstanceQuota(envPath, 12345, "free"); err != nil {
				t.Error(err)
			}
		}()
		go func() {
			defer wg.Done()
			if err := setInstanceMeter(envPath, "https://cloud.example/api/usage", "tok"); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()
	env, err := readEnvFile(envPath)
	if err != nil {
		t.Fatal(err)
	}
	for _, key := range []string{"LCHAT_ADMIN_TOKEN", "LCHAT_BLOB_QUOTA_BYTES", "LCHAT_METER_URL", "LCHAT_METER_TOKEN"} {
		if env[key] == "" {
			t.Errorf("%s lost in concurrent pushes: %v", key, env)
		}
	}
}