package jupyter_test

import (
	"os"
	"strings"
	"testing"

	"github.com/jfraeys/fetch_ml/internal/jupyter"
)

func TestPrepareContainerConfig_PublicJupyter_RegistersKernelAndStartsNotebook(t *testing.T) {
	oldEnv := os.Getenv("FETCHML_JUPYTER_CONDA_ENV")
	oldKernel := os.Getenv("FETCHML_JUPYTER_KERNEL_NAME")
	_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", "base")
	_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", "python")
	t.Cleanup(func() {
		_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", oldEnv)
		_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", oldKernel)
	})

	req := &jupyter.StartRequest{
		Name:      "test",
		Workspace: "/tmp/ws",
		Image:     "quay.io/jupyter/base-notebook:latest",
		Network: jupyter.NetworkConfig{
			HostPort:       8888,
			ContainerPort:  8888,
			EnableToken:    false,
			EnablePassword: false,
		},
		Security: jupyter.SecurityConfig{AllowNetwork: true},
	}

	cfg := jupyter.PrepareContainerConfig("svc", req)
	if len(cfg.Command) < 3 {
		t.Fatalf("expected bootstrap command, got %#v", cfg.Command)
	}
	if cfg.Command[0] != "bash" || cfg.Command[1] != "-lc" {
		t.Fatalf("expected bash -lc wrapper, got %#v", cfg.Command)
	}
	script := cfg.Command[2]
	if !strings.Contains(script, "ipykernel install") {
		t.Fatalf("expected ipykernel install in script, got %q", script)
	}
	if !strings.Contains(script, "start-notebook.sh") {
		t.Fatalf("expected start-notebook.sh in script, got %q", script)
	}
	if !strings.Contains(script, "--ServerApp.token=") {
		t.Fatalf("expected token disable flags in script, got %q", script)
	}
}

func TestPrepareContainerConfig_MLToolsRunner_RegistersKernelAndStartsNotebook(t *testing.T) {
	oldEnv := os.Getenv("FETCHML_JUPYTER_CONDA_ENV")
	oldKernel := os.Getenv("FETCHML_JUPYTER_KERNEL_NAME")
	_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", "ml_env")
	_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", "ml")
	t.Cleanup(func() {
		_ = os.Setenv("FETCHML_JUPYTER_CONDA_ENV", oldEnv)
		_ = os.Setenv("FETCHML_JUPYTER_KERNEL_NAME", oldKernel)
	})

	req := &jupyter.StartRequest{
		Name:      "test",
		Workspace: "/tmp/ws",
		Image:     "localhost/ml-tools-runner:latest",
		Network: jupyter.NetworkConfig{
			HostPort:      8888,
			ContainerPort: 8888,
			EnableToken:   false,
		},
		Security: jupyter.SecurityConfig{AllowNetwork: true},
	}

	cfg := jupyter.PrepareContainerConfig("svc", req)
	if len(cfg.Command) < 3 {
		t.Fatalf("expected bootstrap command, got %#v", cfg.Command)
	}
	if cfg.Command[0] != "bash" || cfg.Command[1] != "-lc" {
		t.Fatalf("expected bash -lc wrapper, got %#v", cfg.Command)
	}
	script := cfg.Command[2]
	if !strings.Contains(script, "ipykernel install") {
		t.Fatalf("expected ipykernel install in script, got %q", script)
	}
	if !strings.Contains(script, "jupyter notebook") {
		t.Fatalf("expected jupyter notebook in script, got %q", script)
	}
	if !strings.Contains(script, "conda run -n ml_env") {
		t.Fatalf("expected conda run to use ml_env in script, got %q", script)
	}
	if strings.Contains(script, "--ServerApp.token=") {
		return
	}
	if !strings.Contains(script, "--NotebookApp.token=") {
		t.Fatalf("expected token disable flags in script, got %q", script)
	}
}