package recon

import (
	"context"
	"fmt"
	"strings"
	"time"

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

// ReconAgent orchestrates security tools and analyzes output to build an AttackSurface.
type ReconAgent struct {
	provider    llm.Provider
	coordinator *tools.Coordinator
	strict      bool
	onErr       func(error)
}

// Option customises ReconAgent construction.
type Option func(*ReconAgent)

// WithStrict makes LLM failures fatal instead of returning a partial surface.
func WithStrict() Option {
	return func(r *ReconAgent) { r.strict = true }
}

// WithErrorSink installs a callback for LLM / parse errors. Useful for
// emitting degraded-mode warnings to the event stream.
func WithErrorSink(fn func(error)) Option {
	return func(r *ReconAgent) { r.onErr = fn }
}

// NewReconAgent creates a new recon agent.
func NewReconAgent(provider llm.Provider, coordinator *tools.Coordinator, opts ...Option) *ReconAgent {
	r := &ReconAgent{
		provider:    provider,
		coordinator: coordinator,
	}
	for _, opt := range opts {
		opt(r)
	}
	return r
}

// ReconPlan defines which tools to run based on target type.
type ReconPlan struct {
	Target    string   `json:"target"`
	ToolOrder []string `json:"tool_order"`
}

// PlanRecon determines which tools to run based on the target type.
func (r *ReconAgent) PlanRecon(target string) ReconPlan {
	plan := ReconPlan{Target: target}

	if isIPTarget(target) {
		// IP target: port scan → probe → scan
		plan.ToolOrder = []string{"naabu", "httpx", "nuclei"}
	} else if isURLTarget(target) {
		// URL target: probe → crawl → history → scan
		plan.ToolOrder = []string{"httpx", "katana", "gau", "nuclei"}
	} else {
		// Domain target: full recon pipeline
		plan.ToolOrder = []string{"subfinder", "dnsx", "naabu", "httpx", "katana", "gau", "nuclei"}
	}

	return plan
}

// Execute runs the recon plan and produces an AttackSurface.
func (r *ReconAgent) Execute(ctx context.Context, plan ReconPlan, scopeDef *scope.ScopeDefinition, campaignID uuid.UUID) (*pipeline.AttackSurface, error) {
	// Run tools
	_, resultCh := r.coordinator.RunSelected(ctx, plan.ToolOrder, plan.Target, scopeDef, tools.Options{})

	// Collect results as they stream in
	var results []*tools.ToolResult
	for result := range resultCh {
		results = append(results, result)
	}

	// Analyze results with LLM
	surface, err := r.Analyze(ctx, results, campaignID)
	if err != nil {
		return nil, fmt.Errorf("analyzing recon results: %w", err)
	}

	return surface, nil
}

// Analyze sends tool results to the LLM for structured analysis.
func (r *ReconAgent) Analyze(ctx context.Context, results []*tools.ToolResult, campaignID uuid.UUID) (*pipeline.AttackSurface, error) {
	// Build context from tool results
	var contextBuilder strings.Builder
	for _, result := range results {
		if result.Error != nil {
			contextBuilder.WriteString(fmt.Sprintf("Tool: %s (FAILED: %s)\n\n", result.ToolName, result.Error))
			continue
		}
		contextBuilder.WriteString(fmt.Sprintf("Tool: %s\nOutput:\n%s\n\n", result.ToolName, result.RawOutput))
	}

	req := llm.CompletionRequest{
		SystemPrompt: reconSystemPrompt,
		Messages: []llm.Message{
			{
				Role: "user",
				Content: fmt.Sprintf(
					"Analyze the following security tool outputs and produce a structured AttackSurface JSON object.\n\n%s\n\nRespond ONLY with valid JSON matching the AttackSurface schema.",
					contextBuilder.String(),
				),
			},
		},
		MaxTokens:         8192,
		Temperature:       0.1,
		CacheSystemPrompt: true,
	}

	resp, err := r.provider.Complete(ctx, req)
	if err != nil {
		return nil, fmt.Errorf("LLM analysis failed: %w", err)
	}

	// Parse the response
	surface, err := ParseAttackSurface(resp.Content)
	if err != nil {
		// Retry with simplified prompt
		retryReq := llm.CompletionRequest{
			SystemPrompt: "You are a JSON parser. Extract structured data from security tool output. Respond ONLY with valid JSON.",
			Messages: []llm.Message{
				{
					Role: "user",
					Content: fmt.Sprintf(
						"Parse this into an AttackSurface JSON with fields: target, subdomains, hosts, endpoints, technologies.\n\n%s",
						contextBuilder.String(),
					),
				},
			},
			MaxTokens:   8192,
			Temperature: 0,
		}

		retryResp, retryErr := r.provider.Complete(ctx, retryReq)
		if retryErr != nil {
			return nil, fmt.Errorf("retry analysis also failed: %w", retryErr)
		}

		surface, err = ParseAttackSurface(retryResp.Content)
		if err != nil {
			if r.strict {
				return nil, fmt.Errorf("recon parse failed after retry: %w", err)
			}
			if r.onErr != nil {
				r.onErr(fmt.Errorf("recon analysis returned empty surface: %w", err))
			}
			// Return partial results rather than error (degraded mode).
			return &pipeline.AttackSurface{
				CampaignID: campaignID,
				CreatedAt:  time.Now(),
			}, nil
		}
	}

	surface.CampaignID = campaignID
	surface.CreatedAt = time.Now()

	return surface, nil
}

const reconSystemPrompt = `You are a specialized security reconnaissance analyst. Your job is to analyze output from security scanning tools and produce a structured attack surface model.

Given the raw output from tools like subfinder, httpx, nuclei, naabu, katana, dnsx, and gau, you must:

1. Identify all discovered subdomains with their IP addresses and sources
2. Map all hosts with their open ports and running services
3. Catalog all discovered web endpoints with parameters
4. Detect technologies and their versions
5. Note any interesting findings or anomalies

Output your analysis as a valid JSON object matching this schema:
{
  "target": "string",
  "subdomains": [{"domain": "string", "ip": "string", "source": "string"}],
  "hosts": [{"ip": "string", "hostnames": ["string"], "open_ports": [int], "services": {}, "os": "string"}],
  "endpoints": [{"url": "string", "method": "string", "status_code": int, "interesting": bool}],
  "technologies": {"key": "version"}
}

Respond ONLY with the JSON object. No markdown, no explanation.`

func isIPTarget(target string) bool {
	parts := strings.Split(target, ".")
	if len(parts) != 4 {
		return false
	}
	for _, p := range parts {
		for _, c := range p {
			if c < '0' || c > '9' {
				return false
			}
		}
	}
	return true
}

func isURLTarget(target string) bool {
	return strings.HasPrefix(target, "http://") || strings.HasPrefix(target, "https://")
}
