package train

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
)

// fakePython stands in for an engine's interpreter: as the trainer it
// prints progress and writes an adapter; as the converter it writes a GGUF;
// told to fail, it fails.
const fakePython = `#!/bin/sh
case "$1" in
*trainer.py)
  out=$(sed -n 's/.*"out": *"\([^"]*\)".*/\1/p' "$2")
  grep -q FAIL "$(dirname "$2")/data.jsonl" && { echo "CUDA out of memory" >&2; exit 1; }
  echo '{"state":"preparing","message":"fetching base"}'
  echo '{"state":"running","step":5,"steps":10,"loss":0.9,"peak_mib":1900}'
  mkdir -p "$out"; echo '{"r":16}' > "$out/adapter_config.json"; echo weights > "$out/adapter_model.safetensors"
  echo '{"state":"trained","step":10,"steps":10,"peak_mib":2100,"base_dir":"/base"}'
  ;;
*convert_lora_to_gguf.py)
  while [ $# -gt 0 ]; do [ "$1" = "--outfile" ] && echo GGUF > "$2"; shift; done
  ;;
esac
`

func runner(t *testing.T, free int) *Runner {
	t.Helper()
	dir := t.TempDir()
	py := filepath.Join(dir, "python3")
	_ = os.WriteFile(py, []byte(fakePython), 0o755)
	uv := filepath.Join(dir, "uv")
	_ = os.WriteFile(uv, []byte("#!/bin/sh\nwhile [ $# -gt 0 ]; do [ \"$1\" = \"--target\" ] && mkdir -p \"$2\"; shift; done\n"), 0o755)
	return New(Env{WorkDir: filepath.Join(dir, "work"), Python: py, UV: uv, ConvDir: dir,
		GPUs: func(context.Context) ([]GPU, error) { return []GPU{{0, 500}, {1, free}}, nil }})
}

type sink struct {
	mu       sync.Mutex
	statuses []Status
	parts    map[string][]byte
	lasts    int
}

func (s *sink) status(st Status) { s.mu.Lock(); s.statuses = append(s.statuses, st); s.mu.Unlock() }
func (s *sink) artifact(kind string, _ int, b []byte, last bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.parts == nil {
		s.parts = map[string][]byte{}
	}
	s.parts[kind] = append(s.parts[kind], b...)
	if last {
		s.lasts++
	}
	return nil
}

func TestAJobTrainsOnTheFreestCardAndSendsWhatItMade(t *testing.T) {
	r := runner(t, 3000)
	s := &sink{}
	err := r.Run(context.Background(), Job{ID: "j1", BaseRepo: "Qwen/x", Data: []byte(`{"q":"a","a":"b"}`), GPUMemoryMiB: 2500, GGUF: true},
		s.status, s.artifact)
	if err != nil {
		t.Fatal(err)
	}
	var running bool
	for _, st := range s.statuses {
		if st.State == "running" && st.Step == 5 && st.Steps == 10 && st.PeakMiB == 1900 {
			running = true
		}
	}
	if !running {
		t.Errorf("statuses %+v", s.statuses)
	}
	if string(s.parts["gguf"]) != "GGUF\n" || s.lasts != 2 {
		t.Errorf("gguf %q, lasts %d", s.parts["gguf"], s.lasts)
	}
	files := map[string]string{}
	gz, err := gzip.NewReader(bytes.NewReader(s.parts["peft"]))
	if err != nil {
		t.Fatal(err)
	}
	tr := tar.NewReader(gz)
	for {
		h, err := tr.Next()
		if err != nil {
			break
		}
		b, _ := io.ReadAll(tr)
		files[h.Name] = string(b)
	}
	if files["adapter_model.safetensors"] != "weights\n" || files["adapter_config.json"] == "" {
		t.Errorf("peft archive %v", files)
	}
	// Nothing of the job is left on the node.
	if entries, _ := os.ReadDir(filepath.Join(r.env.WorkDir, "jobs")); len(entries) != 0 {
		t.Errorf("left behind: %v", entries)
	}
}

func TestAJobWaitsForNoCardItDoesNotFit(t *testing.T) {
	r := runner(t, 1000)
	err := r.Run(context.Background(), Job{ID: "j1", Data: []byte(`{"q":"a","a":"b"}`), GPUMemoryMiB: 2500}, (&sink{}).status, (&sink{}).artifact)
	if err == nil || !strings.Contains(err.Error(), "no card here has 2500 MiB free") {
		t.Errorf("got %v", err)
	}
}

func TestATrainerThatFailsSaysWhy(t *testing.T) {
	r := runner(t, 3000)
	err := r.Run(context.Background(), Job{ID: "j1", Data: []byte(`{"q":"FAIL","a":"b"}`)}, (&sink{}).status, (&sink{}).artifact)
	if err == nil || !strings.Contains(err.Error(), "CUDA out of memory") {
		t.Errorf("got %v", err)
	}
}

func TestOneJobAtATime(t *testing.T) {
	r := runner(t, 3000)
	r.job = "other"
	if err := r.Run(context.Background(), Job{ID: "j2"}, (&sink{}).status, (&sink{}).artifact); err != ErrBusy {
		t.Errorf("got %v", err)
	}
}

func TestCardsAreReadFromNvidiaSMI(t *testing.T) {
	got := parseGPUs("0, 2310\n1, 80000\nbad line\n")
	if len(got) != 2 || got[1].FreeMiB != 80000 || got[0].Index != 0 {
		t.Errorf("parsed %+v", got)
	}
}

func TestOnlyTheConversionCodeIsKeptOfLlamaCpp(t *testing.T) {
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gz)
	for _, name := range []string{"llama.cpp-b1/convert_lora_to_gguf.py", "llama.cpp-b1/conversion/qwen.py",
		"llama.cpp-b1/gguf-py/gguf/__init__.py", "llama.cpp-b1/src/llama.cpp", "llama.cpp-b1/../evil"} {
		_ = tw.WriteHeader(&tar.Header{Name: name, Mode: 0o644, Size: 1, Typeflag: tar.TypeReg})
		_, _ = tw.Write([]byte("x"))
	}
	_ = tw.Close()
	_ = gz.Close()
	dir := filepath.Join(t.TempDir(), "conv")
	if err := untarConversion(&buf, dir); err != nil {
		t.Fatal(err)
	}
	for _, want := range []string{"convert_lora_to_gguf.py", "conversion/qwen.py", "gguf-py/gguf/__init__.py"} {
		if _, err := os.Stat(filepath.Join(dir, want)); err != nil {
			t.Errorf("%s not kept", want)
		}
	}
	if _, err := os.Stat(filepath.Join(dir, "src")); err == nil {
		t.Error("the rest of llama.cpp was kept")
	}
}