package main

import (
	"encoding/json"
	"fmt"
	"sort"
	"time"
)

// RPC methods for the external disks (see storage.go).

func (r *rpcCtx) storageRPC() {
	switch r.m.Method {
	case "storage.status":
		r.ok(storageOverview(r.cfg))
	case "storage.move":
		var p struct{ Runtime, Name, To, Disk string }
		if json.Unmarshal(r.m.Params, &p) != nil || p.Name == "" {
			r.fail("invalid parameters")
			return
		}
		var last time.Time
		err := moveModel(r.cfg, p.Runtime, p.Name, p.To, p.Disk, func(status string, total, done int64) {
			if status == "complete" || time.Since(last) > 500*time.Millisecond {
				r.event(map[string]any{"status": status, "total": total, "completed": done})
				last = time.Now()
			}
		})
		forgetVLLMSizes()
		if err != nil {
			r.fail("%v", err)
			return
		}
		r.ok(map[string]any{"name": p.Name, "to": p.To})
	case "storage.dedupe":
		var p struct{ Runtime, Name, Keep string }
		if json.Unmarshal(r.m.Params, &p) != nil || p.Keep == "" {
			r.fail("keep must be internal or the id of a disk")
			return
		}
		if len(presentDisks()) == 0 {
			r.fail("no external disk is connected")
			return
		}
		storageMu.Lock()
		var err error
		switch p.Runtime {
		case "ollama":
			err = ollamaDedupe(r.cfg, p.Name, p.Keep)
		case "vllm":
			err = vllmDedupe(r.cfg, p.Name, p.Keep)
		default:
			err = fmt.Errorf("%s models have no duplicates handling yet", p.Runtime)
		}
		storageMu.Unlock()
		forgetVLLMSizes()
		if err != nil {
			r.fail("%v", err)
			return
		}
		r.ok(map[string]any{"name": p.Name, "kept": p.Keep})
	case "storage.decide":
		var p struct{ ID, Choice string }
		if json.Unmarshal(r.m.Params, &p) != nil {
			r.fail("invalid parameters")
			return
		}
		if err := decidePlacement(p.ID, p.Choice); err != nil {
			r.fail("%v", err)
			return
		}
		r.ok(map[string]any{"id": p.ID, "choice": p.Choice})
	}
}

// storageOverview is everything the Models page needs to show where models live.
func storageOverview(cfg nodeConfig) map[string]any {
	disks := externalDisks()
	models := []storedModel{}
	if inst, err := (ollamaBackend{}).installed(); err == nil {
		for name, size := range inst {
			loc, disk := ollamaLocation(cfg, name)
			models = append(models, storedModel{"ollama", name, size, loc, disk, movable(cfg, "ollama", loc)})
		}
	}
	if rt := runtimeByID("vllm"); rt != nil && rt.Ready() == nil {
		if inst, err := rt.Installed(); err == nil {
			for name, size := range inst {
				loc, disk := vllmLocation(cfg, name)
				models = append(models, storedModel{"vllm", name, size, loc, disk, movable(cfg, "vllm", loc)})
			}
		}
	}
	sort.Slice(models, func(i, j int) bool {
		if models[i].Runtime != models[j].Runtime {
			return models[i].Runtime < models[j].Runtime
		}
		return models[i].Name < models[j].Name
	})
	return map[string]any{"disks": disks, "configured": len(cfg.ExternalDisks) > 0, "models": models, "duplicates": storageDuplicates(cfg), "pending": pendingPlacements()}
}

func movable(cfg nodeConfig, rt, loc string) bool {
	return len(presentDisks()) > 0 && hasStr(cfg.ExternalRuntimes, rt) && (loc == "internal" || loc == "external" || loc == "mixed")
}

// forgetVLLMSizes drops the cached folder sizes after a move.
func forgetVLLMSizes() {
	if v, ok := runtimeByID("vllm").(*vllmRuntime); ok {
		v.mu.Lock()
		v.sizes = nil
		v.mu.Unlock()
	}
}