package exploit

import (
	"context"
	"fmt"
	"net/url"
	"strings"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/pipeline"
	"github.com/google/uuid"
)

// BOLAHit is one confirmed cross-user access: an authenticated request for an
// object id the session's user does not own that still returned data.
type BOLAHit struct {
	URL      string
	Evidence string
}

// DeriveBOLATargets turns concrete endpoint URLs into id-templated probe
// targets. Any path segment that is a UUID or a 2+ digit number is treated as
// the object id and replaced with the "{id}" placeholder; URLs with no such
// segment (no object to pivot on) are skipped, and duplicates collapse. This is
// how the adaptive loop decides where a harvested id might be replayable.
func DeriveBOLATargets(urls []string) []string {
	seen := make(map[string]struct{})
	var out []string
	for _, raw := range urls {
		u, err := url.Parse(raw)
		if err != nil {
			continue
		}
		segs := strings.Split(u.Path, "/")
		replaced := false
		for i, s := range segs {
			// A concrete object id (uuid / number) or an existing template
			// placeholder ({id}, {{victim_vehicle}}, …) marks the object slot.
			if uuidRe.MatchString(s) || (len(s) >= 2 && isAllDigits(s)) || strings.HasPrefix(s, "{") {
				segs[i] = "{id}"
				replaced = true
				break // template one id segment — the object being addressed
			}
		}
		if !replaced {
			continue
		}
		tmpl := raw
		if u.Host != "" {
			tmpl = u.Scheme + "://" + u.Host + strings.Join(segs, "/")
		}
		if _, dup := seen[tmpl]; dup {
			continue
		}
		seen[tmpl] = struct{}{}
		out = append(out, tmpl)
	}
	return out
}

func isAllDigits(s string) bool {
	for _, r := range s {
		if r < '0' || r > '9' {
			return false
		}
	}
	return s != ""
}

// SweepBOLA is the generalized BOLA/IDOR engine: it replays harvested object
// ids across id-templated endpoints using the caller's own session
// (authHeader, e.g. "Bearer <jwt>"). A 200 for an id that the session's user
// does not own is a broken-object-level-authorization hit. It is bounded by
// maxProbes so a large surface can't explode into thousands of requests, and it
// only replays uuid/numeric refs (an email is not a path-addressable object
// id). Every request goes through the executor's scope-validated httpreq path.
func (e *Executor) SweepBOLA(ctx context.Context, authHeader string, templates []string, refs []ObjectRef, campaignID uuid.UUID) []BOLAHit {
	return e.SweepBOLAWithProbe(ctx, authHeader, templates, refs, campaignID, nil)
}

// SweepBOLAWithProbe is SweepBOLA with a per-probe callback: onProbe is invoked
// once for every replay work-unit right after it runs, with the probed target
// and whether it returned a cross-user 200. It exists so the live dashboard and
// TUI can visualize the sweep fanning out into many concurrent probe workers —
// each callback is a cheap telemetry event, not extra network or LLM work. A
// nil onProbe is the plain SweepBOLA behavior.
func (e *Executor) SweepBOLAWithProbe(ctx context.Context, authHeader string, templates []string, refs []ObjectRef, campaignID uuid.UUID, onProbe func(target string, ok bool)) []BOLAHit {
	const maxProbes = 40
	var hits []BOLAHit
	probes := 0
	for _, tmpl := range templates {
		for _, ref := range refs {
			if ref.Kind == "email" {
				continue
			}
			if probes >= maxProbes {
				return hits
			}
			probes++
			target := strings.Replace(tmpl, "{id}", ref.Value, 1)
			step := pipeline.AttackStep{
				ID:                    uuid.New(),
				Name:                  "adaptive BOLA probe",
				Command:               fmt.Sprintf("httpreq --url %s --header 'Authorization: %s'", target, authHeader),
				ExpectedOutputPattern: "HTTP 200",
			}
			res, _ := e.Execute(ctx, step, campaignID)
			ok := res != nil && res.Success
			if onProbe != nil {
				onProbe(target, ok)
			}
			if ok {
				hits = append(hits, BOLAHit{URL: target, Evidence: res.Output})
			}
		}
	}
	return hits
}