package main

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
)

// vLLM models are plain folders <vllm_dir>/<org>/<repo>. A model moves to an external disk as a whole folder
// (<disk>/wxch-models/vllm/<org>/<repo>) and the machine keeps a link at the old place. The runner mounts the external
// root into the container at the same path, so the link resolves there too (see wxch-runner, extra_roots).

func vllmPaths(cfg nodeConfig, name string) (in string, err error) {
	if !hfRepoRe.MatchString(name) {
		return "", fmt.Errorf("invalid model name")
	}
	return filepath.Join(cfg.VLLMDir, name), nil
}

func vllmDiskRoot(cfg nodeConfig, diskID string) string {
	return filepath.Join(cfg.extStore(diskID), "vllm")
}

func vllmDiskUnit(cfg nodeConfig, diskID, name string) string {
	return filepath.Join(vllmDiskRoot(cfg, diskID), name)
}

// vllmLocation: internal | external (on the disk named) | offline | "" (no such model).
func vllmLocation(cfg nodeConfig, name string) (loc, disk string) {
	in, err := vllmPaths(cfg, name)
	if err != nil {
		return "", ""
	}
	lst, err := os.Lstat(in)
	if err != nil {
		return "", ""
	}
	if lst.Mode()&os.ModeSymlink == 0 {
		return "internal", ""
	}
	tgt, _ := os.Readlink(in)
	id := diskOfPath(cfg, tgt)
	if id == "" {
		return "internal", ""
	}
	if _, err := os.Stat(filepath.Join(in, "config.json")); err != nil {
		return "offline", id
	}
	return "external", id
}

// vllmRelink makes `in` a link to target in one atomic step; a real folder at `in` is removed afterwards.
func vllmRelink(in, target string) error {
	lnk, old := in+".wxch-lnk", in+".wxch-old"
	os.Remove(lnk)
	os.RemoveAll(old)
	if err := os.Symlink(target, lnk); err != nil {
		return err
	}
	if lst, err := os.Lstat(in); err == nil && lst.Mode()&os.ModeSymlink == 0 {
		if err := os.Rename(in, old); err != nil {
			os.Remove(lnk)
			return err
		}
		if err := os.Rename(lnk, in); err != nil {
			os.Rename(old, in)
			return err
		}
		return os.RemoveAll(old)
	}
	if err := os.Rename(lnk, in); err != nil { // replaces an existing link
		os.Remove(lnk)
		return err
	}
	return nil
}

func vllmMove(cfg nodeConfig, name, to, diskID string, progress progressFn) error {
	in, err := vllmPaths(cfg, name)
	if err != nil {
		return err
	}
	lst, err := os.Lstat(in)
	if err != nil {
		return fmt.Errorf("model %s not found", name)
	}
	loc, cur := vllmLocation(cfg, name)
	if loc == "offline" {
		return fmt.Errorf("the model is on the external disk %q, which is not connected", cur)
	}
	var total, done int64
	add := func(n int64) {
		done += n
		progress("moving", total, done)
	}
	if to == "external" {
		if diskID == "" && loc == "external" {
			return nil
		}
		src := in // the real folder to copy from
		if lst.Mode()&os.ModeSymlink != 0 {
			if tgt, _ := os.Readlink(in); loc == "external" {
				src = tgt
			}
		}
		if _, err := os.Stat(filepath.Join(src, "config.json")); err != nil || hasPartial(src) {
			return fmt.Errorf("the model is not completely downloaded")
		}
		size := treeSize(src)
		target, err := pickDisk(cfg, diskID, func(d extDisk) uint64 {
			if loc == "external" && cur == d.ID {
				return 0
			}
			return uint64(size)
		})
		if err != nil {
			return err
		}
		if loc == "external" && cur == target.ID {
			return nil // already there
		}
		dst := vllmDiskUnit(cfg, target.ID, name)
		if _, err := os.Lstat(dst); err == nil {
			return fmt.Errorf("a copy already exists on the external disk %q: resolve the duplicate first", target.ID)
		}
		total = size
		if err := os.MkdirAll(filepath.Dir(dst), 0755); err != nil {
			return err
		}
		tmp := dst + ".wxch-tmp"
		os.RemoveAll(tmp)
		if _, err := copyTree(src, tmp, add); err != nil {
			os.RemoveAll(tmp)
			return err
		}
		if got := treeSize(tmp); got != total {
			os.RemoveAll(tmp)
			return fmt.Errorf("the copy has the wrong size (%d of %d bytes)", got, total)
		}
		if err := os.Rename(tmp, dst); err != nil {
			os.RemoveAll(tmp)
			return err
		}
		if err := vllmRelink(in, dst); err != nil {
			return err
		}
		if loc == "external" { // it came from another disk: that copy is not needed any more
			os.RemoveAll(src)
		}
		progress("complete", total, total)
		return nil
	}
	if loc == "internal" {
		return nil // already on the machine
	}
	src, _ := os.Readlink(in)
	total = treeSize(src)
	if free, _, err := diskFree(cfg.VLLMDir); err == nil && uint64(total)+uint64(cfg.MinFreeGB)<<30 > free {
		return fmt.Errorf("not enough room on the machine: need %.1f GB (+%d GB spare), %.1f GB free", float64(total)/1e9, cfg.MinFreeGB, float64(free)/1e9)
	}
	tmp := in + ".wxch-tmp"
	os.RemoveAll(tmp)
	if _, err := copyTree(src, tmp, add); err != nil {
		os.RemoveAll(tmp)
		return err
	}
	if got := treeSize(tmp); got != total {
		os.RemoveAll(tmp)
		return fmt.Errorf("the copy has the wrong size (%d of %d bytes)", got, total)
	}
	if err := os.Remove(in); err != nil {
		os.RemoveAll(tmp)
		return err
	}
	if err := os.Rename(tmp, in); err != nil {
		os.Symlink(src, in) // put the link back rather than leave the model missing
		return err
	}
	progress("complete", total, total)
	return os.RemoveAll(src)
}

// vllmDiskUnits lists the model folders on a disk (two levels, with a config.json).
func vllmDiskUnits(cfg nodeConfig, diskID string) []string {
	var out []string
	root := vllmDiskRoot(cfg, diskID)
	orgs, _ := os.ReadDir(root)
	for _, o := range orgs {
		if !o.IsDir() {
			continue
		}
		repos, _ := os.ReadDir(filepath.Join(root, o.Name()))
		for _, r := range repos {
			if _, err := os.Stat(filepath.Join(root, o.Name(), r.Name(), "config.json")); err == nil {
				out = append(out, o.Name()+"/"+r.Name())
			}
		}
	}
	sort.Strings(out)
	return out
}

// vllmDuplicates lists models that exist in full in two or more places (the machine, the disks).
func vllmDuplicates(cfg nodeConfig) []dupModel {
	var out []dupModel
	names := map[string]bool{}
	disks := presentDisks()
	for _, d := range disks {
		for _, n := range vllmDiskUnits(cfg, d.ID) {
			names[n] = true
		}
	}
	sorted := make([]string, 0, len(names))
	for n := range names {
		sorted = append(sorted, n)
	}
	sort.Strings(sorted)
	for _, name := range sorted {
		in, err := vllmPaths(cfg, name)
		if err != nil {
			continue
		}
		var places []string
		var size int64
		if lst, err := os.Lstat(in); err == nil && lst.Mode()&os.ModeSymlink == 0 {
			if _, err := os.Stat(filepath.Join(in, "config.json")); err == nil {
				places = append(places, "internal")
				size = treeSize(in)
			}
		}
		for _, d := range disks {
			if _, err := os.Stat(filepath.Join(vllmDiskUnit(cfg, d.ID, name), "config.json")); err == nil {
				places = append(places, d.ID)
				if size == 0 {
					size = treeSize(vllmDiskUnit(cfg, d.ID, name))
				}
			}
		}
		if len(places) >= 2 {
			out = append(out, dupModel{"vllm", name, size, places})
		}
	}
	return out
}

// vllmDedupe keeps one copy: keep is "internal" or a disk id. The other copies are deleted.
func vllmDedupe(cfg nodeConfig, name, keep string) error {
	in, err := vllmPaths(cfg, name)
	if err != nil {
		return err
	}
	disks := presentDisks()
	has := func(id string) bool {
		_, err := os.Stat(filepath.Join(vllmDiskUnit(cfg, id, name), "config.json"))
		return err == nil
	}
	lst, lerr := os.Lstat(in)
	internalReal := lerr == nil && lst.Mode()&os.ModeSymlink == 0
	if keep == "internal" {
		if !internalReal {
			if err := vllmMove(cfg, name, "internal", "", func(string, int64, int64) {}); err != nil {
				return err
			}
		}
		for _, d := range disks {
			if has(d.ID) {
				if err := os.RemoveAll(vllmDiskUnit(cfg, d.ID, name)); err != nil {
					return err
				}
			}
		}
		return nil
	}
	if d, ok := diskByID(keep); !ok || !d.Present || !has(keep) {
		return fmt.Errorf("%s is not on the connected external disk %q", name, keep)
	}
	if err := vllmRelink(in, vllmDiskUnit(cfg, keep, name)); err != nil {
		return err
	}
	for _, d := range disks {
		if d.ID != keep && has(d.ID) {
			if err := os.RemoveAll(vllmDiskUnit(cfg, d.ID, name)); err != nil {
				return err
			}
		}
	}
	return nil
}

// vllmAdoptDisk links the models of a disk into the machine's vLLM folder (nothing is copied).
func vllmAdoptDisk(cfg nodeConfig, d extDisk) {
	for _, name := range vllmDiskUnits(cfg, d.ID) {
		in, err := vllmPaths(cfg, name)
		if err != nil {
			continue
		}
		if _, err := os.Lstat(in); err == nil {
			continue
		}
		if err := os.MkdirAll(filepath.Dir(in), 0755); err != nil {
			logf("adopt %s: %v", name, err)
			continue
		}
		if err := os.Symlink(vllmDiskUnit(cfg, d.ID, name), in); err != nil {
			logf("adopt %s: %v", name, err)
			continue
		}
		logf("adopted %s from the external disk %s", name, d.ID)
	}
}

// vllmRemoveUnit removes a model; for one on a disk that means both the link and the data.
func vllmRemoveUnit(cfg nodeConfig, in string) error {
	if lst, err := os.Lstat(in); err == nil && lst.Mode()&os.ModeSymlink != 0 {
		if tgt, _ := os.Readlink(in); diskOfPath(cfg, tgt) != "" {
			if err := os.RemoveAll(tgt); err != nil {
				return err
			}
		}
		return os.Remove(in)
	}
	return os.RemoveAll(in)
}