package tools

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

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

// WhatWebTool wraps urbanadventurer/WhatWeb (https://github.com/urbanadventurer/WhatWeb)
// for web technology fingerprinting: CMS, frameworks, servers, JS
// libraries, analytics and hundreds of other signatures.
//
// WhatWeb identifies what a target is built from, not what is wrong with
// it, so every detection is surfaced at "info" severity. The classifier
// agent can correlate the reported product versions against known CVEs.
type WhatWebTool struct{}

// NewWhatWebTool constructs the adapter.
func NewWhatWebTool() *WhatWebTool { return &WhatWebTool{} }

// Name implements Tool.
func (w *WhatWebTool) Name() string { return "whatweb" }

// IsAvailable checks for `whatweb` on PATH.
func (w *WhatWebTool) IsAvailable() bool { return IsCommandAvailable("whatweb") }

// Run executes whatweb against a target URL/host.
//
// Supported options:
//
//	timeout    int    — per-invocation timeout in seconds (default 120).
//	aggression int    — aggression level (-a): 1 stealthy, 3 aggressive,
//	                     4 heavy. Empty means whatweb's default (1).
//	user_agent string — custom UA (--user-agent).
//	proxy      string — route requests through a proxy (--proxy).
func (w *WhatWebTool) Run(ctx context.Context, target string, opts Options) (*ToolResult, error) {
	if scopeDef := getScopeFromContext(ctx); scopeDef != nil {
		if err := scope.ValidateAndLog("whatweb", target, *scopeDef); err != nil {
			return nil, fmt.Errorf("scope violation in whatweb: %w", err)
		}
	}

	timeout := time.Duration(opts.GetInt("timeout", 120)) * time.Second

	tmp, err := os.CreateTemp("", "whatweb-*.json")
	if err != nil {
		return &ToolResult{ToolName: "whatweb", Target: target, Error: err}, err
	}
	tmpPath := tmp.Name()
	_ = tmp.Close()
	defer os.Remove(tmpPath)

	args := []string{
		target,
		"--log-json", tmpPath,
		"--color", "never",
		"--no-errors",
	}
	if aggression := opts.GetInt("aggression", 0); aggression > 0 {
		args = append(args, "-a", fmt.Sprintf("%d", aggression))
	}
	if ua := opts.GetString("user_agent", ""); ua != "" {
		args = append(args, "--user-agent", ua)
	}
	if proxy := opts.GetString("proxy", ""); proxy != "" {
		args = append(args, "--proxy", proxy)
	}

	result := RunToolCommand(ctx, "whatweb", target, timeout, "whatweb", args...)
	if result.Error != nil && !strings.Contains(result.Error.Error(), "exit status") {
		return result, result.Error
	}
	result.Error = nil

	body, readErr := os.ReadFile(tmpPath)
	if readErr != nil || len(body) == 0 {
		body = []byte(result.RawOutput)
	}
	result.ParsedFindings = parseWhatWebJSON(body)
	return result, nil
}

// whatwebResult mirrors one record of whatweb's `--log-json` output: a
// flat array of per-target objects. `plugins` maps a detected
// technology name to the evidence whatweb matched for it. whatweb emits
// a trailing empty object in the array, which carries no target.
type whatwebResult struct {
	Target     string                   `json:"target"`
	HTTPStatus int                      `json:"http_status"`
	Plugins    map[string]whatwebPlugin `json:"plugins"`
}

// whatwebPlugin holds the evidence fields a plugin may populate. A bare
// detection (e.g. HTML5) leaves them all empty and is still reported.
type whatwebPlugin struct {
	Version []string `json:"version,omitempty"`
	String  []string `json:"string,omitempty"`
	OS      []string `json:"os,omitempty"`
	Module  []string `json:"module,omitempty"`
	Account []string `json:"account,omitempty"`
}

// parseWhatWebJSON turns the result array into one finding per detected
// technology. The plugin name is the technology; version and string
// evidence are carried through when present. All findings are "info"
// severity since whatweb fingerprints rather than rating risk.
func parseWhatWebJSON(body []byte) []map[string]any {
	if len(body) == 0 {
		return nil
	}
	var results []whatwebResult
	if err := json.Unmarshal(body, &results); err != nil {
		return nil
	}
	var findings []map[string]any
	for _, r := range results {
		if r.Target == "" || len(r.Plugins) == 0 {
			continue
		}
		names := make([]string, 0, len(r.Plugins))
		for name := range r.Plugins {
			names = append(names, name)
		}
		sort.Strings(names)
		for _, name := range names {
			p := r.Plugins[name]
			finding := map[string]any{
				"tool":       "whatweb",
				"url":        r.Target,
				"technology": name,
				"title":      whatwebTitle(name, p.Version),
				"severity":   "info",
			}
			if r.HTTPStatus != 0 {
				finding["http_status"] = r.HTTPStatus
			}
			if len(p.Version) > 0 {
				finding["version"] = strings.Join(p.Version, ", ")
			}
			if len(p.String) > 0 {
				finding["string"] = strings.Join(p.String, ", ")
			}
			if len(p.OS) > 0 {
				finding["os"] = strings.Join(p.OS, ", ")
			}
			findings = append(findings, finding)
		}
	}
	return findings
}

func whatwebTitle(name string, version []string) string {
	if len(version) > 0 {
		return fmt.Sprintf("%s (%s)", name, strings.Join(version, ", "))
	}
	return name
}