package oracle

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

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/swarm/verify"
)

func TestGuarded_RefusesOutOfScope(t *testing.T) {
	var reached bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reached = true
		w.WriteHeader(200)
	}))
	defer srv.Close()

	guard := func(host string) error { return errors.New("out of scope: " + host) }
	doer := Guarded(srv.Client(), guard)

	spec := &verify.ProbeSpec{
		Class: verify.ProbeDifferential,
		BaseRequest: "GET /?id=" + verify.DefaultMarker + " HTTP/1.1\r\nHost: x\r\n\r\n",
		AttackPayload: "2", ControlPayload: "9",
	}
	_, _, err := LiveDiffProbe(doer)(context.Background(), specCand(srv.URL, spec))
	if err == nil || !strings.Contains(err.Error(), "scope guard") {
		t.Fatalf("expected scope-guard refusal, got %v", err)
	}
	if reached {
		t.Error("a guarded out-of-scope request must never reach the server")
	}
}

func TestGuarded_AllowsInScope(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		_, _ = w.Write([]byte("ok"))
	}))
	defer srv.Close()

	guard := func(host string) error { return nil } // everything allowed
	doer := Guarded(srv.Client(), guard)

	spec := &verify.ProbeSpec{
		Class: verify.ProbeDifferential,
		BaseRequest: "GET /?id=" + verify.DefaultMarker + " HTTP/1.1\r\nHost: x\r\n\r\n",
		AttackPayload: "2", ControlPayload: "9",
	}
	attack, _, err := LiveDiffProbe(doer)(context.Background(), specCand(srv.URL, spec))
	if err != nil || attack.Status != 200 {
		t.Fatalf("in-scope probe should succeed: status=%d err=%v", attack.Status, err)
	}
}

func TestGuarded_NilPassThrough(t *testing.T) {
	if Guarded(nil, func(string) error { return nil }) != nil {
		t.Error("nil doer should pass through")
	}
	d := http.DefaultClient
	if Guarded(d, nil) == nil {
		t.Error("nil guard should return the doer unchanged (non-nil)")
	}
}

// BuildLive threads the guard into the oracle set so every oracle is guarded.
func TestBuildLive_GuardApplied(t *testing.T) {
	var reached bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reached = true
	}))
	defer srv.Close()

	oracles := BuildLive(nil, srv.Client(), Options{Guard: func(string) error { return errors.New("blocked") }})
	spec := &verify.ProbeSpec{
		Class: verify.ProbeDifferential,
		BaseRequest: "GET /?id=" + verify.DefaultMarker + " HTTP/1.1\r\nHost: x\r\n\r\n",
		AttackPayload: "2", ControlPayload: "9",
	}
	cand := specCand(srv.URL, spec)
	for _, o := range oracles {
		if o.Applies(cand) {
			v, _, _, err := o.Probe(context.Background(), cand)
			if err == nil {
				t.Errorf("%s probe should have been blocked by the guard", o.Name())
			}
			_ = v
		}
	}
	if reached {
		t.Error("guard must prevent any request from reaching the server")
	}
}

func TestRateLimited_TakesTokenBeforeEachRequest(t *testing.T) {
	var reqs, takes int
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reqs++
		w.WriteHeader(200)
	}))
	defer srv.Close()

	gate := func(ctx context.Context) error { takes++; return nil }
	doer := RateLimited(srv.Client(), gate)
	spec := &verify.ProbeSpec{
		Class: verify.ProbeDifferential,
		BaseRequest: "GET /?id=" + verify.DefaultMarker + " HTTP/1.1\r\nHost: x\r\n\r\n",
		AttackPayload: "2", ControlPayload: "9",
	}
	_, _, err := LiveDiffProbe(doer)(context.Background(), specCand(srv.URL, spec))
	if err != nil {
		t.Fatal(err)
	}
	if takes != 2 || reqs != 2 {
		t.Errorf("takes=%d reqs=%d, want 2 each (attack + control each gated)", takes, reqs)
	}
}

func TestRateLimited_GateErrorBlocksRequest(t *testing.T) {
	var reached bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { reached = true }))
	defer srv.Close()

	gate := func(ctx context.Context) error { return context.Canceled }
	doer := RateLimited(srv.Client(), gate)
	spec := &verify.ProbeSpec{
		Class: verify.ProbeDifferential,
		BaseRequest: "GET /?id=" + verify.DefaultMarker + " HTTP/1.1\r\nHost: x\r\n\r\n",
		AttackPayload: "2", ControlPayload: "9",
	}
	_, _, err := LiveDiffProbe(doer)(context.Background(), specCand(srv.URL, spec))
	if err == nil {
		t.Fatal("gate error should abort the probe")
	}
	if reached {
		t.Error("a gated-out request must not reach the server")
	}
}

func TestRateLimited_NilPassThrough(t *testing.T) {
	if RateLimited(nil, func(context.Context) error { return nil }) != nil {
		t.Error("nil doer passes through")
	}
	if RateLimited(http.DefaultClient, nil) == nil {
		t.Error("nil gate returns doer unchanged (non-nil)")
	}
}