package exploit

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

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

func TestSSTISpecs_Shapes(t *testing.T) {
	specs, err := SSTISpecs("https://t/greet?name=bob", "name", "Bearer tok")
	if err != nil {
		t.Fatal(err)
	}
	if len(specs) != len(sstiEngines) {
		t.Fatalf("got %d specs, want one per engine (%d)", len(specs), len(sstiEngines))
	}
	for _, ts := range specs {
		if ts.Spec.Class != verify.ProbeDifferential || !ts.Spec.Valid() {
			t.Fatalf("%s: invalid differential spec", ts.Technique)
		}
		if ts.Spec.Indicator != "36863" {
			t.Errorf("%s: indicator = %q, want the product 36863", ts.Technique, ts.Spec.Indicator)
		}
		if !strings.Contains(ts.Spec.BaseRequest, "name="+verify.DefaultMarker) {
			t.Errorf("%s: param not marked: %q", ts.Technique, ts.Spec.BaseRequest)
		}
		if !strings.Contains(ts.Spec.BaseRequest, "Authorization: Bearer tok") {
			t.Errorf("%s: credential not carried", ts.Technique)
		}
		// Payloads are query-encoded (no raw spaces/braces that would corrupt the
		// request line).
		if strings.ContainsAny(ts.Spec.AttackPayload, " {}<>") {
			t.Errorf("%s: attack payload not encoded: %q", ts.Technique, ts.Spec.AttackPayload)
		}
	}
}

func TestSSTISpecs_RejectsEmptyParam(t *testing.T) {
	if _, err := SSTISpecs("https://t/x?a=1", "", ""); err == nil {
		t.Error("empty parameter should be rejected")
	}
}

// End to end: a fake Jinja-style engine that evaluates {{a*b}} is proven, while
// a plain reflection (the control) echoes the literal expression — no product.
func TestSSTISpecs_ProvesEvaluatingEngine(t *testing.T) {
	jinja := regexp.MustCompile(`\{\{\s*(\d+)\s*\*\s*(\d+)\s*\}\}`)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		name := r.URL.Query().Get("name")
		// Simulate template rendering: evaluate {{a*b}}; otherwise reflect as-is.
		rendered := jinja.ReplaceAllStringFunc(name, func(m string) string {
			g := jinja.FindStringSubmatch(m)
			a, _ := atoiSafe(g[1])
			b, _ := atoiSafe(g[2])
			return itoa(a * b)
		})
		w.WriteHeader(200)
		_, _ = w.Write([]byte("Hello " + rendered))
	}))
	defer srv.Close()

	specs, err := SSTISpecs(srv.URL+"/greet?name=x", "name", "")
	if err != nil {
		t.Fatal(err)
	}
	d := oracle.NewDifferential(oracle.LiveDiffProbe(srv.Client()))
	proven := ""
	for _, ts := range specs {
		v, _, _, err := d.Probe(context.Background(), verify.Candidate{Target: srv.URL, Spec: ts.Spec})
		if err == nil && v == verify.Proven {
			proven = ts.Technique
		}
	}
	if proven != "ssti:Jinja2/Twig/Nunjucks" {
		t.Errorf("expected the Jinja engine proven, got %q", proven)
	}
}

// A plain reflection that never evaluates must not be flagged as SSTI.
func TestSSTISpecs_DisprovesPlainReflection(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		_, _ = w.Write([]byte("Hello " + r.URL.Query().Get("name"))) // echo, no templating
	}))
	defer srv.Close()

	specs, err := SSTISpecs(srv.URL+"/greet?name=x", "name", "")
	if err != nil {
		t.Fatal(err)
	}
	d := oracle.NewDifferential(oracle.LiveDiffProbe(srv.Client()))
	for _, ts := range specs {
		v, _, _, err := d.Probe(context.Background(), verify.Candidate{Target: srv.URL, Spec: ts.Spec})
		if err == nil && v == verify.Proven {
			t.Fatalf("plain reflection must not prove SSTI (%s)", ts.Technique)
		}
	}
}

func atoiSafe(s string) (int, bool) {
	n := 0
	for _, r := range s {
		if r < '0' || r > '9' {
			return 0, false
		}
		n = n*10 + int(r-'0')
	}
	return n, true
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b []byte
	for n > 0 {
		b = append([]byte{byte('0' + n%10)}, b...)
		n /= 10
	}
	return string(b)
}

func TestSSTIPolyglotSpec_ProvesAndGates(t *testing.T) {
	spec, err := SSTIPolyglotSpec("https://t/g?name=x", "name", "")
	if err != nil {
		t.Fatal(err)
	}
	if spec.Indicator != "36863" || spec.Class != verify.ProbeDifferential {
		t.Fatalf("bad polyglot spec: %+v", spec)
	}
	if strings.ContainsAny(spec.AttackPayload, " {}<>") {
		t.Errorf("polyglot attack payload not encoded: %q", spec.AttackPayload)
	}
}