package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"time"
)

const version = "0.14.4"

// nodeConfig is written by setup/10-models-disk.sh; defaults keep a fresh machine working.
type nodeConfig struct {
	ModelsDir string `json:"models_dir"`
	GGUFDir   string `json:"gguf_dir"`
	OllamaURL string `json:"ollama_url"`
	MinFreeGB int    `json:"min_free_gb"`
	Server    string `json:"server"`     // wss://control.example.org/ws/agent (written by the installer)
	ReserveGB int    `json:"reserve_gb"` // memory always kept free for the owner when admitting models

	// vLLM runtime (Docker): see vllm.go
	VLLMImage        string   `json:"vllm_image"`
	VLLMDir          string   `json:"vllm_dir"`
	VLLMBasePort     int      `json:"vllm_base_port"`
	VLLMMaxLen       int      `json:"vllm_max_len"`
	VLLMGPUUtil      float64  `json:"vllm_gpu_util"` // 0 = computed from the model size
	VLLMKVGB         int      `json:"vllm_kv_gb"`    // KV-cache room added to the weights when computing the reservation
	VLLMExtraArgs    []string `json:"vllm_extra_args"`
	VLLMRunnerSocket string   `json:"vllm_runner_socket"` // default /run/wxch/runner.sock

	// LM Studio runtime (llmster): see lmstudio.go
	LMStudioURL    string `json:"lmstudio_url"`
	LMStudioBin    string `json:"lmstudio_bin"`
	LMStudioModels string `json:"lmstudio_models"` // models folder (default ~/.lmstudio/models)
	LMStudioCtx    int    `json:"lmstudio_ctx"`    // context length used when loading a model
	LMStudioToken  string `json:"lmstudio_token"`  // only if LM Studio authentication is switched on

	// External model disks, added one by one with the add-disk script (see storage.go)
	ExternalRoot     string       `json:"external_root"`     // each disk is mounted at <external_root>/<id>
	ExternalDisks    []extDiskCfg `json:"external_disks"`    // the disks that were added
	ExternalRuntimes []string     `json:"external_runtimes"` // runtimes whose downloads may be placed on an external disk
	MainModels       []string     `json:"main_models"`       // the starter models: they stay on the machine, asked about for 20 s

	// Owner Guard (all optional; see guard.go for the defaults)
	// Optional AI modules (image, video, speech): see modules.go
	ModulesDir  string `json:"modules_dir"`   // default /mnt/models/modules
	EngineMode  string `json:"engine_mode"`   // "host" (default: engines run directly on the machine) or "docker" (through wxch-runner)
	ModuleIdleS int    `json:"module_idle_s"` // an engine used by jobs and then idle for this long is stopped (default 600, -1 = never)

	OwnerGuardOff    bool `json:"owner_guard_off"`
	OwnerActiveIdleS int  `json:"owner_active_idle_s"`
	OwnerReturnIdleS int  `json:"owner_return_idle_s"`
	OwnerGraceS      int  `json:"owner_grace_s"`
}

// cfgOverride lets tests inject a node configuration.
var cfgOverride *nodeConfig

func loadNodeConfig() nodeConfig {
	if cfgOverride != nil {
		return *cfgOverride
	}
	c := nodeConfig{ModelsDir: "/mnt/models/ollama", GGUFDir: "/mnt/models/gguf", OllamaURL: "http://127.0.0.1:11434", MinFreeGB: 10, ReserveGB: 8,
		VLLMImage: "nvcr.io/nvidia/vllm:26.09-py3", VLLMDir: "/mnt/models/vllm", VLLMBasePort: 8100, VLLMMaxLen: 32768, VLLMKVGB: 8,
		LMStudioURL: "http://127.0.0.1:1234", LMStudioCtx: 32768,
		ExternalRoot: "/mnt/wxch-ext", ExternalRuntimes: []string{"ollama", "vllm"},
		MainModels: []string{"qwen2.5:0.5b", "lmstudio-community/Qwen2.5-0.5B-Instruct-GGUF@Q4_K_M", "Qwen/Qwen2.5-0.5B-Instruct"}}
	if b, err := os.ReadFile("/etc/wxch/node.json"); err == nil {
		json.Unmarshal(b, &c)
	}
	return c
}

var modelNameRe = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._/-]{0,100}(:[a-zA-Z0-9._-]{1,64})?$`)

type rpcCtx struct {
	m     Msg
	send  func(Msg)
	cfg   nodeConfig
	httpc *http.Client
}

func (r *rpcCtx) ok(v any) {
	b, _ := json.Marshal(v)
	r.send(Msg{T: "rpc_result", ID: r.m.ID, OK: true, Data: b})
}
func (r *rpcCtx) fail(format string, a ...any) {
	r.send(Msg{T: "rpc_result", ID: r.m.ID, Err: fmt.Sprintf(format, a...)})
}
func (r *rpcCtx) event(v any) {
	b, _ := json.Marshal(v)
	r.send(Msg{T: "rpc_event", ID: r.m.ID, Data: b})
}

func handleRPC(m Msg, send func(Msg)) {
	r := &rpcCtx{m: m, send: send, cfg: loadNodeConfig(), httpc: &http.Client{Timeout: 15 * time.Second}}
	defer func() {
		if e := recover(); e != nil {
			r.fail("internal error: %v", e)
		}
	}()
	switch m.Method {
	case "sys.info":
		r.sysInfo()
	case "models.list":
		r.modelsList()
	case "models.pull":
		r.modelsPull()
	case "models.delete":
		r.modelsDelete()
	case "models.unload":
		r.modelsUnload()
	case "models.import":
		r.modelsImport()
	case "storage.status", "storage.move", "storage.dedupe", "storage.decide":
		r.storageRPC()
	case "modules.list", "modules.install", "modules.remove", "modules.engine", "modules.file_delete":
		r.modulesRPC()
	case "owner.status", "owner.pause", "owner.resume":
		r.ownerRPC()
	case "limits.get", "limits.set":
		r.limitsRPC()
	default:
		r.fail("unknown method %q", m.Method)
	}
}

// ---- helpers ----

func (r *rpcCtx) getJSON(path string, out any) error {
	resp, err := r.httpc.Get(r.cfg.OllamaURL + path)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
		return fmt.Errorf("ollama %s: %s", resp.Status, strings.TrimSpace(string(b)))
	}
	return json.NewDecoder(resp.Body).Decode(out)
}

func (r *rpcCtx) postJSON(method, path string, body any) (*http.Response, error) {
	b, _ := json.Marshal(body)
	req, _ := http.NewRequest(method, r.cfg.OllamaURL+path, bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	return r.httpc.Do(req)
}

func memInfo() (total, avail uint64) {
	b, err := os.ReadFile("/proc/meminfo")
	if err != nil {
		return
	}
	for _, line := range strings.Split(string(b), "\n") {
		f := strings.Fields(line)
		if len(f) < 2 {
			continue
		}
		v, _ := strconv.ParseUint(f[1], 10, 64)
		switch f[0] {
		case "MemTotal:":
			total = v * 1024
		case "MemAvailable:":
			avail = v * 1024
		}
	}
	return
}

type loadedModel struct {
	Name      string `json:"name"`
	Size      int64  `json:"size"`
	ExpiresAt string `json:"expires_at"`
}

func (r *rpcCtx) loaded() map[string]loadedModel {
	var ps struct {
		Models []loadedModel `json:"models"`
	}
	out := map[string]loadedModel{}
	if r.getJSON("/api/ps", &ps) == nil {
		for _, m := range ps.Models {
			out[m.Name] = m
		}
	}
	return out
}

// ---- methods ----

func (r *rpcCtx) sysInfo() {
	host, _ := os.Hostname()
	kernel, _ := os.ReadFile("/proc/sys/kernel/osrelease")
	total, avail := memInfo()
	info := map[string]any{
		"agent_version": version,
		"hostname":      host,
		"kernel":        strings.TrimSpace(string(kernel)),
		"arch":          runtime.GOARCH,
		"mem_total":     total,
		"mem_avail":     avail,
		"models_dir":    r.cfg.ModelsDir,
		"gguf_dir":      r.cfg.GGUFDir,
		"min_free_gb":   r.cfg.MinFreeGB,
	}
	info["models_dir_mounted"] = isMountPoint(filepath.Dir(r.cfg.ModelsDir))
	if free, tot, err := diskFree(r.cfg.ModelsDir); err == nil {
		info["disk_free"], info["disk_total"] = free, tot
	} else if free, tot, err := diskFree(filepath.Dir(r.cfg.ModelsDir)); err == nil {
		info["disk_free"], info["disk_total"] = free, tot
	}
	var v struct {
		Version string `json:"version"`
	}
	if err := r.getJSON("/api/version", &v); err == nil {
		info["ollama"] = v.Version
	} else {
		info["ollama"] = nil
		info["ollama_error"] = err.Error()
	}
	info["runtimes"] = runtimeStatuses()
	r.ok(info)
}

func (r *rpcCtx) modelsList() {
	var tags struct {
		Models []struct {
			Name       string `json:"name"`
			Size       int64  `json:"size"`
			ModifiedAt string `json:"modified_at"`
			Details    struct {
				Family string `json:"family"`
				Params string `json:"parameter_size"`
				Quant  string `json:"quantization_level"`
			} `json:"details"`
		} `json:"models"`
	}
	if err := r.getJSON("/api/tags", &tags); err != nil {
		r.fail("ollama not reachable (is it installed? run setup/30-ollama.sh): %v", err)
		return
	}
	ld := r.loaded()
	type row struct {
		Name     string `json:"name"`
		Size     int64  `json:"size"`
		Modified string `json:"modified"`
		Family   string `json:"family"`
		Params   string `json:"params"`
		Quant    string `json:"quant"`
		Loaded   bool   `json:"loaded"`
		Expires  string `json:"expires,omitempty"`
		Location string `json:"location,omitempty"` // internal | external | mixed | offline (external disk)
		Disk     string `json:"disk,omitempty"`     // the external disk, when there is one
	}
	rows := []row{}
	for _, m := range tags.Models {
		l, isLoaded := ld[m.Name]
		loc, disk := ollamaLocation(r.cfg, m.Name)
		rows = append(rows, row{m.Name, m.Size, m.ModifiedAt, m.Details.Family, m.Details.Params, m.Details.Quant, isLoaded, l.ExpiresAt, loc, disk})
	}
	type gg struct {
		Path string `json:"path"`
		Size int64  `json:"size"`
	}
	ggufs := []gg{}
	filepath.WalkDir(r.cfg.GGUFDir, func(p string, d os.DirEntry, err error) error {
		if err != nil {
			return nil
		}
		if d.IsDir() && strings.Count(strings.TrimPrefix(p, r.cfg.GGUFDir), string(os.PathSeparator)) > 3 {
			return filepath.SkipDir
		}
		if !d.IsDir() && strings.HasSuffix(strings.ToLower(p), ".gguf") {
			if fi, e := d.Info(); e == nil {
				ggufs = append(ggufs, gg{p, fi.Size()})
			}
		}
		return nil
	})
	r.ok(map[string]any{"models": rows, "gguf": ggufs, "other": otherModelRows()})
}

// streamLines runs fn for each NDJSON line of an Ollama streaming response.
func (r *rpcCtx) pullStream(name string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 12*time.Hour)
	defer cancel()
	b, _ := json.Marshal(map[string]any{"name": name, "stream": true})
	req, _ := http.NewRequestWithContext(ctx, "POST", r.cfg.OllamaURL+"/api/pull", bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	resp, err := (&http.Client{}).Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		bb, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
		return fmt.Errorf("ollama %s: %s", resp.Status, strings.TrimSpace(string(bb)))
	}
	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 1<<20), 1<<20)
	var lastStatus string
	var lastSend time.Time
	for sc.Scan() {
		var l struct {
			Status    string `json:"status"`
			Total     int64  `json:"total"`
			Completed int64  `json:"completed"`
			Error     string `json:"error"`
		}
		if json.Unmarshal(sc.Bytes(), &l) != nil {
			continue
		}
		if l.Error != "" {
			return fmt.Errorf("%s", l.Error)
		}
		if l.Status != lastStatus || time.Since(lastSend) > 500*time.Millisecond {
			r.event(map[string]any{"status": l.Status, "total": l.Total, "completed": l.Completed})
			lastStatus, lastSend = l.Status, time.Now()
		}
	}
	return sc.Err()
}

func (r *rpcCtx) modelsPull() {
	if name, rid, ok := parseRuntimeParams(r.m.Params); ok {
		if rt := otherRuntime(rid); rt != nil {
			r.pullVia(rt, name)
			return
		} else if rid != "" && rid != "ollama" {
			r.fail("unknown runtime %q", rid)
			return
		}
	}
	var p struct{ Name string }
	if json.Unmarshal(r.m.Params, &p) != nil || !modelNameRe.MatchString(p.Name) {
		r.fail("invalid model name")
		return
	}
	if free, _, err := diskFree(r.cfg.ModelsDir); err == nil && free < uint64(r.cfg.MinFreeGB)<<30 {
		r.fail("not enough free disk space: %.1f GB free, need at least %d GB (is the model disk mounted?)", float64(free)/(1<<30), r.cfg.MinFreeGB)
		return
	}
	if err := r.pullStream(p.Name); err != nil {
		r.fail("pull failed: %v", err)
		return
	}
	r.place("ollama", p.Name)
	r.ok(map[string]any{"name": p.Name})
}

func (r *rpcCtx) modelsDelete() {
	if name, rid, ok := parseRuntimeParams(r.m.Params); ok {
		if rt := otherRuntime(rid); rt != nil {
			r.deleteVia(rt, name)
			return
		}
	}
	var p struct{ Name string }
	if json.Unmarshal(r.m.Params, &p) != nil || !modelNameRe.MatchString(p.Name) {
		r.fail("invalid model name")
		return
	}
	if _, busy := r.loaded()[p.Name]; busy {
		r.fail("model is loaded in memory; unload it first")
		return
	}
	blobs, _ := ollamaBlobs(r.cfg.ModelsDir, p.Name)
	resp, err := r.postJSON("DELETE", "/api/delete", map[string]string{"name": p.Name})
	if err != nil {
		r.fail("%v", err)
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		bb, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		r.fail("delete failed: %s %s", resp.Status, strings.TrimSpace(string(bb)))
		return
	}
	ollamaForgetAfterDelete(r.cfg, p.Name, blobs)
	r.ok(map[string]any{"name": p.Name})
}

func (r *rpcCtx) modelsUnload() {
	if name, rid, ok := parseRuntimeParams(r.m.Params); ok {
		if rt := otherRuntime(rid); rt != nil {
			r.unloadVia(rt, name)
			return
		}
	}
	var p struct{ Name string }
	if json.Unmarshal(r.m.Params, &p) != nil || !modelNameRe.MatchString(p.Name) {
		r.fail("invalid model name")
		return
	}
	resp, err := r.postJSON("POST", "/api/generate", map[string]any{"model": p.Name, "keep_alive": 0})
	if err != nil {
		r.fail("%v", err)
		return
	}
	resp.Body.Close()
	if resp.StatusCode != 200 {
		r.fail("unload failed: %s", resp.Status)
		return
	}
	r.ok(map[string]any{"name": p.Name})
}

func (r *rpcCtx) modelsImport() {
	var p struct{ File, Name string }
	if json.Unmarshal(r.m.Params, &p) != nil || !modelNameRe.MatchString(p.Name) {
		r.fail("invalid model name")
		return
	}
	file, err := filepath.EvalSymlinks(p.File)
	root, err2 := filepath.EvalSymlinks(r.cfg.GGUFDir)
	if err != nil || err2 != nil || !strings.HasPrefix(file, root+string(os.PathSeparator)) ||
		!strings.HasSuffix(strings.ToLower(file), ".gguf") || strings.ContainsAny(file, "\r\n\x00") {
		r.fail("file must be a .gguf under %s", r.cfg.GGUFDir)
		return
	}
	dir, err := os.MkdirTemp("", "wxch-import")
	if err != nil {
		r.fail("%v", err)
		return
	}
	defer os.RemoveAll(dir)
	mf := filepath.Join(dir, "Modelfile")
	if err := os.WriteFile(mf, []byte("FROM "+file+"\n"), 0600); err != nil {
		r.fail("%v", err)
		return
	}
	cmd := exec.Command("ollama", "create", p.Name, "-f", mf)
	if u, err := url.Parse(r.cfg.OllamaURL); err == nil {
		cmd.Env = append(os.Environ(), "OLLAMA_HOST="+u.Host)
	}
	out, err := cmd.StdoutPipe()
	if err != nil {
		r.fail("%v", err)
		return
	}
	cmd.Stderr = cmd.Stdout
	if err := cmd.Start(); err != nil {
		r.fail("cannot run ollama: %v", err)
		return
	}
	sc := bufio.NewScanner(out)
	var lastLine string
	for sc.Scan() {
		lastLine = strings.TrimSpace(sc.Text())
		if lastLine != "" {
			r.event(map[string]any{"status": lastLine})
		}
	}
	if err := cmd.Wait(); err != nil {
		r.fail("import failed: %v %s", err, lastLine)
		return
	}
	r.ok(map[string]any{"name": p.Name})
}

// ---- self update ----

// updateExe is the file that gets replaced; a variable so a test can point it elsewhere.
var updateExe = os.Executable

// selfUpdate downloads the agent matching this machine from the control server, checks that it runs and swaps it in
// place. It returns the new version; the caller restarts (exits, so systemd starts the new file). Every step is logged,
// so "nothing happened" can be told from "it failed" in `journalctl -u wxch-agent | grep update`.
func selfUpdate(serverWS string) (string, error) {
	u, err := url.Parse(serverWS)
	if err != nil {
		return "", err
	}
	scheme := "https"
	if u.Scheme == "ws" {
		scheme = "http"
	}
	src := fmt.Sprintf("%s://%s/dl/wxch-agent-%s", scheme, u.Host, runtime.GOARCH)
	exe, err := updateExe()
	if err != nil {
		return "", err
	}
	exe, _ = filepath.EvalSymlinks(exe)
	log.Printf("update: running %s (%s), downloading %s", exe, version, src)
	resp, err := (&http.Client{Timeout: 5 * time.Minute}).Get(src)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("download %s: %s", src, resp.Status)
	}
	tmp := exe + ".new"
	f, err := os.OpenFile(tmp, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		return "", fmt.Errorf("cannot write next to %s: %w", exe, err)
	}
	n, err := io.Copy(f, io.LimitReader(resp.Body, 100<<20))
	if err != nil {
		f.Close()
		os.Remove(tmp)
		return "", err
	}
	if resp.ContentLength > 0 && n != resp.ContentLength {
		f.Close()
		os.Remove(tmp)
		return "", fmt.Errorf("the download was cut short (%d of %d bytes)", n, resp.ContentLength)
	}
	if err := f.Close(); err != nil {
		os.Remove(tmp)
		return "", err
	}
	log.Printf("update: downloaded %d bytes", n)
	out, err := exec.Command(tmp, "-version").Output()
	newVer := string(bytes.TrimSpace(out))
	if err != nil || newVer == "" {
		os.Remove(tmp)
		return "", fmt.Errorf("the downloaded agent does not run on this machine (%v)", err)
	}
	if err := os.Rename(tmp, exe); err != nil {
		os.Remove(tmp)
		return "", err
	}
	log.Printf("update: %s replaced, %s -> %s", exe, version, newVer)
	return newVer, nil
}