package oracle

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

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

func specCand(target string, s *verify.ProbeSpec) verify.Candidate {
	return verify.Candidate{Target: target, Spec: s}
}

// LiveDiffProbe should send attack + control and return their responses; the
// server returns victim data for the attack id and a 403 for the control.
func TestLiveDiffProbe_SendsAttackAndControl(t *testing.T) {
	var seen []string
	var mu sync.Mutex
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		id := r.URL.Query().Get("id")
		mu.Lock()
		seen = append(seen, id)
		mu.Unlock()
		if id == "2" {
			w.WriteHeader(200)
			_, _ = w.Write([]byte(`{"user":"victim"}`))
			return
		}
		w.WriteHeader(403)
		_, _ = w.Write([]byte(`forbidden`))
	}))
	defer srv.Close()

	spec := &verify.ProbeSpec{
		Class:          verify.ProbeDifferential,
		BaseRequest:    "GET /api/user?id=" + verify.DefaultMarker + " HTTP/1.1\nHost: x\n\n",
		AttackPayload:  "2",
		ControlPayload: "999999",
	}
	attack, control, err := LiveDiffProbe(srv.Client())(context.Background(), specCand(srv.URL, spec))
	if err != nil {
		t.Fatal(err)
	}
	if attack.Status != 200 || !strings.Contains(attack.Body, "victim") {
		t.Errorf("attack = %+v", attack)
	}
	if control.Status != 403 {
		t.Errorf("control = %+v", control)
	}
	if len(seen) != 2 || seen[0] != "2" || seen[1] != "999999" {
		t.Errorf("server saw ids %v, want [2 999999]", seen)
	}
}

func TestLiveDiffProbe_RejectsWrongSpec(t *testing.T) {
	_, _, err := LiveDiffProbe(http.DefaultClient)(context.Background(), specCand("http://t", nil))
	if err == nil {
		t.Error("nil spec should error")
	}
}

// LiveTimeProbe should inject the delay seconds and the server sleeps that long.
func TestLiveTimeProbe_InjectsDelay(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// payload like SLEEP(0.2) — extract the number and sleep that many secs.
		q := r.URL.Query().Get("q")
		if i := strings.Index(q, "SLEEP("); i >= 0 {
			rest := q[i+len("SLEEP("):]
			if j := strings.Index(rest, ")"); j >= 0 {
				if f, err := strconv.ParseFloat(rest[:j], 64); err == nil {
					time.Sleep(time.Duration(f * float64(time.Second)))
				}
			}
		}
		w.WriteHeader(200)
	}))
	defer srv.Close()

	spec := &verify.ProbeSpec{
		Class:         verify.ProbeTiming,
		BaseRequest:   "GET /?q=" + verify.DefaultMarker + " HTTP/1.1\nHost: x\n\n",
		DelayTemplate: "SLEEP(" + verify.DelayMarker + ")",
	}
	probe := LiveTimeProbe(srv.Client())

	base, err := probe(context.Background(), specCand(srv.URL, spec), 0)
	if err != nil {
		t.Fatal(err)
	}
	delayed, err := probe(context.Background(), specCand(srv.URL, spec), 250*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	if delayed-base < 180*time.Millisecond {
		t.Errorf("delayed (%v) should exceed baseline (%v) by ~250ms", delayed, base)
	}
}

// LiveDelivery should inject the OOB URL into the payload and send it.
func TestLiveDelivery_InjectsOOBURL(t *testing.T) {
	var gotQuery string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotQuery = r.URL.RawQuery
		w.WriteHeader(200)
	}))
	defer srv.Close()

	spec := &verify.ProbeSpec{
		Class:       verify.ProbeOOB,
		BaseRequest: "GET /fetch?url=" + verify.DefaultMarker + " HTTP/1.1\nHost: x\n\n",
		OOBTemplate: "http://" + verify.OOBMarker + "/x",
	}
	err := LiveDelivery(srv.Client())(context.Background(), specCand(srv.URL, spec), "abc.oast.fun")
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(gotQuery, "abc.oast.fun") {
		t.Errorf("server query %q should contain the injected OOB host", gotQuery)
	}
}

func TestRequireSpec(t *testing.T) {
	ok := &verify.ProbeSpec{Class: verify.ProbeTiming, BaseRequest: "GET /?q=" + verify.DefaultMarker + " HTTP/1.1\n\n", DelayTemplate: "SLEEP(" + verify.DelayMarker + ")"}
	applies := RequireSpec(verify.ProbeTiming)
	if !applies(verify.Candidate{Spec: ok}) {
		t.Error("valid timing spec should apply")
	}
	if applies(verify.Candidate{Spec: &verify.ProbeSpec{Class: verify.ProbeOOB}}) {
		t.Error("wrong class should not apply")
	}
	if applies(verify.Candidate{}) {
		t.Error("no spec should not apply")
	}
}

// A marker in the request BODY (JSON API param) is injected and sent — not just
// markers in the URL/query.
func TestLiveDiffProbe_InjectsBodyMarker(t *testing.T) {
	var bodies []string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		bodies = append(bodies, string(b))
		w.WriteHeader(200)
	}))
	defer srv.Close()

	spec := &verify.ProbeSpec{
		Class: verify.ProbeDifferential,
		BaseRequest: "POST /api/order HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\n\r\n" +
			`{"order_id":"` + verify.DefaultMarker + `"}`,
		AttackPayload:  "1337",
		ControlPayload: "000",
	}
	_, _, err := LiveDiffProbe(srv.Client())(context.Background(), specCand(srv.URL, spec))
	if err != nil {
		t.Fatal(err)
	}
	if len(bodies) != 2 || !strings.Contains(bodies[0], `"order_id":"1337"`) || !strings.Contains(bodies[1], `"order_id":"000"`) {
		t.Errorf("body markers not injected: %v", bodies)
	}
}