package cli

import (
	"bytes"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// TestMain lets the test binary stand in for llama-server: run with
// LLMGW_FAKE_LLAMA set, it prints what llama.cpp prints and answers.
func TestMain(m *testing.M) {
	if mode := os.Getenv("LLMGW_FAKE_LLAMA"); mode != "" {
		fakeLlama(mode)
		return
	}
	os.Exit(m.Run())
}

func fakeLlama(mode string) {
	port := ""
	for i, a := range os.Args {
		if a == "--port" && i+1 < len(os.Args) {
			port = os.Args[i+1]
		}
	}
	if mode == "gpu" {
		fmt.Println("llama_model_load_from_file_impl: using device CUDA0 (Quadro RTX 6000) - 23800 MiB free")
		fmt.Println("load_tensors: offloaded 25/25 layers to GPU")
	} else {
		fmt.Println("warning: no usable GPU found, --gpu-layers option will be ignored")
		fmt.Println("load_tensors: offloaded 0/25 layers to GPU")
	}
	http.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })
	http.HandleFunc("/completion", func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte(`{"content":"!"}`)) })
	_ = http.ListenAndServe("127.0.0.1:"+port, nil) //nolint:gosec // a test's stand-in
}

func TestEngineCheckComputesATokenOnTheGPU(t *testing.T) {
	// nvidia-smi present, so the check runs.
	dir := t.TempDir()
	smi := filepath.Join(dir, "nvidia-smi")
	if err := os.WriteFile(smi, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil {
		t.Fatal(err)
	}
	t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH"))
	model := filepath.Join(dir, "tiny.gguf")
	_ = os.WriteFile(model, []byte("GGUF"), 0o644)
	self, _ := os.Executable()

	t.Setenv("LLMGW_FAKE_LLAMA", "gpu")
	var out bytes.Buffer
	if err := computeOnGPU(t.Context(), self, model, &out); err != nil || !strings.Contains(out.String(), "on CUDA") ||
		!strings.Contains(out.String(), "25 of 25") {
		t.Fatalf("on the GPU: %v %s", err, out.String())
	}

	t.Setenv("LLMGW_FAKE_LLAMA", "cpu")
	out.Reset()
	if err := computeOnGPU(t.Context(), self, model, &out); err == nil || !strings.Contains(err.Error(), "computed on the CPU") {
		t.Fatalf("a build that falls back to the CPU: %v %s", err, out.String())
	}
}