package tools

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

// fakeTool is a lightweight Tool implementation that records whether
// Run was invoked and returns a caller-controlled result/error pair.
// We deliberately bypass the real coordinator's tool registry by
// constructing a Coordinator with a hand-rolled tools map so the test
// does not depend on subfinder / nuclei / etc. being installed.
type fakeTool struct {
	name      string
	available bool
	runFunc   func(ctx context.Context, target string, opts Options) (*ToolResult, error)
	runs      int32
}

func (f *fakeTool) Name() string      { return f.name }
func (f *fakeTool) IsAvailable() bool { return f.available }
func (f *fakeTool) Run(ctx context.Context, target string, opts Options) (*ToolResult, error) {
	atomic.AddInt32(&f.runs, 1)
	if f.runFunc != nil {
		return f.runFunc(ctx, target, opts)
	}
	return &ToolResult{ToolName: f.name, Target: target, RawOutput: "ok"}, nil
}

// newFakeCoordinator builds a Coordinator whose registry contains only
// the supplied tools, sidestepping NewCoordinator's hard-coded list.
// This keeps the unit test hermetic.
func newFakeCoordinator(toolsIn ...Tool) *Coordinator {
	c := &Coordinator{tools: map[string]Tool{}}
	for _, t := range toolsIn {
		c.tools[t.Name()] = t
	}
	return c
}

func drain(ch <-chan *ToolResult) []*ToolResult {
	var got []*ToolResult
	for r := range ch {
		got = append(got, r)
	}
	return got
}

func TestSetHooks_NilSafe(t *testing.T) {
	c := newFakeCoordinator(&fakeTool{name: "fake", available: true})
	c.SetHooks(nil) // must not panic
	_, ch := c.RunSelected(context.Background(), []string{"fake"}, "tgt", &scope.ScopeDefinition{}, Options{})
	if got := drain(ch); len(got) != 1 {
		t.Errorf("expected 1 result, got %d", len(got))
	}
}

func TestRunSelected_OnStartBeforeOnDone(t *testing.T) {
	var (
		mu         sync.Mutex
		order      []string
		gotResults []*ToolResult
	)
	tool := &fakeTool{
		name:      "fake",
		available: true,
		runFunc: func(ctx context.Context, target string, opts Options) (*ToolResult, error) {
			// Sleep so we can verify OnStart fired before OnDone in
			// wall-clock order, not just by lucky scheduling.
			time.Sleep(10 * time.Millisecond)
			return &ToolResult{ToolName: "fake", Target: target}, nil
		},
	}
	c := newFakeCoordinator(tool)
	c.SetHooks(&ToolHooks{
		OnStart: func(name, target string) {
			mu.Lock()
			defer mu.Unlock()
			order = append(order, "start:"+name)
		},
		OnDone: func(name, target string, result *ToolResult, err error) {
			mu.Lock()
			defer mu.Unlock()
			order = append(order, "done:"+name)
			gotResults = append(gotResults, result)
		},
	})

	_, ch := c.RunSelected(context.Background(), []string{"fake"}, "tgt", &scope.ScopeDefinition{}, Options{})
	_ = drain(ch)

	mu.Lock()
	defer mu.Unlock()
	if len(order) != 2 || order[0] != "start:fake" || order[1] != "done:fake" {
		t.Errorf("event order wrong: %v", order)
	}
	if len(gotResults) != 1 || gotResults[0].ToolName != "fake" {
		t.Errorf("OnDone did not receive the tool result: %+v", gotResults)
	}
}

func TestRunSelected_OnSkipForUnregisteredTool(t *testing.T) {
	var skipped []string
	c := newFakeCoordinator(&fakeTool{name: "present", available: true})
	c.SetHooks(&ToolHooks{
		OnSkip: func(name, target, reason string) {
			skipped = append(skipped, name+":"+reason)
		},
	})

	_, ch := c.RunSelected(context.Background(), []string{"ghost"}, "tgt", &scope.ScopeDefinition{}, Options{})
	_ = drain(ch)

	if len(skipped) != 1 || skipped[0] != "ghost:tool not registered" {
		t.Errorf("expected ghost skip reason, got %v", skipped)
	}
}

func TestRunSelected_OnSkipForUnavailableBinary(t *testing.T) {
	var skipped []string
	tool := &fakeTool{name: "fake", available: false}
	c := newFakeCoordinator(tool)
	c.SetHooks(&ToolHooks{
		OnSkip: func(name, target, reason string) {
			skipped = append(skipped, name+":"+reason)
		},
	})

	_, ch := c.RunSelected(context.Background(), []string{"fake"}, "tgt", &scope.ScopeDefinition{}, Options{})
	_ = drain(ch)

	if len(skipped) != 1 || skipped[0] != "fake:binary not found in PATH" {
		t.Errorf("expected fake skip with PATH reason, got %v", skipped)
	}
	if atomic.LoadInt32(&tool.runs) != 0 {
		t.Error("unavailable tool must not be invoked")
	}
}

func TestRunSelected_OnDoneReceivesError(t *testing.T) {
	wantErr := errors.New("boom")
	tool := &fakeTool{
		name:      "fake",
		available: true,
		runFunc: func(ctx context.Context, target string, opts Options) (*ToolResult, error) {
			return nil, wantErr
		},
	}
	c := newFakeCoordinator(tool)

	var (
		mu     sync.Mutex
		gotErr error
		gotRes *ToolResult
	)
	c.SetHooks(&ToolHooks{
		OnDone: func(name, target string, result *ToolResult, err error) {
			mu.Lock()
			defer mu.Unlock()
			gotErr = err
			gotRes = result
		},
	})

	_, ch := c.RunSelected(context.Background(), []string{"fake"}, "tgt", &scope.ScopeDefinition{}, Options{})
	_ = drain(ch)

	mu.Lock()
	defer mu.Unlock()
	if !errors.Is(gotErr, wantErr) {
		t.Errorf("OnDone err = %v, want %v", gotErr, wantErr)
	}
	if gotRes == nil || gotRes.Error == nil {
		t.Errorf("OnDone should receive a result with Error populated, got %+v", gotRes)
	}
}

func TestRunAll_HooksFireForEveryAvailableTool(t *testing.T) {
	a := &fakeTool{name: "a", available: true}
	b := &fakeTool{name: "b", available: true}
	skipped := &fakeTool{name: "skip-me", available: false}
	c := newFakeCoordinator(a, b, skipped)

	var (
		mu     sync.Mutex
		starts = map[string]int{}
		dones  = map[string]int{}
		skips  = map[string]int{}
	)
	c.SetHooks(&ToolHooks{
		OnStart: func(name, _ string) { mu.Lock(); starts[name]++; mu.Unlock() },
		OnDone:  func(name, _ string, _ *ToolResult, _ error) { mu.Lock(); dones[name]++; mu.Unlock() },
		OnSkip:  func(name, _, _ string) { mu.Lock(); skips[name]++; mu.Unlock() },
	})

	_, ch := c.RunAll(context.Background(), "tgt", &scope.ScopeDefinition{}, Options{})
	_ = drain(ch)

	mu.Lock()
	defer mu.Unlock()
	if starts["a"] != 1 || starts["b"] != 1 {
		t.Errorf("OnStart counts wrong: %v", starts)
	}
	if dones["a"] != 1 || dones["b"] != 1 {
		t.Errorf("OnDone counts wrong: %v", dones)
	}
	if skips["skip-me"] != 1 {
		t.Errorf("OnSkip should fire once for skip-me, got %v", skips)
	}
}