// Package train runs LoRA training jobs on a node (BACKLOG F4): with an
// installed engine's Python — its torch and transformers, already matched to
// the node's CUDA — and a small overlay of what training needs besides, so a
// node that serves can train without a second multi-gigabyte environment.
package train

import (
	"archive/tar"
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	_ "embed"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
)

//go:embed trainer.py
var trainerPy []byte

// overlayPackages are what training needs besides an engine's Python, pinned
// and installed without their dependencies: those are the engine's own, and
// letting an installer resolve them pulled a second torch.
var overlayPackages = []string{"peft==0.21.2", "accelerate==1.15.0"}

// Job is one training job, as the control plane sent it.
type Job struct {
	ID           string
	BaseRepo     string
	BaseRevision string
	// Data is NDJSON of {"q", "a"}.
	Data []byte
	// Params is the trainer's settings as JSON.
	Params       string
	GPUMemoryMiB int
	GGUF         bool
}

// Status is how a job is going.
type Status struct {
	State   string
	Message string
	Step    int
	Steps   int
	Loss    float64
	PeakMiB int
}

// Env is where a node keeps what training needs.
type Env struct {
	// EnginesDir is where engines are installed (the agent's engines_dir).
	EnginesDir string
	// WorkDir is the trainer's own directory: its overlay, the llama.cpp
	// conversion code, and each job's files while it runs.
	WorkDir string
	// RunAs is the user jobs run as: the engines' user.
	RunAs string
	Log   *slog.Logger
	// Python, UV, ConvDir and GPUs replace what would be found, for tests.
	Python, UV, ConvDir string
	GPUs                func(ctx context.Context) ([]GPU, error)
}

// GPU is one card and how much of it is free.
type GPU struct {
	Index   int
	FreeMiB int
}

// Runner runs one job at a time.
type Runner struct {
	env    Env
	mu     sync.Mutex
	cancel context.CancelFunc
	job    string
}

// New is a runner.
func New(env Env) *Runner {
	if env.Log == nil {
		env.Log = slog.New(slog.DiscardHandler)
	}
	return &Runner{env: env}
}

// Cancel stops a job if it is the one running.
func (r *Runner) Cancel(id string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.job == id && r.cancel != nil {
		r.cancel()
	}
}

// ErrBusy is a node already training.
var ErrBusy = errors.New("this node is already training")

// Run runs a job to its end, telling status as it goes and artifact each
// piece of what it made. It answers the error it failed with.
func (r *Runner) Run(ctx context.Context, job Job, status func(Status), artifact func(kind string, seq int, data []byte, last bool) error) error {
	r.mu.Lock()
	if r.job != "" {
		r.mu.Unlock()
		return ErrBusy
	}
	ctx, cancel := context.WithCancel(ctx)
	r.job, r.cancel = job.ID, cancel
	r.mu.Unlock()
	defer func() {
		cancel()
		r.mu.Lock()
		r.job, r.cancel = "", nil
		r.mu.Unlock()
	}()

	status(Status{State: "preparing", Message: "preparing the trainer"})
	py, err := r.python()
	if err != nil {
		return err
	}
	overlay, err := r.overlay(ctx, py)
	if err != nil {
		return fmt.Errorf("installing %s: %w", strings.Join(overlayPackages, ", "), err)
	}
	conv := ""
	if job.GGUF {
		if conv, err = r.conversion(ctx); err != nil {
			return fmt.Errorf("fetching llama.cpp's conversion code: %w", err)
		}
	}
	gpu, err := r.pickGPU(ctx, job.GPUMemoryMiB)
	if err != nil {
		return err
	}

	dir := filepath.Join(r.env.WorkDir, "jobs", job.ID)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	defer os.RemoveAll(dir)
	out := filepath.Join(dir, "adapter")
	var params map[string]any
	if job.Params != "" {
		if err := json.Unmarshal([]byte(job.Params), &params); err != nil {
			return fmt.Errorf("the job's params: %w", err)
		}
	}
	spec, _ := json.Marshal(map[string]any{"base_repo": job.BaseRepo, "base_revision": job.BaseRevision,
		"data": filepath.Join(dir, "data.jsonl"), "out": out, "params": params})
	for name, b := range map[string][]byte{"data.jsonl": job.Data, "job.json": spec, "trainer.py": trainerPy} {
		if err := os.WriteFile(filepath.Join(dir, name), b, 0o644); err != nil {
			return err
		}
	}

	var baseDir string
	err = r.python3(ctx, py, dir, []string{"PYTHONPATH=" + overlay, "CUDA_VISIBLE_DEVICES=" + strconv.Itoa(gpu)},
		[]string{filepath.Join(dir, "trainer.py"), filepath.Join(dir, "job.json")},
		func(line map[string]any) {
			s := Status{State: str(line["state"]), Message: str(line["message"]), Step: num(line["step"]),
				Steps: num(line["steps"]), PeakMiB: num(line["peak_mib"])}
			if v, ok := line["loss"].(float64); ok {
				s.Loss = v
			}
			if b := str(line["base_dir"]); b != "" {
				baseDir = b
			}
			if s.State == "trained" {
				s.State = "running"
				s.Message = "trained; packing the adapter"
			}
			status(s)
		})
	if err != nil {
		return err
	}

	if job.GGUF {
		status(Status{State: "running", Message: "converting the adapter to GGUF for llama.cpp"})
		gguf := filepath.Join(dir, "adapter.gguf")
		err := r.python3(ctx, py, dir, []string{"PYTHONPATH=" + strings.Join([]string{filepath.Join(conv, "gguf-py"), conv, overlay}, ":")},
			[]string{filepath.Join(conv, "convert_lora_to_gguf.py"), "--base", baseDir, "--outtype", "f16", "--outfile", gguf, out}, nil)
		if err != nil {
			return fmt.Errorf("converting to GGUF: %w", err)
		}
		if err := sendFile(gguf, "gguf", artifact); err != nil {
			return err
		}
	}
	var tgz bytes.Buffer
	if err := tarGz(&tgz, out); err != nil {
		return err
	}
	return sendBytes(tgz.Bytes(), "peft", artifact)
}

func str(v any) string {
	s, _ := v.(string)
	return s
}

func num(v any) int {
	f, _ := v.(float64)
	return int(f)
}

// python3 runs a Python script as the engines' user, passing each JSON line
// it prints to line, and answers the last of what it wrote to stderr when it
// fails.
func (r *Runner) python3(ctx context.Context, py, dir string, env, args []string, line func(map[string]any)) error {
	cmd := exec.CommandContext(ctx, py, args...)
	cmd.Dir = dir
	cmd.Env = append(os.Environ(), "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True", "TOKENIZERS_PARALLELISM=false",
		"HF_HUB_DISABLE_PROGRESS_BARS=1", "PYTHONUNBUFFERED=1")
	cmd.Env = append(cmd.Env, env...)
	home, give, why := supervisor.AsEngineUser(cmd, r.env.RunAs, r.env.WorkDir)
	if why != "" {
		r.env.Log.Warn(why)
	}
	if home != "" {
		cmd.Env = append(cmd.Env, "HF_HOME="+filepath.Join(home, ".cache", "huggingface"))
	} else {
		cmd.Env = append(cmd.Env, "HF_HOME="+filepath.Join(r.env.WorkDir, "hf"))
	}
	if err := give(dir); err != nil {
		return fmt.Errorf("handing %s to %s: %w", dir, r.env.RunAs, err)
	}
	// The whole process group goes when the job is cancelled: torch starts
	// workers of its own.
	if cmd.SysProcAttr == nil {
		cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	}
	cmd.Cancel = func() error { return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) }
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	var stderr tailBuffer
	cmd.Stderr = &stderr
	if err := cmd.Start(); err != nil {
		return err
	}
	sc := bufio.NewScanner(stdout)
	sc.Buffer(make([]byte, 64<<10), 1<<20)
	for sc.Scan() {
		var m map[string]any
		if line != nil && json.Unmarshal(sc.Bytes(), &m) == nil {
			line(m)
		}
	}
	err = cmd.Wait()
	if ctx.Err() != nil {
		return errors.New("cancelled")
	}
	if err != nil {
		return fmt.Errorf("%v: %s", err, stderr.last())
	}
	return nil
}

// tailBuffer keeps the end of what a process wrote: the reason it failed is
// at the end, and the rest is progress bars.
type tailBuffer struct{ b []byte }

func (t *tailBuffer) Write(p []byte) (int, error) {
	t.b = append(t.b, p...)
	if len(t.b) > 16<<10 {
		t.b = t.b[len(t.b)-16<<10:]
	}
	return len(p), nil
}

func (t *tailBuffer) last() string {
	lines := strings.Split(strings.TrimSpace(string(t.b)), "\n")
	if len(lines) > 3 {
		lines = lines[len(lines)-3:]
	}
	return strings.Join(lines, " | ")
}

// python is the newest installed vLLM's interpreter: its torch is built for
// this node's CUDA.
func (r *Runner) python() (string, error) {
	if r.env.Python != "" {
		return r.env.Python, nil
	}
	found, _ := filepath.Glob(filepath.Join(r.env.EnginesDir, "vllm", "*", "bin", "python3"))
	if len(found) == 0 {
		return "", errors.New("training needs vLLM installed on this node (its torch is what trains); " +
			"llmgw ctl engine-build, or serving any vLLM model here, installs it")
	}
	sort.Strings(found)
	return found[len(found)-1], nil
}

func (r *Runner) uv() (string, error) {
	if r.env.UV != "" {
		return r.env.UV, nil
	}
	var found []string
	_ = filepath.WalkDir(filepath.Join(r.env.EnginesDir, "uv"), func(p string, d os.DirEntry, err error) error {
		if err == nil && !d.IsDir() && d.Name() == "uv" {
			found = append(found, p)
		}
		return nil
	})
	if len(found) == 0 {
		return "", errors.New("no uv under " + filepath.Join(r.env.EnginesDir, "uv"))
	}
	sort.Strings(found)
	return found[len(found)-1], nil
}

// overlay installs the pinned packages once per Python and set of pins.
func (r *Runner) overlay(ctx context.Context, py string) (string, error) {
	sum := sha256.Sum256([]byte(py + "\x00" + strings.Join(overlayPackages, " ")))
	dir := filepath.Join(r.env.WorkDir, "overlay-"+hex.EncodeToString(sum[:6]))
	if _, err := os.Stat(filepath.Join(dir, ".done")); err == nil {
		return dir, nil
	}
	uv, err := r.uv()
	if err != nil {
		return "", err
	}
	tmp := dir + ".new"
	_ = os.RemoveAll(tmp)
	args := append([]string{"pip", "install", "--quiet", "--python", py, "--no-deps", "--target", tmp}, overlayPackages...)
	cmd := exec.CommandContext(ctx, uv, args...)
	if out, err := cmd.CombinedOutput(); err != nil {
		return "", fmt.Errorf("%v: %s", err, strings.TrimSpace(string(out)))
	}
	_ = os.WriteFile(filepath.Join(tmp, ".done"), nil, 0o644)
	_ = os.RemoveAll(dir)
	if err := os.Rename(tmp, dir); err != nil {
		return "", err
	}
	// Readable by the engines' user, which runs the trainer.
	_ = filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error {
		if err == nil {
			mode := os.FileMode(0o644)
			if d.IsDir() {
				mode = 0o755
			}
			_ = os.Chmod(p, mode)
		}
		return nil
	})
	return dir, nil
}

// llamaCppTag is the release of the newest llama.cpp installed: its
// conversion code is the one whose GGUF that llama.cpp reads.
func (r *Runner) llamaCppTag() (string, error) {
	found, _ := filepath.Glob(filepath.Join(r.env.EnginesDir, "llamacpp", "b*"))
	if len(found) == 0 {
		return "", errors.New("a GGUF adapter needs llama.cpp installed on this node, to match its release")
	}
	sort.Strings(found)
	tag, _, _ := strings.Cut(filepath.Base(found[len(found)-1]), "-")
	return tag, nil
}

// conversionSource is where llama.cpp's source is fetched from, by tag.
var conversionSource = "https://codeload.github.com/ggml-org/llama.cpp/tar.gz/refs/tags/"

// conversion fetches llama.cpp's LoRA conversion code for the installed
// release, once.
func (r *Runner) conversion(ctx context.Context) (string, error) {
	if r.env.ConvDir != "" {
		return r.env.ConvDir, nil
	}
	tag, err := r.llamaCppTag()
	if err != nil {
		return "", err
	}
	dir := filepath.Join(r.env.WorkDir, "llama.cpp-"+tag)
	if _, err := os.Stat(filepath.Join(dir, "convert_lora_to_gguf.py")); err == nil {
		return dir, nil
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, conversionSource+tag, nil)
	if err != nil {
		return "", err
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return "", fmt.Errorf("%s answered %s", conversionSource+tag, res.Status)
	}
	tmp := dir + ".new"
	_ = os.RemoveAll(tmp)
	if err := untarConversion(res.Body, tmp); err != nil {
		return "", err
	}
	_ = os.RemoveAll(dir)
	return dir, os.Rename(tmp, dir)
}

// untarConversion keeps, of llama.cpp's source, the conversion scripts and
// the Python packages they import.
func untarConversion(r io.Reader, dir string) error {
	gz, err := gzip.NewReader(r)
	if err != nil {
		return err
	}
	tr := tar.NewReader(gz)
	kept := 0
	for {
		h, err := tr.Next()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return err
		}
		_, rel, ok := strings.Cut(h.Name, "/")
		if !ok || h.Typeflag != tar.TypeReg || strings.Contains(rel, "..") {
			continue
		}
		if rel != "convert_lora_to_gguf.py" && rel != "convert_hf_to_gguf.py" &&
			!strings.HasPrefix(rel, "conversion/") && !strings.HasPrefix(rel, "gguf-py/") {
			continue
		}
		p := filepath.Join(dir, filepath.FromSlash(rel))
		if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
			return err
		}
		f, err := os.OpenFile(p, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
		if err != nil {
			return err
		}
		_, err = io.Copy(f, io.LimitReader(tr, 64<<20))
		f.Close()
		if err != nil {
			return err
		}
		kept++
	}
	if kept == 0 {
		return errors.New("no conversion code in llama.cpp's source")
	}
	return nil
}

// pickGPU is the card with the most memory free, if that is enough.
func (r *Runner) pickGPU(ctx context.Context, need int) (int, error) {
	list := r.env.GPUs
	if list == nil {
		list = nvidiaSMI
	}
	gpus, err := list(ctx)
	if err != nil {
		return 0, fmt.Errorf("reading the cards' free memory: %w", err)
	}
	if len(gpus) == 0 {
		return 0, errors.New("this node has no GPU to train on")
	}
	best := gpus[0]
	for _, g := range gpus[1:] {
		if g.FreeMiB > best.FreeMiB {
			best = g
		}
	}
	if best.FreeMiB < need {
		return 0, fmt.Errorf("no card here has %d MiB free (the most free is GPU %d with %d MiB); "+
			"a smaller base, or a card freed for it, would do", need, best.Index, best.FreeMiB)
	}
	return best.Index, nil
}

func nvidiaSMI(ctx context.Context) ([]GPU, error) {
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	out, err := exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=index,memory.free", "--format=csv,noheader,nounits").Output()
	if err != nil {
		return nil, err
	}
	return parseGPUs(string(out)), nil
}

func parseGPUs(s string) []GPU {
	var out []GPU
	for _, line := range strings.Split(strings.TrimSpace(s), "\n") {
		a, b, ok := strings.Cut(line, ",")
		if !ok {
			continue
		}
		i, err1 := strconv.Atoi(strings.TrimSpace(a))
		f, err2 := strconv.Atoi(strings.TrimSpace(b))
		if err1 == nil && err2 == nil {
			out = append(out, GPU{Index: i, FreeMiB: f})
		}
	}
	return out
}

// tarGz packs a directory.
func tarGz(w io.Writer, dir string) error {
	gz := gzip.NewWriter(w)
	tw := tar.NewWriter(gz)
	err := filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error {
		if err != nil || d.IsDir() {
			return err
		}
		rel, _ := filepath.Rel(dir, p)
		info, err := d.Info()
		if err != nil {
			return err
		}
		if err := tw.WriteHeader(&tar.Header{Name: filepath.ToSlash(rel), Mode: 0o644, Size: info.Size(), ModTime: info.ModTime()}); err != nil {
			return err
		}
		f, err := os.Open(p)
		if err != nil {
			return err
		}
		defer f.Close()
		_, err = io.Copy(tw, f)
		return err
	})
	if err != nil {
		return err
	}
	if err := tw.Close(); err != nil {
		return err
	}
	return gz.Close()
}

// artifactChunk is how much of an artifact goes in one message.
const artifactChunk = 1 << 20

func sendFile(path, kind string, send func(string, int, []byte, bool) error) error {
	b, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	return sendBytes(b, kind, send)
}

func sendBytes(b []byte, kind string, send func(string, int, []byte, bool) error) error {
	seq := 0
	for {
		n := min(artifactChunk, len(b))
		last := n == len(b)
		if err := send(kind, seq, b[:n], last); err != nil {
			return err
		}
		if last {
			return nil
		}
		b = b[n:]
		seq++
	}
}