package main

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

// External disks for model files, added one at a time with the add-disk script (any number of them).
//
// add-disk registers a disk in /etc/wxch/node.json (id + UUID), mounts it on demand at <external_root>/<id> with an
// automount unit and creates <mount>/wxch-models/ (a marker file and one folder per runtime). The disk keeps its other
// files untouched. The agent never mounts anything (it has no root rights): it notices the disk by its UUID link, touches
// the marker (which triggers the automount) and reports the result.
//
// Layout on a disk:  <mount>/wxch-models/.wxch-models      marker
//                    <mount>/wxch-models/ollama/blobs/     blobs of Ollama models (+ index.json: model -> blobs)
//                    <mount>/wxch-models/vllm/<org>/<repo>/  vLLM models (plain Hugging Face folders)
//
// A model "lives" on a disk when its files were moved there and the machine keeps a symbolic link in the runtime's own
// folder, so no runtime needs a settings change. When the disk is unplugged the links dangle: those models are left out
// of the installed list (so nobody is routed to them) and come back when the disk does.

const (
	extMarker   = ".wxch-models"
	extStoreDir = "wxch-models"
)

// extDiskCfg is one entry of "external_disks" in node.json.
type extDiskCfg struct {
	ID    string `json:"id"`    // short name, also the mount folder: <external_root>/<id>
	UUID  string `json:"uuid"`  // filesystem UUID (blkid)
	Label string `json:"label"` // name shown to people (default: the id)
}

// extDisk is a configured disk and what the agent sees of it right now.
type extDisk struct {
	ID      string `json:"id"`
	Label   string `json:"label"`
	Present bool   `json:"present"`
	Free    uint64 `json:"free,omitempty"`
	Total   uint64 `json:"total,omitempty"`
	Why     string `json:"why,omitempty"` // when not present: not_plugged | not_prepared
	dir     string // mount point
	store   string // <dir>/wxch-models
}

func (c nodeConfig) extRoot() string {
	if c.ExternalRoot == "" {
		return "/mnt/wxch-ext"
	}
	return c.ExternalRoot
}
func (c nodeConfig) extDiskDir(id string) string { return filepath.Join(c.extRoot(), id) }
func (c nodeConfig) extStore(id string) string   { return filepath.Join(c.extDiskDir(id), extStoreDir) }

// diskOfPath says which configured disk a path lies on ("" if none), whether or not it is connected.
func diskOfPath(cfg nodeConfig, p string) string {
	for _, d := range cfg.ExternalDisks {
		if under(p, cfg.extDiskDir(d.ID)) {
			return d.ID
		}
	}
	return ""
}

// hooks for tests
var (
	extDeviceExists = func(uuid string) bool { _, err := os.Stat("/dev/disk/by-uuid/" + uuid); return err == nil }
	extIsMount      = isMountPoint
	extProbeWait    = 20 * time.Second
)

// probeDisk looks at one disk now. Touching the marker makes systemd mount a plugged-in disk; it is done with a timeout
// so a slow disk cannot stall the agent, and only when the disk is really there (an absent disk would make systemd retry
// the mount over and over).
func probeDisk(cfg nodeConfig, c extDiskCfg) extDisk {
	d := extDisk{ID: c.ID, Label: c.Label, dir: cfg.extDiskDir(c.ID), store: cfg.extStore(c.ID)}
	if d.Label == "" {
		d.Label = c.ID
	}
	if !extDeviceExists(c.UUID) && !extIsMount(d.dir) {
		d.Why = "not_plugged"
		return d
	}
	found := make(chan bool, 1)
	go func() {
		_, err := os.Stat(filepath.Join(d.store, extMarker))
		found <- err == nil
	}()
	select {
	case ok := <-found:
		if !ok || !extIsMount(d.dir) {
			d.Why = "not_prepared"
			return d
		}
	case <-time.After(extProbeWait):
		d.Why = "not_prepared"
		return d
	}
	d.Present = true
	d.Free, d.Total, _ = diskFree(d.dir)
	return d
}

func probeAll(cfg nodeConfig) []extDisk {
	out := make([]extDisk, 0, len(cfg.ExternalDisks))
	for _, c := range cfg.ExternalDisks {
		if c.ID != "" {
			out = append(out, probeDisk(cfg, c))
		}
	}
	return out
}

var extState struct {
	sync.Mutex
	disks map[string]extDisk
}

// externalDisks is the last result of the watcher (cheap to call), sorted by id.
func externalDisks() []extDisk {
	extState.Lock()
	defer extState.Unlock()
	out := make([]extDisk, 0, len(extState.disks))
	for _, d := range extState.disks {
		out = append(out, d)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out
}

func presentDisks() []extDisk {
	var out []extDisk
	for _, d := range externalDisks() {
		if d.Present {
			out = append(out, d)
		}
	}
	return out
}

func diskByID(id string) (extDisk, bool) {
	extState.Lock()
	defer extState.Unlock()
	d, ok := extState.disks[id]
	return d, ok
}

// setExternalDisks stores a fresh probe and returns what was known before.
func setExternalDisks(list []extDisk) map[string]extDisk {
	extState.Lock()
	defer extState.Unlock()
	was := extState.disks
	extState.disks = map[string]extDisk{}
	for _, d := range list {
		extState.disks[d.ID] = d
	}
	return was
}

// watchExternal polls the disks every few seconds; plugging one in adopts the models that are on it.
func watchExternal() {
	first := true
	for {
		cfg := loadNodeConfig()
		now := probeAll(cfg)
		was := setExternalDisks(now)
		for _, d := range now {
			w, known := was[d.ID]
			if d.Present && (first || !known || !w.Present) {
				logf("external disk %s connected at %s (%.0f GB free)", d.ID, d.dir, float64(d.Free)/1e9)
				go adoptDisk(loadNodeConfig(), d)
			} else if !d.Present && known && w.Present {
				logf("external disk %s is gone (%s)", d.ID, d.Why)
			}
		}
		first = false
		time.Sleep(5 * time.Second)
	}
}

// pickDisk chooses where a download or move goes: the connected disk with the most room that can take the data
// (needFor tells how much each disk would have to receive). want names a disk ("" = any).
func pickDisk(cfg nodeConfig, want string, needFor func(extDisk) uint64) (extDisk, error) {
	var best extDisk
	found := false
	var worstNeed uint64
	disks := presentDisks()
	if want != "" {
		d, ok := diskByID(want)
		if !ok || !d.Present {
			return extDisk{}, fmt.Errorf("the external disk %q is not connected", want)
		}
		disks = []extDisk{d}
	}
	if len(disks) == 0 {
		return extDisk{}, fmt.Errorf("no external disk is connected")
	}
	for _, d := range disks {
		need := needFor(d)
		if need > worstNeed {
			worstNeed = need
		}
		if need+uint64(cfg.MinFreeGB)<<30 > d.Free && d.Free > 0 {
			continue
		}
		if !found || d.Free > best.Free {
			best, found = d, true
		}
	}
	if !found {
		return extDisk{}, fmt.Errorf("not enough room on the external disk(s): need %.1f GB (+%d GB spare)", float64(worstNeed)/1e9, cfg.MinFreeGB)
	}
	return best, nil
}

// ---- helpers shared by the runtimes' store code ----

var storageMu sync.Mutex // one move / dedupe at a time

func under(path, dir string) bool {
	rel, err := filepath.Rel(dir, path)
	return err == nil && rel != "." && !strings.HasPrefix(rel, "..") && !filepath.IsAbs(rel)
}

// copyFileSynced copies src to dst (mode 0644, so the Ollama user can read it) and syncs it to the disk.
func copyFileSynced(src, dst string, add func(int64)) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	buf := make([]byte, 4<<20)
	for {
		n, rerr := in.Read(buf)
		if n > 0 {
			if _, werr := out.Write(buf[:n]); werr != nil {
				out.Close()
				return werr
			}
			if add != nil {
				add(int64(n))
			}
		}
		if rerr == io.EOF {
			break
		}
		if rerr != nil {
			out.Close()
			return rerr
		}
	}
	if err := out.Sync(); err != nil {
		out.Close()
		return err
	}
	return out.Close()
}

// copyTree copies a folder (regular files only, symlinks are skipped) and returns the bytes copied.
func copyTree(src, dst string, add func(int64)) (int64, error) {
	var total int64
	err := filepath.WalkDir(src, func(p string, d os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		rel, _ := filepath.Rel(src, p)
		target := filepath.Join(dst, rel)
		if d.IsDir() {
			return os.MkdirAll(target, 0755)
		}
		if !d.Type().IsRegular() {
			return nil
		}
		fi, err := d.Info()
		if err != nil {
			return err
		}
		if err := copyFileSynced(p, target, add); err != nil {
			return err
		}
		total += fi.Size()
		return nil
	})
	return total, err
}

func treeSize(dir string) int64 {
	var n int64
	filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error {
		if err == nil && d.Type().IsRegular() {
			if fi, e := d.Info(); e == nil {
				n += fi.Size()
			}
		}
		return nil
	})
	return n
}

func hasStr(list []string, s string) bool {
	for _, x := range list {
		if strings.EqualFold(x, s) {
			return true
		}
	}
	return false
}

// ---- where a download goes ----

// placeAskWait is how long the owner has to answer; a variable for tests.
var placeAskWait = 20 * time.Second

// isMainModel: the starter models stay on the machine unless the owner says otherwise.
func isMainModel(cfg nodeConfig, name string) bool { return hasStr(cfg.MainModels, name) }

type placeAsk struct {
	ID       string
	Runtime  string
	Name     string
	Deadline time.Time
	ch       chan string
}

var pendingAsks = struct {
	sync.Mutex
	m map[string]*placeAsk
}{m: map[string]*placeAsk{}}

// askPlacement lets the owner choose in the web UI; with no answer in `wait` the model stays on the machine.
func askPlacement(rt, name string, wait time.Duration) string {
	id := make([]byte, 4)
	rand.Read(id)
	a := &placeAsk{ID: hex.EncodeToString(id), Runtime: rt, Name: name, Deadline: time.Now().Add(wait), ch: make(chan string, 1)}
	pendingAsks.Lock()
	pendingAsks.m[a.ID] = a
	pendingAsks.Unlock()
	defer func() {
		pendingAsks.Lock()
		delete(pendingAsks.m, a.ID)
		pendingAsks.Unlock()
	}()
	select {
	case c := <-a.ch:
		return c
	case <-time.After(wait):
		return "internal"
	}
}

func decidePlacement(id, choice string) error {
	if choice != "internal" && choice != "external" {
		return fmt.Errorf("choice must be internal or external")
	}
	pendingAsks.Lock()
	a := pendingAsks.m[id]
	pendingAsks.Unlock()
	if a == nil {
		return fmt.Errorf("the question has already been answered or timed out")
	}
	select {
	case a.ch <- choice:
	default:
	}
	return nil
}

type pendingInfo struct {
	ID      string `json:"id"`
	Runtime string `json:"runtime"`
	Name    string `json:"name"`
	LeftMS  int64  `json:"left_ms"` // time left to answer, as of this sample (the clocks of machine and browser may differ)
}

func pendingPlacements() []pendingInfo {
	pendingAsks.Lock()
	defer pendingAsks.Unlock()
	var out []pendingInfo
	for _, a := range pendingAsks.m {
		out = append(out, pendingInfo{a.ID, a.Runtime, a.Name, time.Until(a.Deadline).Milliseconds()})
	}
	return out
}

// placeAfterPull decides where a finished download should live and moves it there.
// mode: "auto" (an external disk for ordinary models, a 20 s question for the starter models), "internal" or "external".
// A failed move is reported but leaves the model working where it was downloaded.
func placeAfterPull(rtID, name, mode string, progress progressFn) error {
	cfg := loadNodeConfig()
	if mode == "internal" {
		return nil
	}
	if !hasStr(cfg.ExternalRuntimes, rtID) {
		if mode == "external" {
			return fmt.Errorf("%s models cannot be placed on an external disk yet", rtID)
		}
		return nil
	}
	if len(presentDisks()) == 0 {
		if mode == "external" {
			return fmt.Errorf("no external disk is connected")
		}
		if len(cfg.ExternalDisks) > 0 {
			progress("no external disk is connected: the model stays on the machine", 0, 0)
		}
		return nil
	}
	to := "external"
	if mode == "auto" && isMainModel(cfg, name) {
		progress("starter model: choose where to keep it (20 s, default: the machine)", 0, 0)
		to = askPlacement(rtID, name, placeAskWait)
	}
	if to == "internal" {
		progress("kept on the machine", 0, 0)
		return nil
	}
	progress("moving to an external disk", 0, 0)
	if err := moveModel(cfg, rtID, name, "external", "", progress); err != nil {
		return err
	}
	progress("stored on an external disk", 0, 0)
	return nil
}

// moveModel moves one model between the machine and an external disk (diskID "" = the best connected one), or from
// one external disk to another.
func moveModel(cfg nodeConfig, rtID, name, to, diskID string, progress progressFn) error {
	if to != "internal" && to != "external" {
		return fmt.Errorf("target must be internal or external")
	}
	if to == "external" {
		if diskID != "" {
			if d, ok := diskByID(diskID); !ok || !d.Present {
				return fmt.Errorf("the external disk %q is not connected", diskID)
			}
		} else if len(presentDisks()) == 0 {
			return fmt.Errorf("no external disk is connected")
		}
	}
	if l, err := runtimeLoaded(rtID); err == nil {
		if _, busy := l[name]; busy {
			return fmt.Errorf("the model is loaded in memory: unload it first")
		}
	}
	storageMu.Lock()
	defer storageMu.Unlock()
	switch rtID {
	case "ollama":
		return ollamaMove(cfg, name, to, diskID, progress)
	case "vllm":
		return vllmMove(cfg, name, to, diskID, progress)
	}
	return fmt.Errorf("%s models cannot be moved yet", rtID)
}

// runtimeLoaded is a variable so tests need no running runtime.
var runtimeLoaded = func(id string) (map[string]int64, error) {
	rt := runtimeByID(id)
	if rt == nil {
		return nil, fmt.Errorf("unknown runtime")
	}
	return rt.Loaded()
}

// ---- overview for the UI ----

type storedModel struct {
	Runtime  string `json:"runtime"`
	Name     string `json:"name"`
	Size     int64  `json:"size"`
	Location string `json:"location"`       // internal | external | mixed | offline
	Disk     string `json:"disk,omitempty"` // for external / offline: the disk id
	Movable  bool   `json:"movable"`
}

// dupModel is a model that exists in full in two or more places (the machine and/or external disks).
type dupModel struct {
	Runtime string   `json:"runtime"`
	Name    string   `json:"name"`
	Size    int64    `json:"size"`
	Places  []string `json:"places"` // "internal" and/or disk ids
}

// modelLocation returns where a model's files are ("" when the runtime is not handled here).
func modelLocation(cfg nodeConfig, rtID, name string) (loc, disk string) {
	switch rtID {
	case "ollama":
		return ollamaLocation(cfg, name)
	case "vllm":
		return vllmLocation(cfg, name)
	}
	return "", ""
}

func storageDuplicates(cfg nodeConfig) []dupModel {
	if len(presentDisks()) == 0 {
		return nil
	}
	return append(ollamaDuplicates(cfg), vllmDuplicates(cfg)...)
}

// adoptDisk makes the models that are on a freshly connected disk usable on this machine.
func adoptDisk(cfg nodeConfig, d extDisk) {
	storageMu.Lock()
	defer storageMu.Unlock()
	if cur, ok := diskByID(d.ID); !ok || !cur.Present {
		return
	}
	vllmAdoptDisk(cfg, d)
	ollamaAdoptDisk(cfg, d)
}