package cli

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/doctor"
)

func newDoctorCmd() *cobra.Command {
	var (
		opts    doctor.Options
		asJSON  bool
		verbose bool
		fabrics []string
	)
	cmd := &cobra.Command{
		Use:   "doctor",
		Short: "Check this node's GPUs, NVLink, InfiniBand and engines for what llmgw needs",
		Long: `Reads this node's GPUs, their NVLink and PCIe topology, its InfiniBand and
RoCE ports, the kernel modules RDMA needs, the agent's configuration and
whether vLLM can be installed here, and says what is wrong and how to fix
it — most serious first.

Run it on a node before it joins, and whenever a multi-node replica will not
come up. --json writes everything it read as well as what it concluded; send
that file when the report is not enough to go on.`,
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			doctor.SetConfigLoader(config.LoadAgent)
			for _, spec := range fabrics {
				name, hcas, _ := strings.Cut(spec, ":")
				if name == "" || hcas == "" {
					return fmt.Errorf("--fabric %q: give it as name:mlx5_0,mlx5_1", spec)
				}
				if opts.Fabrics == nil {
					opts.Fabrics = map[string][]string{}
				}
				opts.Fabrics[name] = strings.Split(hcas, ",")
			}
			switch opts.VLLMVersion {
			case "":
				opts.VLLMVersion = config.DefaultEngines["vllm"]
			case "none":
				opts.VLLMVersion = ""
			}
			ctx, cancel := context.WithTimeout(cmd.Context(), 60*time.Second)
			defer cancel()
			facts := doctor.Collect(ctx, opts)
			findings := doctor.Evaluate(facts)
			if asJSON {
				enc := json.NewEncoder(os.Stdout)
				enc.SetIndent("", "  ")
				return enc.Encode(struct {
					Findings  []doctor.Finding `json:"findings"`
					Readiness doctor.Readiness `json:"readiness"`
					Facts     doctor.Facts     `json:"facts"`
				}{findings, doctor.ReadinessOf(findings), facts})
			}
			printDoctor(os.Stdout, facts, findings, verbose)
			for _, f := range findings {
				if f.Level == doctor.Fail {
					return fmt.Errorf("%d problem(s) need fixing", countLevel(findings, doctor.Fail))
				}
			}
			return nil
		},
	}
	f := cmd.Flags()
	f.StringVar(&opts.ConfigPath, "config", "", "agent configuration (default /etc/llmgw/agent.yaml when present)")
	f.StringVar(&opts.IBRoot, "ib-root", discovery.SysfsRoot, "sysfs InfiniBand directory")
	f.StringVar(&opts.NvidiaSMI, "nvidia-smi", "", "path to nvidia-smi")
	f.StringVar(&opts.VLLMVersion, "vllm", "", "the vLLM version to check for (default: the one this llmgw installs; none to skip)")
	f.StringArrayVar(&fabrics, "fabric", nil, "a fabric about to be declared, name:mlx5_0,mlx5_1, checked as if it were; repeatable")
	f.BoolVar(&asJSON, "json", false, "write every fact and finding as JSON, to send back")
	f.BoolVar(&verbose, "verbose", false, "also print the raw nvidia-smi topology")
	return cmd
}

func countLevel(fs []doctor.Finding, l doctor.Level) int {
	n := 0
	for _, f := range fs {
		if f.Level == l {
			n++
		}
	}
	return n
}

func printDoctor(out io.Writer, f doctor.Facts, findings []doctor.Finding, verbose bool) {
	fmt.Fprintf(out, "%s · %s %s %s · driver %s\n\n", f.Hostname, f.Host.OSID, f.Host.OSVersion,
		f.Host.Arch, orDash(f.Host.NVIDIADriver))

	// The node first; then what replicas spanning nodes would need, which
	// keeps a node out only when a fabric was declared for it.
	multi := false
	for _, pass := range []bool{false, true} {
		for _, x := range findings {
			if x.MultiNode() != pass {
				continue
			}
			if pass && !multi {
				fmt.Fprintln(out, "\nMulti-node (fabrics, RDMA) — for replicas spanning nodes only:")
				multi = true
			}
			fmt.Fprintf(out, "%-4s  %-6s  %s\n", x.Level, x.Area, x.Message)
			if x.Fix != "" && (x.Level == doctor.Fail || x.Level == doctor.Warn) {
				fmt.Fprintf(out, "              → %s\n", x.Fix)
			}
		}
	}
	r := doctor.ReadinessOf(findings)
	fmt.Fprintf(out, "\nThis node:  %s\nMulti-node: %s\n", r.Node, r.MultiNode)

	w := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
	if len(f.GPUs) > 0 {
		fmt.Fprintln(w, "\nGPU\tNAME\tCLASS\tNUMA\tNVLINK PEERS\tNEAREST ADAPTER")
		for _, g := range f.GPUs {
			// A MIG instance as 0/mig2, under its card.
			fmt.Fprintf(w, "%s\t%s\t%s\t%d\t%s\t%s\n", discovery.DeviceLabel(g.Index), g.Name, g.Class(), g.NUMANode,
				orDash(ints(g.NVLinkPeers)), orDash(nearest(g.NICPaths)))
		}
	}
	if len(f.HCAs) > 0 {
		fmt.Fprintln(w, "\nADAPTER\tSTATE\tPHYS\tRATE\tLINK\tNUMA\tROCE GID\tERRORS")
		for _, h := range f.HCAs {
			for _, p := range h.Ports {
				gid := "-"
				if p.RoCEv2GID >= 0 {
					gid = fmt.Sprint(p.RoCEv2GID)
				}
				fmt.Fprintf(w, "%s/%d\t%s\t%s\t%d Gb/s\t%s\t%d\t%s\t%s\n", h.Name, p.Num, p.State,
					orDash(p.PhysState), p.RateGbps, p.LinkLayer, h.NUMANode, gid, orDash(errorsOf(p.Counters)))
			}
		}
	}
	_ = w.Flush()

	for _, fabric := range sortedKeys(f.NCCL) {
		fmt.Fprintf(out, "\nA shard on %s is given:\n", fabric)
		env := f.NCCL[fabric]
		for _, k := range sortedKeys(env) {
			fmt.Fprintf(out, "  %s=%s\n", k, env[k])
		}
	}
	if len(f.Engines) > 0 {
		fmt.Fprintln(out, "\nEngines:")
		for _, k := range sortedKeys(f.Engines) {
			fmt.Fprintf(out, "  %s: %s\n", k, f.Engines[k])
		}
	}
	if verbose && f.TopoRaw != "" {
		fmt.Fprintf(out, "\nnvidia-smi topo -m:\n%s\n", f.TopoRaw)
	}
	fmt.Fprintf(out, "\n%d failed, %d warnings. `llmgw doctor --json > report.json` keeps everything read, to send back.\n",
		countLevel(findings, doctor.Fail), countLevel(findings, doctor.Warn))
}

func nearest(paths map[string]string) string {
	best, rel := "", ""
	for h, r := range paths {
		if best == "" || discovery.PathRank(r) < discovery.PathRank(rel) ||
			(discovery.PathRank(r) == discovery.PathRank(rel) && h < best) {
			best, rel = h, r
		}
	}
	if best == "" {
		return ""
	}
	return best + " (" + rel + ")"
}

func errorsOf(c map[string]uint64) string {
	var out []string
	for _, k := range sortedKeys(c) {
		if c[k] > 0 && k != "port_xmit_data" && k != "port_rcv_data" {
			out = append(out, fmt.Sprintf("%s=%d", k, c[k]))
		}
	}
	return strings.Join(out, " ")
}

func ints(v []int) string {
	s := make([]string, len(v))
	for i, x := range v {
		s[i] = fmt.Sprint(x)
	}
	return strings.Join(s, ",")
}

func orDash(s string) string {
	if s == "" {
		return "-"
	}
	return s
}

func sortedKeys[V any](m map[string]V) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}