package main

import (
	"context"
	"encoding/json"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

const statsInterval = 5 * time.Second

// gpuProc is one process currently using the GPU for compute.
type gpuProc struct {
	PID  int
	Name string
}

type gpuStat struct {
	Util      int       `json:"util"`
	Temp      int       `json:"temp"`
	PowerW    float64   `json:"power_w"`
	SMMHz     int       `json:"sm_mhz"`
	PState    string    `json:"pstate"`
	VRAMTotal uint64    `json:"vram_total,omitempty"` // discrete GPUs only (0 on unified-memory machines)
	VRAMUsed  uint64    `json:"vram_used,omitempty"`
	Apps      []string  `json:"apps"` // processes currently using the GPU (compute)
	procs     []gpuProc // same, with pids (not sent)
}

type ollamaStat struct {
	Up     bool          `json:"up"`
	Loaded []loadedModel `json:"loaded"`
}

type modelRef struct {
	Name    string `json:"name"`
	Size    int64  `json:"size"`
	Runtime string `json:"runtime,omitempty"`
}

type rtStat struct {
	Ready bool   `json:"ready"`
	Error string `json:"error,omitempty"`
}

type diskStat struct {
	Path    string `json:"path"`
	Free    uint64 `json:"free"`
	Total   uint64 `json:"total"`
	Mounted bool   `json:"mounted"`
}

// stats is one telemetry sample. Memory is CPU+GPU unified on GB10, so mem_* is the number that matters.
type stats struct {
	TS               int64             `json:"ts"`
	Version          string            `json:"agent_version"`
	Uptime           int64             `json:"uptime_s"`
	CPUs             int               `json:"cpus"`
	CPUPct           float64           `json:"cpu_pct"`
	Load             [3]float64        `json:"load"`
	MemTotal         uint64            `json:"mem_total"`
	MemAvail         uint64            `json:"mem_avail"`
	SwapUsed         uint64            `json:"swap_used"`
	GPU              *gpuStat          `json:"gpu,omitempty"`
	TempCPU          float64           `json:"temp_cpu,omitempty"`
	TempNVMe         float64           `json:"temp_nvme,omitempty"`
	Disk             *diskStat         `json:"disk,omitempty"`
	RootDisk         *diskStat         `json:"root_disk,omitempty"` // the system disk (/), which may be a different disk from the models
	Ollama           ollamaStat        `json:"ollama"`
	Installed        []modelRef        `json:"installed"`
	LoadedAll        []modelRef        `json:"loaded_models"` // everything in memory, across runtimes
	Runtimes         map[string]rtStat `json:"runtimes"`
	Mode             string            `json:"mode"` // owner guard: available | draining | owner | paused
	ModeReason       string            `json:"mode_reason,omitempty"`
	External         []extDisk         `json:"external,omitempty"`          // the external disks that were added and how they are
	PendingPlacement []pendingInfo     `json:"pending_placement,omitempty"` // questions the owner may answer (starter models, 20 s)
	ModeKind         string            `json:"mode_kind,omitempty"`         // desktop | gpu | cooldown | paused: no program names, safe to show to other users
	ModeRetryS       int64             `json:"mode_retry_s,omitempty"`      // estimated seconds until available again if the owner stays away; 0 = unknown
	Modules          []string          `json:"modules,omitempty"`           // installed AI modules (capability ids)
	ModuleJobs       int               `json:"module_jobs,omitempty"`       // jobs running now
	Limits           *limitsStat       `json:"limits,omitempty"`            // the CPU/GPU load limits and what they are doing now
	Paused           bool              `json:"paused"`
	OwnerIdle        int64             `json:"owner_idle_s"` // seconds since keyboard/mouse on the desktop; -1 unknown
}

var (
	lastSampleMu  sync.Mutex
	lastOwnerIdle int64 = -1
	lastForeign   []string
)

type cpuTimes struct{ idle, total uint64 }

func readCPUTimes() (cpuTimes, bool) {
	b, err := os.ReadFile("/proc/stat")
	if err != nil {
		return cpuTimes{}, false
	}
	line := strings.SplitN(string(b), "\n", 2)[0]
	f := strings.Fields(line)
	if len(f) < 6 || f[0] != "cpu" {
		return cpuTimes{}, false
	}
	var t cpuTimes
	for i, s := range f[1:] {
		v, _ := strconv.ParseUint(s, 10, 64)
		t.total += v
		if i == 3 || i == 4 { // idle + iowait
			t.idle += v
		}
	}
	return t, true
}

func readFloat(path string) (float64, bool) {
	b, err := os.ReadFile(path)
	if err != nil {
		return 0, false
	}
	v, err := strconv.ParseFloat(strings.TrimSpace(string(b)), 64)
	return v, err == nil
}

func maxThermal() float64 {
	zones, _ := filepath.Glob("/sys/class/thermal/thermal_zone*/temp")
	var m float64
	for _, z := range zones {
		if v, ok := readFloat(z); ok && v/1000 > m {
			m = v / 1000
		}
	}
	return m
}

func nvmeTemp() float64 {
	names, _ := filepath.Glob("/sys/class/hwmon/hwmon*/name")
	for _, n := range names {
		if b, err := os.ReadFile(n); err == nil && strings.TrimSpace(string(b)) == "nvme" {
			if v, ok := readFloat(filepath.Join(filepath.Dir(n), "temp1_input")); ok {
				return v / 1000
			}
		}
	}
	return 0
}

func runCmd(timeout time.Duration, name string, args ...string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	out, err := exec.CommandContext(ctx, name, args...).Output()
	return string(out), err
}

func readGPU() *gpuStat {
	out, err := runCmd(3*time.Second, "nvidia-smi",
		"--query-gpu=utilization.gpu,temperature.gpu,power.draw,clocks.current.sm,pstate", "--format=csv,noheader,nounits")
	if err != nil {
		return nil
	}
	f := strings.Split(strings.TrimSpace(strings.SplitN(out, "\n", 2)[0]), ",")
	if len(f) < 5 {
		return nil
	}
	num := func(s string) float64 { v, _ := strconv.ParseFloat(strings.TrimSpace(s), 64); return v }
	g := &gpuStat{Util: int(num(f[0])), Temp: int(num(f[1])), PowerW: num(f[2]), SMMHz: int(num(f[3])), PState: strings.TrimSpace(f[4]), Apps: []string{}}
	if vm, ok := readGPUMem(); ok {
		g.VRAMTotal, g.VRAMUsed = vm.Total, vm.Used
	}
	if apps, err := runCmd(3*time.Second, "nvidia-smi", "--query-compute-apps=pid,process_name", "--format=csv,noheader"); err == nil {
		for _, l := range strings.Split(apps, "\n") {
			parts := strings.SplitN(strings.TrimSpace(l), ",", 2)
			if len(parts) != 2 {
				continue
			}
			pid, _ := strconv.Atoi(strings.TrimSpace(parts[0]))
			name := strings.TrimSpace(parts[1])
			g.Apps = append(g.Apps, filepath.Base(name))
			g.procs = append(g.procs, gpuProc{PID: pid, Name: name})
		}
	}
	return g
}

// ownerIdle asks GNOME how long the desktop has had no keyboard/mouse input (seconds), -1 if unknown.
func ownerIdle() int64 {
	bus := "/run/user/" + strconv.Itoa(os.Getuid()) + "/bus"
	if _, err := os.Stat(bus); err != nil {
		return -1
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, "gdbus", "call", "--session", "--dest", "org.gnome.Mutter.IdleMonitor",
		"--object-path", "/org/gnome/Mutter/IdleMonitor/Core", "--method", "org.gnome.Mutter.IdleMonitor.GetIdletime")
	cmd.Env = append(os.Environ(), "DBUS_SESSION_BUS_ADDRESS=unix:path="+bus)
	out, err := cmd.Output()
	if err != nil {
		return -1
	}
	// "(uint64 603487,)"
	s := string(out)
	i := strings.Index(s, "uint64 ")
	if i < 0 {
		return -1
	}
	s = s[i+len("uint64 "):]
	if j := strings.IndexAny(s, ",)"); j >= 0 {
		s = s[:j]
	}
	ms, err := strconv.ParseInt(strings.TrimSpace(s), 10, 64)
	if err != nil {
		return -1
	}
	return ms / 1000
}

func sample(cfg nodeConfig, prev *cpuTimes) stats {
	s := stats{TS: time.Now().Unix(), Version: version, CPUs: runtime.NumCPU(), OwnerIdle: -1}
	if b, err := os.ReadFile("/proc/uptime"); err == nil {
		v, _ := strconv.ParseFloat(strings.Fields(string(b))[0], 64)
		s.Uptime = int64(v)
	}
	if b, err := os.ReadFile("/proc/loadavg"); err == nil {
		f := strings.Fields(string(b))
		for i := 0; i < 3 && i < len(f); i++ {
			s.Load[i], _ = strconv.ParseFloat(f[i], 64)
		}
	}
	if t, ok := readCPUTimes(); ok {
		if prev.total > 0 && t.total > prev.total {
			s.CPUPct = 100 * (1 - float64(t.idle-prev.idle)/float64(t.total-prev.total))
		}
		*prev = t
	}
	s.MemTotal, s.MemAvail = memInfo()
	if b, err := os.ReadFile("/proc/meminfo"); err == nil {
		var st, sf uint64
		for _, l := range strings.Split(string(b), "\n") {
			f := strings.Fields(l)
			if len(f) >= 2 {
				v, _ := strconv.ParseUint(f[1], 10, 64)
				switch f[0] {
				case "SwapTotal:":
					st = v
				case "SwapFree:":
					sf = v
				}
			}
		}
		s.SwapUsed = (st - sf) * 1024
	}
	s.GPU = readGPU()
	if free, total, err := diskFree("/"); err == nil {
		s.RootDisk = &diskStat{Path: "/", Free: free, Total: total, Mounted: true}
	}
	s.TempCPU = maxThermal()
	s.TempNVMe = nvmeTemp()
	if free, total, err := diskFree(cfg.ModelsDir); err == nil {
		s.Disk = &diskStat{Path: cfg.ModelsDir, Free: free, Total: total, Mounted: isMountPoint(filepath.Dir(cfg.ModelsDir))}
	} else if free, total, err := diskFree(filepath.Dir(cfg.ModelsDir)); err == nil {
		s.Disk = &diskStat{Path: cfg.ModelsDir, Free: free, Total: total, Mounted: isMountPoint(filepath.Dir(cfg.ModelsDir))}
	}
	s.Ollama.Loaded = []loadedModel{}
	hc := &http.Client{Timeout: 2 * time.Second}
	if resp, err := hc.Get(cfg.OllamaURL + "/api/ps"); err == nil {
		var ps struct {
			Models []loadedModel `json:"models"`
		}
		if resp.StatusCode == 200 && json.NewDecoder(resp.Body).Decode(&ps) == nil {
			s.Ollama.Up = true
			if ps.Models != nil {
				s.Ollama.Loaded = ps.Models
			}
		}
		resp.Body.Close()
	}
	s.Installed, s.LoadedAll = []modelRef{}, []modelRef{}
	s.Runtimes = map[string]rtStat{}
	for _, rt := range allRuntimes {
		if err := rt.Ready(); err != nil {
			s.Runtimes[rt.ID()] = rtStat{Error: err.Error()}
			continue
		}
		s.Runtimes[rt.ID()] = rtStat{Ready: true}
		if inst, err := rt.Installed(); err == nil {
			for n, sz := range inst {
				s.Installed = append(s.Installed, modelRef{Name: n, Size: sz, Runtime: rt.ID()})
			}
		}
		if l, err := rt.Loaded(); err == nil {
			for n, sz := range l {
				s.LoadedAll = append(s.LoadedAll, modelRef{Name: n, Size: sz, Runtime: rt.ID()})
			}
		}
	}
	s.OwnerIdle = ownerIdle()

	// Owner Guard decides from the signals of this very sample
	var foreign []string
	if s.GPU != nil {
		foreign = otherGPUProcs(s.GPU.procs, procCgroup)
	}
	lastSampleMu.Lock()
	lastOwnerIdle, lastForeign = s.OwnerIdle, foreign
	lastSampleMu.Unlock()
	guardEvaluate(s.OwnerIdle, foreign)
	mode, reason, paused := guardSnapshot()
	s.Mode, s.ModeReason, s.Paused = string(mode), reason, paused
	if mode != modeAvailable {
		s.ModeKind, s.ModeRetryS = guardWhy(s.OwnerIdle, foreign)
	}
	ls := limitsSnapshot()
	s.Limits = &ls
	s.Modules, s.ModuleJobs = installedModuleIDs(cfg), jobsRunning()
	s.External = externalDisks()
	s.PendingPlacement = pendingPlacements()
	return s
}

// telemetryLoop sends a sample every few seconds until stop is closed (connection ended).
func telemetryLoop(send func(Msg), stop <-chan struct{}) {
	cfg := loadNodeConfig()
	var prev cpuTimes
	sample(cfg, &prev) // primes the CPU delta
	// first report right after connecting, so the gateway knows what this machine has without waiting a full interval
	select {
	case <-stop:
		return
	case <-time.After(1500 * time.Millisecond):
		b, _ := json.Marshal(sample(cfg, &prev))
		send(Msg{T: "stats", Data: b})
	}
	t := time.NewTicker(statsInterval)
	defer t.Stop()
	for {
		select {
		case <-stop:
			return
		case <-t.C:
			b, _ := json.Marshal(sample(cfg, &prev))
			send(Msg{T: "stats", Data: b})
		}
	}
}