package cli

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

// A node's overview: what a bare `sudo llmgw` opens on a machine that runs
// only llmgw-agent. Everything on it is read here — the agent's own files,
// systemd, the cards and the disks — because a node has no admin token, and
// the control plane may be the very thing it cannot reach.

// nodeCmdTimeout bounds every command and check the overview runs: a screen
// refreshed every three seconds cannot wait on a tool that hangs.
const nodeCmdTimeout = 2 * time.Second

// nodeControlPlanesFile is where the agent keeps the control planes the
// leader last named (controlPlanesFile in the agent package).
var nodeControlPlanesFile = "/var/lib/llmgw-agent/control-planes.json"

// nodeProblemLines is how many of the agent's warnings the overview shows.
const nodeProblemLines = 6

type nodeOverviewMsg nodeOverview

// nodeOverview is everything the overview shows, read in one go.
type nodeOverview struct {
	at      time.Time
	version string
	agent   agentService
	planes  []planeCheck
	gpus    []nodeGPU
	// gpuNote says why there are no GPU rows: no card, or no answer.
	gpuNote string
	// names is the cards' names by index, which do not change and are asked
	// for once rather than every refresh.
	names       map[int]string
	engines     []engineRun
	enginesNote string
	disks       []diskUse
	problems    []string
	// problemsNote replaces the problems when the journal cannot be read.
	problemsNote string
}

// agentService is what systemd says of llmgw-agent.
type agentService struct {
	state, sub string
	since      time.Time
	// note replaces the rest when systemd cannot say.
	note string
}

// planeCheck is one control plane the agent dials, and whether it answers.
type planeCheck struct {
	addr string
	// learned is an address the leader named rather than one in agent.yaml.
	learned bool
	err     string
}

type nodeGPU struct {
	index int
	name  string
	// util is -1 where the driver does not report it.
	util                  float64
	memUsedMB, memTotalMB float64
	tempC                 float64
	models                []string
}

// engineRun is one engine the supervisor started here, from its state file.
type engineRun struct {
	replica, model, engine string
	gpus                   []int
	port, pid              int
	alive                  bool
	started                time.Time
}

type diskUse struct {
	what, path  string
	free, total uint64
	err         string
}

// nodeProbe reads this machine. Every path and tool is a field, so a test
// reads a directory of its own instead.
type nodeProbe struct {
	configured []string
	planesFile string
	workDir    string
	weightsDir string
	enginesDir string
	root       bool
	names      map[int]string

	run    func(ctx context.Context, name string, args ...string) ([]byte, error)
	dial   func(ctx context.Context, addr string) error
	gpus   func(ctx context.Context) []discovery.GPUSensor
	cards  func(ctx context.Context) ([]discovery.GPU, error)
	alive  func(pid int) bool
	statfs func(path string) (free, total uint64, err error)
	now    func() time.Time
}

// newNodeProbe reads the machine agent.yaml describes, as the agent would.
func newNodeProbe(cfg any, names map[int]string) nodeProbe {
	c, ok := cfg.(agent.Config)
	if !ok {
		c = agent.DefaultConfig()
	}
	var configured []string
	if c.Server != "" {
		configured = append(configured, c.Server)
	}
	configured = append(configured, c.Servers...)
	// The same defaults the agent applies as it starts.
	weights := c.WeightsCache
	if weights == "" && len(c.ModelDirs) > 0 {
		weights = c.ModelDirs[0]
	}
	if weights == "" {
		weights = agent.DefaultWeightsCache
	}
	enginesDir := c.EnginesDir
	if enginesDir == "" {
		enginesDir = provision.DefaultRoot
	}
	work := c.Superv.WorkDir
	if work == "" {
		work = "/var/lib/llmgw-agent"
	}
	return nodeProbe{
		configured: configured, planesFile: nodeControlPlanesFile, workDir: work,
		weightsDir: weights, enginesDir: enginesDir, root: os.Geteuid() == 0, names: names,
		run: runLocal, dial: dialTCP, gpus: discovery.SampleGPUs,
		cards: func(ctx context.Context) ([]discovery.GPU, error) { return discovery.DiscoverGPUsSMI(ctx, "") },
		alive: pidAlive, statfs: statfs, now: time.Now,
	}
}

// fetchNodeOverview reads the machine in the background.
func (m consoleModel) fetchNodeOverview() tea.Cmd {
	p := newNodeProbe(m.cfg, m.nodeOv.names)
	return func() tea.Msg { return nodeOverviewMsg(p.gather(context.Background())) }
}

// nodeTick refreshes the overview while it is on screen, one reading at a
// time; the other screens of a node watch nothing.
func (m consoleModel) nodeTick() (consoleModel, tea.Cmd) {
	tick := tea.Tick(3*time.Second, func(time.Time) tea.Msg { return consoleTickMsg{} })
	if m.menu != cOverview || m.nodeBusy {
		return m, tick
	}
	m.nodeBusy = true
	return m, tea.Batch(m.fetchNodeOverview(), tick)
}

// gather reads everything at once, each part in its own goroutine, so the
// slowest part bounds the whole.
func (p nodeProbe) gather(ctx context.Context) nodeOverview {
	ov := nodeOverview{version: version.Version}
	var wg sync.WaitGroup
	do := func(f func()) {
		wg.Add(1)
		go func() { defer wg.Done(); f() }()
	}
	var engines []engineRun
	do(func() { ov.agent = p.agentState(ctx) })
	do(func() { ov.planes = p.checkPlanes(ctx) })
	do(func() { ov.gpus, ov.names, ov.gpuNote = p.readGPUs(ctx) })
	do(func() { engines, ov.enginesNote = readShardStates(p.workDir, p.alive) })
	do(func() { ov.disks = p.readDisks() })
	do(func() { ov.problems, ov.problemsNote = p.readProblems(ctx) })
	wg.Wait()
	ov.engines = engines
	ov.gpus = modelsOnGPUs(ov.gpus, engines)
	ov.at = p.now()
	return ov
}

// --- the agent ---

func (p nodeProbe) agentState(ctx context.Context) agentService {
	out, err := p.run(ctx, "systemctl", "show", "llmgw-agent",
		"-p", "LoadState", "-p", "ActiveState", "-p", "SubState", "-p", "ActiveEnterTimestamp")
	if err != nil && len(out) == 0 {
		return agentService{note: "systemctl is not available: " + err.Error()}
	}
	return parseSystemctlShow(string(out))
}

// parseSystemctlShow reads `systemctl show -p …`: one key=value a line.
func parseSystemctlShow(out string) agentService {
	kv := map[string]string{}
	for _, line := range strings.Split(out, "\n") {
		if k, v, ok := strings.Cut(strings.TrimSpace(line), "="); ok {
			kv[k] = v
		}
	}
	if kv["LoadState"] == "not-found" {
		return agentService{note: "llmgw-agent is not installed as a service here"}
	}
	s := agentService{state: kv["ActiveState"], sub: kv["SubState"]}
	if s.state == "" {
		s.note = "systemd does not say"
	}
	s.since = parseSystemdTime(kv["ActiveEnterTimestamp"])
	return s
}

// parseSystemdTime reads systemd's "Mon 2026-10-19 08:12:03 CEST", in this
// machine's zone; zero when it is empty or n/a.
func parseSystemdTime(v string) time.Time {
	v = strings.TrimSpace(v)
	if v == "" || v == "n/a" {
		return time.Time{}
	}
	for _, layout := range []string{"Mon 2006-01-02 15:04:05 MST", "Mon 2006-01-02 15:04:05 -0700", "Mon 2006-01-02 15:04:05"} {
		if t, err := time.ParseInLocation(layout, v, time.Local); err == nil {
			return t
		}
	}
	return time.Time{}
}

// --- the control planes ---

// readControlPlanes is every address the agent dials, in its order: the
// configured ones, then those the leader named that agent.yaml lacks.
func readControlPlanes(file string, configured []string) []planeCheck {
	var out []planeCheck
	seen := map[string]bool{}
	add := func(s string, learned bool) {
		hp := hostPort(s)
		if hp == "" || seen[hp] {
			return
		}
		seen[hp] = true
		out = append(out, planeCheck{addr: hp, learned: learned})
	}
	for _, s := range configured {
		add(s, false)
	}
	if b, err := os.ReadFile(file); err == nil {
		var told []string
		if json.Unmarshal(b, &told) == nil {
			for _, s := range told {
				add(s, true)
			}
		}
	}
	return out
}

// hostPort is the host:port a control plane's address dials, whether it was
// written as a URL or not.
func hostPort(s string) string {
	s = strings.TrimSpace(s)
	port := ""
	switch {
	case strings.HasPrefix(s, "https://"):
		s, port = strings.TrimPrefix(s, "https://"), "443"
	case strings.HasPrefix(s, "http://"):
		s, port = strings.TrimPrefix(s, "http://"), "80"
	}
	s, _, _ = strings.Cut(s, "/")
	if _, _, err := net.SplitHostPort(s); err != nil && port != "" && s != "" {
		s = net.JoinHostPort(strings.Trim(s, "[]"), port)
	}
	return s
}

func (p nodeProbe) checkPlanes(ctx context.Context) []planeCheck {
	planes := readControlPlanes(p.planesFile, p.configured)
	var wg sync.WaitGroup
	for i := range planes {
		wg.Add(1)
		go func(c *planeCheck) {
			defer wg.Done()
			if err := p.dial(ctx, c.addr); err != nil {
				c.err = err.Error()
			}
		}(&planes[i])
	}
	wg.Wait()
	return planes
}

func dialTCP(ctx context.Context, addr string) error {
	d := net.Dialer{Timeout: nodeCmdTimeout}
	c, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return err
	}
	return c.Close()
}

// --- the GPUs ---

// readGPUs samples the cards, naming them from discovery the first time.
func (p nodeProbe) readGPUs(ctx context.Context) ([]nodeGPU, map[int]string, string) {
	ctx, cancel := context.WithTimeout(ctx, nodeCmdTimeout)
	defer cancel()
	sensors := p.gpus(ctx)
	if len(sensors) == 0 {
		// Without nvidia-smi discovery says nothing and no error: a machine
		// without a card. With it and no answer, a driver in trouble.
		if _, err := p.cards(ctx); err != nil {
			return nil, p.names, "The GPUs do not answer: " + err.Error()
		}
		return nil, p.names, "No GPU on this machine."
	}
	names := p.names
	for _, s := range sensors {
		if names[s.Index] == "" {
			names = map[int]string{}
			if cards, err := p.cards(ctx); err == nil {
				for _, c := range cards {
					if c.MIG == nil {
						names[c.Index] = c.Name
					}
				}
			}
			break
		}
	}
	var out []nodeGPU
	for _, s := range sensors {
		g := nodeGPU{index: s.Index, name: names[s.Index], util: s.UtilPercent, tempC: s.TempC,
			memUsedMB: s.MemoryUsedMB, memTotalMB: s.MemoryTotalMB}
		if s.MemoryTotalMB == 0 {
			g.util = -1
		}
		out = append(out, g)
	}
	return out, names, ""
}

// cardOf is the card a device number is on: itself, or a MIG instance's card.
func cardOf(i int) int {
	if discovery.IsMIGIndex(i) {
		return (i - discovery.MIGIndex(0, 0)) / 100
	}
	return i
}

// modelsOnGPUs names, on each card, the models whose engines hold it.
func modelsOnGPUs(gpus []nodeGPU, runs []engineRun) []nodeGPU {
	for i := range gpus {
		gpus[i].models = nil
		for _, r := range runs {
			for _, d := range r.gpus {
				if cardOf(d) == gpus[i].index && !slices.Contains(gpus[i].models, r.model) {
					gpus[i].models = append(gpus[i].models, r.model)
				}
			}
		}
	}
	return gpus
}

// --- the engines ---

// shardStateFile is the part of the supervisor's shard-*.json the overview
// reads (shardState in the supervisor package).
type shardStateFile struct {
	ShardID   string            `json:"shard_id"`
	ReplicaID string            `json:"replica_id"`
	PID       int               `json:"pid"`
	StartedAt time.Time         `json:"started_at"`
	Spec      engines.ShardSpec `json:"spec"`
}

// readShardStates is the engines the supervisor keeps state files for in
// its work directory, with whether each process is still there.
func readShardStates(dir string, alive func(int) bool) ([]engineRun, string) {
	entries, err := os.ReadDir(dir)
	switch {
	case errors.Is(err, os.ErrPermission):
		return nil, "The engines' state is root's: run with sudo."
	case errors.Is(err, os.ErrNotExist):
		return nil, ""
	case err != nil:
		return nil, err.Error()
	}
	var runs []engineRun
	note := ""
	for _, e := range entries {
		name := e.Name()
		if !strings.HasPrefix(name, "shard-") || !strings.HasSuffix(name, ".json") {
			continue
		}
		b, err := os.ReadFile(filepath.Join(dir, name))
		if errors.Is(err, os.ErrPermission) {
			note = "The engines' state is root's: run with sudo."
			continue
		}
		var st shardStateFile
		if err != nil || json.Unmarshal(b, &st) != nil {
			continue
		}
		replica := st.ReplicaID
		if replica == "" {
			replica = st.Spec.ReplicaID
		}
		model := st.Spec.Profile.ServedModelName
		if model == "" {
			model = modelOfReplica(replica)
		}
		runs = append(runs, engineRun{replica: replica, model: model, engine: st.Spec.Profile.Engine,
			gpus: st.Spec.GPUIndices, port: st.Spec.HTTPPort, pid: st.PID, alive: st.PID > 0 && alive(st.PID),
			started: st.StartedAt})
	}
	sort.Slice(runs, func(i, j int) bool { return runs[i].replica < runs[j].replica })
	return runs, note
}

// modelOfReplica is the model a replica id names: the control plane calls a
// replica the model's name and a suffix.
func modelOfReplica(id string) string {
	if i := strings.LastIndexByte(id, '-'); i > 0 {
		return id[:i]
	}
	return id
}

func pidAlive(pid int) bool {
	err := syscall.Kill(pid, 0)
	return err == nil || errors.Is(err, syscall.EPERM)
}

// --- the disks ---

func (p nodeProbe) readDisks() []diskUse {
	disks := []diskUse{{what: "weights", path: p.weightsDir}, {what: "engines", path: p.enginesDir}}
	for i := range disks {
		// In its own goroutine: statfs on a network mount that has gone away
		// does not return.
		type res struct {
			free, total uint64
			err         error
		}
		ch := make(chan res, 1)
		go func(path string) {
			f, t, err := p.statfs(path)
			ch <- res{f, t, err}
		}(disks[i].path)
		select {
		case r := <-ch:
			disks[i].free, disks[i].total = r.free, r.total
			if r.err != nil {
				disks[i].err = r.err.Error()
				if errors.Is(r.err, os.ErrNotExist) {
					disks[i].err = "not there yet"
				}
			}
		case <-time.After(nodeCmdTimeout):
			disks[i].err = "does not answer"
		}
	}
	return disks
}

func statfs(path string) (free, total uint64, err error) {
	var st syscall.Statfs_t
	if err := syscall.Statfs(path, &st); err != nil {
		return 0, 0, &os.PathError{Op: "statfs", Path: path, Err: err}
	}
	//nolint:unconvert // the field types differ between systems
	return uint64(st.Bavail) * uint64(st.Bsize), uint64(st.Blocks) * uint64(st.Bsize), nil
}

// --- the problems ---

func (p nodeProbe) readProblems(ctx context.Context) ([]string, string) {
	out, err := p.run(ctx, "journalctl", "-u", "llmgw-agent", "-p", "warning", "--since", "-1h",
		"--no-pager", "-o", "cat", "-n", strconv.Itoa(nodeProblemLines))
	if err != nil && len(out) == 0 {
		return nil, "The journal cannot be read: " + err.Error()
	}
	lines := parseJournal(string(out), nodeProblemLines)
	if len(lines) == 0 && !p.root {
		return nil, "None that this user can see: run with sudo to read the agent's journal."
	}
	return lines, ""
}

// parseJournal keeps the last n lines journalctl printed, without its own
// remarks.
func parseJournal(out string, n int) []string {
	var lines []string
	sc := bufio.NewScanner(strings.NewReader(out))
	sc.Buffer(make([]byte, 64*1024), 1024*1024)
	for sc.Scan() {
		l := strings.TrimSpace(sc.Text())
		if l == "" || strings.HasPrefix(l, "-- ") || strings.HasPrefix(l, "Hint: ") {
			continue
		}
		lines = append(lines, l)
	}
	if len(lines) > n {
		lines = lines[len(lines)-n:]
	}
	return lines
}

// runLocal runs a tool on this machine, bounded by nodeCmdTimeout. A tool
// that is not installed is said so, not waited for.
func runLocal(ctx context.Context, name string, args ...string) ([]byte, error) {
	if _, err := exec.LookPath(name); err != nil {
		return nil, fmt.Errorf("%s is not installed", name)
	}
	ctx, cancel := context.WithTimeout(ctx, nodeCmdTimeout)
	defer cancel()
	//nolint:gosec // fixed tools and arguments, from this file
	return exec.CommandContext(ctx, name, args...).Output()
}

// --- the panel ---

func (m consoleModel) nodeOverviewPanel(pw int) (string, string, string) {
	keys := "↑↓ choose · PgUp/PgDn scroll · Q leave · refreshed every three seconds"
	if m.nodeOv.at.IsZero() {
		return "Overview", "\n  " + m.spin.View() + " reading this machine…", keys
	}
	return "Overview", renderNodeOverview(m.nodeOv, pw, time.Now()), keys
}

// renderNodeOverview draws the overview from what was read.
func renderNodeOverview(ov nodeOverview, pw int, now time.Time) string {
	var b strings.Builder
	head := func(s string) { b.WriteString("\n  " + lipgloss.NewStyle().Bold(true).Render(s) + "\n") }
	kv := func(k, v string) { fmt.Fprintf(&b, "    %-14s %s\n", sDim.Render(k), v) }
	line := func(s string) { b.WriteString("    " + s + "\n") }

	head("Agent")
	a := ov.agent
	switch {
	case a.note != "":
		kv("llmgw-agent", sWarn.Render(a.note))
	default:
		state := a.state
		if a.sub != "" {
			state += " (" + a.sub + ")"
		}
		if a.state == "active" {
			state = sOK.Render(state)
		} else {
			state = sBad.Render(state)
		}
		if !a.since.IsZero() {
			state += sDim.Render(" since " + a.since.Format("2006-01-02 15:04") + ", " + uptime(now.Sub(a.since)) + " ago")
		}
		kv("llmgw-agent", state)
	}
	kv("llmgw", ov.version)

	head("Control planes")
	if len(ov.planes) == 0 {
		line(sWarn.Render("None configured: agent.yaml names no server."))
	}
	for _, c := range ov.planes {
		state := sOK.Render("reachable")
		if c.err != "" {
			state = sBad.Render("✗ ") + truncate(c.err, max(20, pw-30))
		}
		if c.learned {
			state += sDim.Render("  (named by the leader)")
		}
		fmt.Fprintf(&b, "    %-22s %s\n", c.addr, state)
	}

	head("GPUs")
	if len(ov.gpus) == 0 {
		line(sDim.Render(ov.gpuNote))
	} else {
		var rows [][]string
		for _, g := range ov.gpus {
			busy, mem, temp := "—", "—", "—"
			if g.util >= 0 {
				busy = fmt.Sprintf("%.0f%%", g.util)
			}
			if g.memTotalMB > 0 {
				mem = fmt.Sprintf("%.1f/%.0f GB", g.memUsedMB/1024, g.memTotalMB/1024)
			}
			if g.tempC > 0 {
				temp = fmt.Sprintf("%.0f°C", g.tempC)
			}
			models := strings.Join(g.models, ", ")
			if models == "" {
				models = sDim.Render("free")
			}
			rows = append(rows, []string{strconv.Itoa(g.index), g.name, busy, mem, temp, models})
		}
		b.WriteString(indent(table([]string{"GPU", "CARD", "BUSY", "MEMORY", "TEMP", "MODELS"}, rows, -1, pw-2, false), "  ") + "\n")
	}

	head("Engines running here")
	switch {
	case len(ov.engines) > 0:
		var rows [][]string
		for _, r := range ov.engines {
			// The process by its pid, marked when it is no longer there.
			proc := sOK.Render(strconv.Itoa(r.pid))
			if !r.alive {
				proc = sBad.Render(strconv.Itoa(r.pid) + " gone")
			}
			gpus := strings.Join(discovery.DeviceLabels(r.gpus), ",")
			if gpus == "" {
				gpus = sDim.Render("CPU")
			}
			up := "—"
			if !r.started.IsZero() {
				up = uptime(now.Sub(r.started))
			}
			port := "—"
			if r.port > 0 {
				port = strconv.Itoa(r.port)
			}
			rows = append(rows, []string{r.replica, r.model, r.engine, gpus, port, proc, up})
		}
		b.WriteString(indent(table([]string{"REPLICA", "MODEL", "ENGINE", "GPUS", "PORT", "PID", "UP"}, rows, -1, pw-2, false), "  ") + "\n")
		if ov.enginesNote != "" {
			line(sWarn.Render(ov.enginesNote))
		}
	case ov.enginesNote != "":
		line(sWarn.Render(ov.enginesNote))
	default:
		line(sDim.Render("None: the control plane has placed nothing here."))
	}

	head("Disk")
	for _, d := range ov.disks {
		use := sWarn.Render(d.err)
		if d.err == "" {
			use = fmt.Sprintf("%s free of %s", diskSize(d.free), diskSize(d.total))
			if d.total > 0 && d.free*10 < d.total {
				use = sWarn.Render(use)
			}
		}
		fmt.Fprintf(&b, "    %-14s %-28s %s\n", sDim.Render(d.what), d.path, use)
	}

	head("Recent problems")
	switch {
	case ov.problemsNote != "":
		line(sDim.Render(ov.problemsNote))
	case len(ov.problems) == 0:
		line(sOK.Render("None in the last hour."))
	default:
		for _, l := range ov.problems {
			l = truncateWidth(l, max(20, pw-4))
			if strings.Contains(l, "level=ERROR") {
				l = sBad.Render(l)
			} else {
				l = sWarn.Render(l)
			}
			line(l)
		}
	}
	b.WriteString("\n  " + sDim.Render("read at "+ov.at.Format("15:04:05")) + "\n")
	return b.String()
}

// uptime is a duration as people say it: 3d4h, 2h15m, 40s.
func uptime(d time.Duration) string {
	d = d.Round(time.Second)
	switch {
	case d < 0:
		return "0s"
	case d >= 24*time.Hour:
		return fmt.Sprintf("%dd%dh", d/(24*time.Hour), (d%(24*time.Hour))/time.Hour)
	case d >= time.Hour:
		return fmt.Sprintf("%dh%dm", d/time.Hour, (d%time.Hour)/time.Minute)
	case d >= time.Minute:
		return fmt.Sprintf("%dm%ds", d/time.Minute, (d%time.Minute)/time.Second)
	}
	return d.String()
}

// diskSize is a byte count in the unit disks are sold in.
func diskSize(n uint64) string {
	switch {
	case n >= 1<<40:
		return fmt.Sprintf("%.1f TB", float64(n)/(1<<40))
	case n >= 1<<30:
		return fmt.Sprintf("%.0f GB", float64(n)/(1<<30))
	}
	return fmt.Sprintf("%.0f MB", float64(n)/(1<<20))
}