package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"
)

// Ollama keeps one store: <models_dir>/manifests/<registry>/<namespace>/<name>/<tag> (small JSON) and
// <models_dir>/blobs/sha256-<hex> (the data, shared between models). A model moves to an external disk by moving its
// blobs and leaving a symbolic link under the same name in blobs/ (Ollama follows it). The manifest stays where it is.
// The agent's user has an ACL on blobs/ (set by the installer), so it can swap files there; nothing else of Ollama's
// store is touched.

var (
	ollamaDigestRe = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
	ollamaPartRe   = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$`)
)

type ollamaBlob struct {
	Digest string
	Size   int64
}

// normOllamaName gives the model its explicit tag ("qwen2.5" -> "qwen2.5:latest").
func normOllamaName(name string) string {
	if i := strings.LastIndex(name, ":"); i > 0 && !strings.Contains(name[i:], "/") {
		return name
	}
	return name + ":latest"
}

func ollamaManifestPath(root, name string) (string, error) {
	name = normOllamaName(name)
	i := strings.LastIndex(name, ":")
	base, tag := name[:i], name[i+1:]
	parts := strings.Split(base, "/")
	var rel []string
	switch len(parts) {
	case 1:
		rel = []string{"registry.ollama.ai", "library", parts[0], tag}
	case 2:
		rel = []string{"registry.ollama.ai", parts[0], parts[1], tag}
	case 3:
		rel = []string{parts[0], parts[1], parts[2], tag}
	default:
		return "", fmt.Errorf("invalid model name")
	}
	for _, p := range rel {
		if !ollamaPartRe.MatchString(p) || p == ".." {
			return "", fmt.Errorf("invalid model name")
		}
	}
	return filepath.Join(root, "manifests", filepath.FromSlash(path.Join(rel...))), nil
}

func ollamaBlobPath(root, digest string) string {
	return filepath.Join(root, "blobs", strings.Replace(digest, ":", "-", 1))
}

// ollamaBlobs reads the manifest of an installed model.
func ollamaBlobs(root, name string) ([]ollamaBlob, error) {
	mp, err := ollamaManifestPath(root, name)
	if err != nil {
		return nil, err
	}
	b, err := os.ReadFile(mp)
	if err != nil {
		return nil, err
	}
	var m struct {
		Config ollamaBlob   `json:"config"`
		Layers []ollamaBlob `json:"layers"`
	}
	if err := json.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	var out []ollamaBlob
	for _, bl := range append([]ollamaBlob{m.Config}, m.Layers...) {
		if !ollamaDigestRe.MatchString(bl.Digest) {
			return nil, fmt.Errorf("manifest has an invalid digest")
		}
		out = append(out, bl)
	}
	return out, nil
}

// ollamaBlobLoc: internal | external (disk) | offline (a link to a disk that is not connected) | missing.
func ollamaBlobLoc(cfg nodeConfig, digest string) (loc, disk string, size int64) {
	p := ollamaBlobPath(cfg.ModelsDir, digest)
	lst, err := os.Lstat(p)
	if err != nil {
		return "missing", "", 0
	}
	if lst.Mode()&os.ModeSymlink != 0 {
		tgt, _ := os.Readlink(p)
		id := diskOfPath(cfg, tgt)
		if id == "" {
			return "internal", "", lst.Size() // somebody else's link: not ours to treat as external
		}
		fi, err := os.Stat(p)
		if err != nil {
			return "offline", id, 0
		}
		return "external", id, fi.Size()
	}
	return "internal", "", lst.Size()
}

// ollamaLocation: internal | external (on the disk named) | mixed | offline | "" (not an installed model).
func ollamaLocation(cfg nodeConfig, name string) (loc, disk string) {
	bl, err := ollamaBlobs(cfg.ModelsDir, name)
	if err != nil {
		return "", ""
	}
	ext, off := 0, 0
	disks := map[string]bool{}
	var firstOff, firstDisk string
	for _, b := range bl {
		switch l, d, _ := ollamaBlobLoc(cfg, b.Digest); l {
		case "external":
			ext++
			disks[d] = true
			firstDisk = d
		case "offline":
			off++
			if firstOff == "" {
				firstOff = d
			}
		}
	}
	switch {
	case off > 0:
		return "offline", firstOff
	case ext == 0:
		return "internal", ""
	case ext == len(bl) && len(disks) == 1:
		return "external", firstDisk
	}
	return "mixed", ""
}

// ollamaDropOffline removes the models whose files are on a disk that is not connected.
func ollamaDropOffline(cfg nodeConfig, inst map[string]int64) map[string]int64 {
	for name := range inst {
		if loc, _ := ollamaLocation(cfg, name); loc == "offline" {
			delete(inst, name)
		}
	}
	return inst
}

// ---- the index on each disk: which models' blobs it holds (lets a reinstalled machine take them over) ----

type ollamaIndex struct {
	Models map[string][]string `json:"models"` // model -> blob digests
}

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

func ollamaExtBlob(cfg nodeConfig, diskID, digest string) string {
	return filepath.Join(ollamaStore(cfg, diskID), "blobs", strings.Replace(digest, ":", "-", 1))
}

func loadOllamaIndex(cfg nodeConfig, diskID string) ollamaIndex {
	idx := ollamaIndex{Models: map[string][]string{}}
	if b, err := os.ReadFile(filepath.Join(ollamaStore(cfg, diskID), "index.json")); err == nil {
		json.Unmarshal(b, &idx)
		if idx.Models == nil {
			idx.Models = map[string][]string{}
		}
	}
	return idx
}

func saveOllamaIndex(cfg nodeConfig, diskID string, idx ollamaIndex) error {
	dir := ollamaStore(cfg, diskID)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}
	b, _ := json.MarshalIndent(idx, "", "  ")
	tmp := filepath.Join(dir, "index.json.tmp")
	if err := os.WriteFile(tmp, append(b, '\n'), 0644); err != nil {
		return err
	}
	return os.Rename(tmp, filepath.Join(dir, "index.json"))
}

func (idx ollamaIndex) usedByOthers(name, digest string) bool {
	for n, ds := range idx.Models {
		if n == name {
			continue
		}
		for _, d := range ds {
			if d == digest {
				return true
			}
		}
	}
	return false
}

// ---- moving ----

// ollamaMove moves a model to an external disk (diskID "" = the best one; a model that already is on a disk stays
// there then) or back to the machine.
func ollamaMove(cfg nodeConfig, name, to, diskID string, progress progressFn) error {
	name = normOllamaName(name)
	blobs, err := ollamaBlobs(cfg.ModelsDir, name)
	if err != nil {
		return fmt.Errorf("model %s not found in Ollama", name)
	}
	loc, cur := ollamaLocation(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)
	}
	digests := make([]string, 0, len(blobs))
	for _, b := range blobs {
		digests = append(digests, b.Digest)
	}

	if to == "external" {
		if diskID == "" && loc == "external" {
			return nil // already on a disk and nobody asked for a particular one
		}
		needFor := func(d extDisk) uint64 {
			var need int64
			for _, b := range blobs {
				if l, bd, _ := ollamaBlobLoc(cfg, b.Digest); l == "external" && bd == d.ID {
					continue
				}
				if fi, err := os.Stat(ollamaExtBlob(cfg, d.ID, b.Digest)); err == nil && fi.Size() == b.Size {
					continue
				}
				need += b.Size
			}
			return uint64(need)
		}
		target, err := pickDisk(cfg, diskID, needFor)
		if err != nil {
			return err
		}
		total = int64(needFor(target))
		if err := os.MkdirAll(filepath.Join(ollamaStore(cfg, target.ID), "blobs"), 0755); err != nil {
			return err
		}
		old := map[string]bool{}
		for _, b := range blobs {
			if _, bd, _ := ollamaBlobLoc(cfg, b.Digest); bd != "" && bd != target.ID {
				old[bd] = true
			}
			if err := ollamaBlobToDisk(cfg, target.ID, b.Digest, add); err != nil {
				return fmt.Errorf("moving %s: %v", b.Digest[:19], err)
			}
		}
		idx := loadOllamaIndex(cfg, target.ID)
		idx.Models[name] = digests
		if err := saveOllamaIndex(cfg, target.ID, idx); err != nil {
			return err
		}
		for od := range old { // the copy on the disk it came from is not needed any more
			ollamaDropCopy(cfg, od, name)
		}
		progress("complete", total, total)
		return nil
	}

	if loc == "internal" {
		return nil // already on the machine
	}
	disks := map[string]bool{}
	for _, b := range blobs {
		if l, bd, size := ollamaBlobLoc(cfg, b.Digest); l == "external" {
			total += size
			disks[bd] = true
		}
	}
	if free, _, err := diskFree(cfg.ModelsDir); 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)
	}
	for _, b := range blobs {
		if err := ollamaBlobToInternal(cfg, b.Digest, add); err != nil {
			return fmt.Errorf("moving %s: %v", b.Digest[:19], err)
		}
	}
	for d := range disks {
		ollamaDropCopy(cfg, d, name)
	}
	progress("complete", total, total)
	return nil
}

// ollamaBlobToDisk copies one blob to the disk (unless an identical one is already there) and swaps the machine's
// file or link for a link to it in one atomic rename, so Ollama never sees the blob missing.
func ollamaBlobToDisk(cfg nodeConfig, diskID, digest string, add func(int64)) error {
	p := ollamaBlobPath(cfg.ModelsDir, digest)
	lst, err := os.Lstat(p)
	if err != nil {
		return err
	}
	src := p
	if lst.Mode()&os.ModeSymlink != 0 {
		tgt, _ := os.Readlink(p)
		id := diskOfPath(cfg, tgt)
		if id == "" || id == diskID {
			return nil // already there (or a link that is not ours)
		}
		src = tgt
	}
	sfi, err := os.Stat(src)
	if err != nil {
		return err
	}
	dst := ollamaExtBlob(cfg, diskID, digest)
	if fi, err := os.Stat(dst); err != nil || fi.Size() != sfi.Size() {
		tmp := dst + ".wxch-tmp"
		if err := copyFileSynced(src, tmp, add); err != nil {
			os.Remove(tmp)
			return err
		}
		if err := os.Rename(tmp, dst); err != nil {
			os.Remove(tmp)
			return err
		}
	}
	lnk := p + ".wxch-lnk"
	os.Remove(lnk)
	if err := os.Symlink(dst, lnk); err != nil {
		return err
	}
	if err := os.Rename(lnk, p); err != nil {
		os.Remove(lnk)
		return err
	}
	return nil
}

// ollamaBlobToInternal copies a blob back from its disk and replaces the link with the real file.
func ollamaBlobToInternal(cfg nodeConfig, digest string, add func(int64)) error {
	p := ollamaBlobPath(cfg.ModelsDir, digest)
	lst, err := os.Lstat(p)
	if err != nil || lst.Mode()&os.ModeSymlink == 0 {
		return nil // already a real file
	}
	tgt, _ := os.Readlink(p)
	if diskOfPath(cfg, tgt) == "" {
		return nil
	}
	tmp := p + ".wxch-tmp"
	if err := copyFileSynced(tgt, tmp, add); err != nil {
		os.Remove(tmp)
		return err
	}
	if err := os.Rename(tmp, p); err != nil {
		os.Remove(tmp)
		return err
	}
	return nil
}

// ollamaDropCopy deletes a model's copy on a disk (and its index entry) unless the machine still uses it or another
// model on that disk shares the blob.
func ollamaDropCopy(cfg nodeConfig, diskID, name string) {
	if d, ok := diskByID(diskID); !ok || !d.Present {
		return
	}
	idx := loadOllamaIndex(cfg, diskID)
	digests, ok := idx.Models[name]
	if !ok {
		return
	}
	for _, dg := range digests {
		if idx.usedByOthers(name, dg) {
			continue
		}
		if lst, err := os.Lstat(ollamaBlobPath(cfg.ModelsDir, dg)); err == nil && lst.Mode()&os.ModeSymlink != 0 {
			if tgt, _ := os.Readlink(ollamaBlobPath(cfg.ModelsDir, dg)); diskOfPath(cfg, tgt) == diskID {
				continue // the machine's model still runs from this copy
			}
		}
		os.Remove(ollamaExtBlob(cfg, diskID, dg))
	}
	delete(idx.Models, name)
	saveOllamaIndex(cfg, diskID, idx)
}

// ---- duplicates, adoption, cleanup ----

func hasFullCopy(cfg nodeConfig, diskID string, blobs []ollamaBlob) bool {
	for _, b := range blobs {
		if fi, err := os.Stat(ollamaExtBlob(cfg, diskID, b.Digest)); err != nil || fi.Size() != b.Size {
			return false
		}
	}
	return len(blobs) > 0
}

// ollamaDuplicates lists models that exist in full in two or more places (the machine, the disks).
func ollamaDuplicates(cfg nodeConfig) []dupModel {
	disks := presentDisks()
	names := map[string]bool{}
	for _, d := range disks {
		for n := range loadOllamaIndex(cfg, d.ID).Models {
			names[n] = true
		}
	}
	sorted := make([]string, 0, len(names))
	for n := range names {
		sorted = append(sorted, n)
	}
	sort.Strings(sorted)
	var out []dupModel
	for _, name := range sorted {
		blobs, err := ollamaBlobs(cfg.ModelsDir, name)
		if err != nil {
			continue // not registered with Ollama: adoption's business, not a duplicate
		}
		var places []string
		if loc, _ := ollamaLocation(cfg, name); loc == "internal" {
			places = append(places, "internal")
		}
		for _, d := range disks {
			if hasFullCopy(cfg, d.ID, blobs) {
				places = append(places, d.ID)
			}
		}
		if len(places) >= 2 {
			var size int64
			for _, b := range blobs {
				size += b.Size
			}
			out = append(out, dupModel{"ollama", name, size, places})
		}
	}
	return out
}

// ollamaDedupe keeps one copy: keep is "internal" or a disk id. The other copies are deleted.
func ollamaDedupe(cfg nodeConfig, name, keep string) error {
	name = normOllamaName(name)
	quiet := func(string, int64, int64) {}
	if keep == "internal" {
		if err := ollamaMove(cfg, name, "internal", "", quiet); err != nil {
			return err
		}
	} else {
		if d, ok := diskByID(keep); !ok || !d.Present {
			return fmt.Errorf("the external disk %q is not connected", keep)
		}
		if _, ok := loadOllamaIndex(cfg, keep).Models[name]; !ok {
			return fmt.Errorf("%s is not on the external disk %q", name, keep)
		}
		if err := ollamaMove(cfg, name, "external", keep, quiet); err != nil {
			return err
		}
	}
	for _, d := range presentDisks() {
		if d.ID != keep {
			ollamaDropCopy(cfg, d.ID, name)
		}
	}
	return nil
}

// ollamaAdoptDisk makes the models of a disk visible to Ollama when it has no manifest for them (after a reinstall or on
// another machine): the blobs get links, then Ollama is asked to pull, which downloads only the small manifest because
// every blob is already there.
func ollamaAdoptDisk(cfg nodeConfig, d extDisk) {
	rt := ollamaRuntime{}
	idx := loadOllamaIndex(cfg, d.ID)
	names := make([]string, 0, len(idx.Models))
	for n := range idx.Models {
		names = append(names, n)
	}
	sort.Strings(names)
	for _, name := range names {
		if _, err := ollamaBlobs(cfg.ModelsDir, name); err == nil {
			continue
		}
		complete := true
		for _, dg := range idx.Models[name] {
			p := ollamaBlobPath(cfg.ModelsDir, dg)
			if _, err := os.Lstat(p); err == nil {
				continue
			}
			ext := ollamaExtBlob(cfg, d.ID, dg)
			if _, err := os.Stat(ext); err != nil {
				complete = false
				continue
			}
			if err := os.Symlink(ext, p); err != nil {
				logf("adopt %s: cannot link %s: %v", name, dg[:19], err)
				complete = false
			}
		}
		if !complete || rt.Ready() != nil {
			continue
		}
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
		err := rt.Pull(ctx, name, func(string, int64, int64) {})
		cancel()
		if err != nil {
			logf("adopt %s: ollama could not register it: %v", name, err)
			continue
		}
		logf("adopted %s from the external disk %s", name, d.ID)
	}
}

// ollamaForgetAfterDelete removes the copies on the disks of a model Ollama has just deleted.
func ollamaForgetAfterDelete(cfg nodeConfig, name string, blobs []ollamaBlob) {
	name = normOllamaName(name)
	for _, d := range presentDisks() {
		idx := loadOllamaIndex(cfg, d.ID)
		if _, ok := idx.Models[name]; !ok {
			continue
		}
		for _, b := range blobs {
			if _, err := os.Lstat(ollamaBlobPath(cfg.ModelsDir, b.Digest)); err == nil {
				continue // still used by another model
			}
			if !idx.usedByOthers(name, b.Digest) {
				os.Remove(ollamaExtBlob(cfg, d.ID, b.Digest))
			}
		}
		delete(idx.Models, name)
		saveOllamaIndex(cfg, d.ID, idx)
	}
}