package supervise

import (
	"context"
	"errors"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestSplitCommandLine(t *testing.T) {
	cases := []struct {
		in   string
		want []string
	}{
		{`python main.py --port 8188`, []string{"python", "main.py", "--port", "8188"}},
		{`"C:\Program Files\py\python.exe" main.py`, []string{`C:\Program Files\py\python.exe`, "main.py"}},
		{`  spaced   out  `, []string{"spaced", "out"}},
		{`a "b c" d`, []string{"a", "b c", "d"}},
		{"", nil},
	}
	for _, c := range cases {
		got, err := splitCommandLine(c.in)
		if err != nil {
			t.Errorf("splitCommandLine(%q): %v", c.in, err)
			continue
		}
		if len(got) != len(c.want) {
			t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want)
			continue
		}
		for i := range got {
			if got[i] != c.want[i] {
				t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want)
				break
			}
		}
	}
	if _, err := splitCommandLine(`"unterminated`); err == nil {
		t.Error("expected error for unterminated quote")
	}
}

// TestHelperProcess is the child executed by the lifecycle tests: it just
// sleeps. The env marker is set only around the spawn, so in the normal
// test run this returns immediately.
func TestHelperProcess(t *testing.T) {
	if os.Getenv("GO_HELPER_PROCESS") != "1" {
		return
	}
	time.Sleep(30 * time.Second)
	os.Exit(0)
}

func newHelper(t *testing.T, name string) *Process {
	t.Helper()
	// The probe always fails: nothing external serves the port, so
	// EnsureRunning spawns the helper child.
	p, err := New(name, `"`+os.Args[0]+`" -test.run=TestHelperProcess`, "",
		func(context.Context) error { return errors.New("nothing there") }, 5*time.Second, slog.Default())
	if err != nil {
		t.Fatal(err)
	}
	return p
}

// startHelper spawns the child with the marker set; exec.Command inherits
// the environment at spawn time, so it can be unset right after.
func startHelper(t *testing.T, p *Process) {
	t.Helper()
	os.Setenv("GO_HELPER_PROCESS", "1")
	defer os.Unsetenv("GO_HELPER_PROCESS")
	if err := p.EnsureRunning(); err != nil {
		t.Fatal(err)
	}
}

func waitStopped(t *testing.T, p *Process, timeout time.Duration) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for p.Running() && time.Now().Before(deadline) {
		time.Sleep(50 * time.Millisecond)
	}
	if p.Running() {
		t.Fatal("process still running")
	}
}

func TestEnsureRunningAndStop(t *testing.T) {
	p := newHelper(t, "helper")
	if p.Running() {
		t.Fatal("Running before start")
	}
	startHelper(t, p)
	if !p.Running() {
		t.Fatal("not Running after EnsureRunning")
	}
	if err := p.EnsureRunning(); err != nil {
		t.Fatal("second EnsureRunning must be a no-op")
	}
	p.Stop()
	waitStopped(t, p, 5*time.Second)
}

func TestEnsureRunningPrefersExternalServer(t *testing.T) {
	// The port is already served (e.g. the ComfyUI desktop app): no child
	// is spawned, the supervisor reports ready, and Stop is a no-op.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()
	probe := func(ctx context.Context) error {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		resp.Body.Close()
		return nil
	}
	p, err := New("external", `"`+os.Args[0]+`"`, "", probe, 5*time.Second, slog.Default())
	if err != nil {
		t.Fatal(err)
	}
	if err := p.EnsureRunning(); err != nil {
		t.Fatal(err)
	}
	if p.Running() {
		t.Fatal("spawned a child even though the port is already served")
	}
	if !p.Ready() {
		t.Fatal("external server should count as ready")
	}
	p.Stop() // must not touch the external server
}

func TestWaitReady(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()
	probe := func(ctx context.Context) error {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		resp.Body.Close()
		return nil
	}
	p, err := New("ready", `"`+os.Args[0]+`"`, "", probe, 5*time.Second, slog.Default())
	if err != nil {
		t.Fatal(err)
	}
	if err := p.WaitReady(context.Background()); err != nil {
		t.Fatal(err)
	}

	failing, err := New("failing", `"`+os.Args[0]+`"`, "",
		func(context.Context) error { return errors.New("no") }, 500*time.Millisecond, slog.Default())
	if err != nil {
		t.Fatal(err)
	}
	if err := failing.WaitReady(context.Background()); err == nil {
		t.Fatal("expected timeout error from WaitReady")
	}
}

func TestWatchIdleStops(t *testing.T) {
	p := newHelper(t, "idle")
	startHelper(t, p)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go p.WatchIdle(ctx, 200*time.Millisecond, func() bool { return true })
	waitStopped(t, p, 10*time.Second)
}

func TestWatchIdleRespectsBusyGPU(t *testing.T) {
	p := newHelper(t, "busy")
	startHelper(t, p)
	defer p.Stop()
	ctx, cancel := context.WithCancel(context.Background())
	go p.WatchIdle(ctx, 100*time.Millisecond, func() bool { return false })
	time.Sleep(600 * time.Millisecond)
	cancel()
	if !p.Running() {
		t.Fatal("process was stopped while the GPU was busy")
	}
}

func TestComfyLayoutAndDefaultCommand(t *testing.T) {
	// Nested layout (ComfyUI/main.py under dir) wins.
	dir := t.TempDir()
	nested := filepath.Join(dir, "ComfyUI", "main.py")
	if err := os.MkdirAll(filepath.Dir(nested), 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(nested, []byte("x"), 0o644); err != nil {
		t.Fatal(err)
	}
	python, script := ComfyLayout("windows", dir)
	if want := filepath.Join(dir, ".venv", "Scripts", "python.exe"); python != want {
		t.Errorf("python = %s, want %s", python, want)
	}
	if script != nested {
		t.Errorf("script = %s, want %s", script, nested)
	}
	cmd := DefaultComfyCommand("windows", dir, "http://127.0.0.1:8189")
	want := `"` + filepath.Join(dir, ".venv", "Scripts", "python.exe") + `" ` + filepath.Join("ComfyUI", "main.py") + " --port 8189"
	if cmd != want {
		t.Errorf("cmd = %q, want %q", cmd, want)
	}

	// Flat layout (main.py directly under dir) is found too.
	flat := t.TempDir()
	if err := os.WriteFile(filepath.Join(flat, "main.py"), []byte("x"), 0o644); err != nil {
		t.Fatal(err)
	}
	if _, script := ComfyLayout("linux", flat); script != filepath.Join(flat, "main.py") {
		t.Errorf("flat script = %s", script)
	}
	cmd = DefaultComfyCommand("linux", flat, "http://comfy.internal")
	if strings.Contains(cmd, "--port") {
		t.Errorf("cmd = %q, want no --port for a port-less URL", cmd)
	}
	if !strings.HasSuffix(cmd, `" main.py`) {
		t.Errorf("cmd = %q, want quoted python + relative main.py", cmd)
	}

	// Neither exists yet: default to the documented nested form so the
	// startup warning points there.
	empty := t.TempDir()
	if _, script := ComfyLayout("windows", empty); script != filepath.Join(empty, "ComfyUI", "main.py") {
		t.Errorf("missing-layout script = %s", script)
	}
}

func TestStripANSI(t *testing.T) {
	cases := map[string]string{
		"\x1b[32m[INFO]\x1b[0m Starting server":   "[INFO] Starting server",
		"\x1b[1m\x1b[31m[ERROR]\x1b[0m boom":      "[ERROR] boom",
		"plain line":                              "plain line",
		"\x1b[33mWARN\x1b[0m: \x1b[1mbold\x1b[0m": "WARN: bold",
	}
	for in, want := range cases {
		if got := stripANSI(in); got != want {
			t.Errorf("stripANSI(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestDesktopSharedDir(t *testing.T) {
	root := t.TempDir()
	comfy := filepath.Join(root, "ComfyUI-Installs", "rtx5080", "ComfyUI")
	if err := os.MkdirAll(comfy, 0o755); err != nil {
		t.Fatal(err)
	}
	if got := DesktopSharedDir(comfy); got != "" {
		t.Fatalf("no shared dir yet: got %q, want empty", got)
	}
	shared := filepath.Join(root, "ComfyUI-Shared")
	if err := os.MkdirAll(filepath.Join(shared, "models"), 0o755); err != nil {
		t.Fatal(err)
	}
	if got := DesktopSharedDir(comfy); got != shared {
		t.Fatalf("got %q, want %q", got, shared)
	}
	if got := DesktopSharedDir(filepath.Join(root, "plain", "ComfyUI")); got != "" {
		t.Fatalf("non-desktop layout: got %q, want empty", got)
	}
}