package main

import (
	"encoding/xml"
	"fmt"
	"sort"
	"strings"
)

// ---------------------------------------------------------------------------
// The ProxyJump topology
// ---------------------------------------------------------------------------

// GraphNode is one host in the jump topology.
type GraphNode struct {
	Name     string `json:"name"`
	Declared bool   `json:"declared"`  // has a Host block of its own
	HostName string `json:"hostname"`  // resolved HostName, when known
	Level    int    `json:"level"`     // 0 = reached directly, N = N jumps in
	InCycle  bool   `json:"in_cycle"`  // part of a ProxyJump loop
	IsTarget bool   `json:"is_target"` // something jumps through it
}

// GraphEdge is "From is reached through To".
type GraphEdge struct {
	From string `json:"from"` // the host being connected to
	Via  string `json:"via"`  // the jump host used to reach it
}

// Graph is the whole topology.
type Graph struct {
	Nodes []GraphNode `json:"nodes"`
	Edges []GraphEdge `json:"edges"`
}

// BuildGraph walks every declared host's ProxyJump chain.
func BuildGraph(cfg *Config, q Query) *Graph {
	nodes := map[string]*GraphNode{}
	edgeSet := map[GraphEdge]bool{}

	touch := func(name string) *GraphNode {
		key := strings.ToLower(name)
		if n, ok := nodes[key]; ok {
			return n
		}
		declared, _ := hostIsDeclared(cfg, name)
		nq := q
		nq.Host = name
		res := Resolve(cfg, nq)
		hn := res.Get("hostname")
		if hn != "" {
			hn = expandTokens(hn, strings.ToLower(name), nq)
		}
		n := &GraphNode{Name: name, Declared: declared, HostName: hn}
		nodes[key] = n
		return n
	}

	var walk func(host string, path []string, onPath map[string]bool)
	walk = func(host string, path []string, onPath map[string]bool) {
		if len(path) > 32 {
			return
		}
		touch(host)
		hops := jumpChainFor(cfg, q, host)
		if len(hops) == 0 {
			return
		}
		// ProxyJump a,b means: reach a first, then b, then the host. So the
		// host is reached via the LAST hop and each hop via the one before it.
		prev := host
		for i := len(hops) - 1; i >= 0; i-- {
			h := hops[i]
			touch(h.Host).IsTarget = true
			e := GraphEdge{From: prev, Via: h.Host}
			if !edgeSet[e] {
				edgeSet[e] = true
			}
			prev = h.Host
		}
		next := hops[0].Host
		if onPath[strings.ToLower(next)] {
			// Only the ring is in the cycle. Hosts that merely lead into it
			// are reachable, they just cannot get past the loop.
			ring := append(append([]string{}, path...), next)
			for i, p := range ring {
				if strings.EqualFold(p, next) {
					ring = ring[i:]
					break
				}
			}
			for _, name := range ring {
				if n, ok := nodes[strings.ToLower(name)]; ok {
					n.InCycle = true
				}
			}
			return
		}
		onPath[strings.ToLower(next)] = true
		walk(next, append(path, next), onPath)
		delete(onPath, strings.ToLower(next))
	}

	for _, h := range concreteHosts(cfg) {
		walk(h, []string{h}, map[string]bool{strings.ToLower(h): true})
	}

	g := &Graph{}
	for _, e := range sortedEdges(edgeSet) {
		g.Edges = append(g.Edges, e)
	}
	// Level = how many jumps deep. A node with no outgoing "via" edge is at
	// level 0 (reached directly); everything else is one deeper than its jump.
	via := map[string]string{}
	for _, e := range g.Edges {
		via[strings.ToLower(e.From)] = strings.ToLower(e.Via)
	}
	var level func(name string, seen map[string]bool) int
	level = func(name string, seen map[string]bool) int {
		if seen[name] {
			return 0 // a cycle has no honest depth
		}
		v, ok := via[name]
		if !ok {
			return 0
		}
		seen[name] = true
		return 1 + level(v, seen)
	}
	var names []string
	for k := range nodes {
		names = append(names, k)
	}
	sort.Strings(names)
	used := map[int]bool{}
	for _, k := range names {
		nodes[k].Level = level(k, map[string]bool{})
		used[nodes[k].Level] = true
	}
	// Squeeze out levels nothing landed on so the drawing has no empty columns.
	var lv []int
	for l := range used {
		lv = append(lv, l)
	}
	sort.Ints(lv)
	remap := map[int]int{}
	for i, l := range lv {
		remap[l] = i
	}
	for _, k := range names {
		nodes[k].Level = remap[nodes[k].Level]
		g.Nodes = append(g.Nodes, *nodes[k])
	}
	sort.SliceStable(g.Nodes, func(i, j int) bool {
		if g.Nodes[i].Level != g.Nodes[j].Level {
			return g.Nodes[i].Level < g.Nodes[j].Level
		}
		return g.Nodes[i].Name < g.Nodes[j].Name
	})
	return g
}

func sortedEdges(set map[GraphEdge]bool) []GraphEdge {
	out := make([]GraphEdge, 0, len(set))
	for e := range set {
		out = append(out, e)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].From != out[j].From {
			return out[i].From < out[j].From
		}
		return out[i].Via < out[j].Via
	})
	return out
}

// ---------------------------------------------------------------------------
// SVG rendering
// ---------------------------------------------------------------------------

const (
	svgNodeH   = 34
	svgRowGap  = 22
	svgColGap  = 90
	svgMargin  = 24
	svgCharW   = 7.2
	svgPadding = 14
)

// RenderSVG draws the topology. Jump hosts sit on the left, the hosts reached
// through them on the right, with an arrow pointing the way a connection
// travels.
func RenderSVG(g *Graph, title string) string {
	// Columns are jump depth: level 0 (reached directly, i.e. the jump hosts
	// themselves) on the left, each extra jump one column to the right.
	maxLevel := 0
	for _, n := range g.Nodes {
		if n.Level > maxLevel {
			maxLevel = n.Level
		}
	}
	colWidth := make([]float64, maxLevel+1)
	byCol := make([][]GraphNode, maxLevel+1)
	for _, n := range g.Nodes {
		byCol[n.Level] = append(byCol[n.Level], n)
		w := float64(len(nodeLabel(n)))*svgCharW + 2*svgPadding
		if w > colWidth[n.Level] {
			colWidth[n.Level] = w
		}
	}
	colX := make([]float64, maxLevel+1)
	x := float64(svgMargin)
	for lv := 0; lv <= maxLevel; lv++ {
		colX[lv] = x
		x += colWidth[lv] + svgColGap
	}
	width := x - svgColGap + svgMargin
	if width < 360 {
		width = 360
	}

	pos := map[string][4]float64{} // x, y, w, h
	rows := 0
	for lv := maxLevel; lv >= 0; lv-- {
		y := float64(svgMargin + 46)
		for _, n := range byCol[lv] {
			pos[strings.ToLower(n.Name)] = [4]float64{colX[lv], y, colWidth[lv], svgNodeH}
			y += svgNodeH + svgRowGap
		}
		if len(byCol[lv]) > rows {
			rows = len(byCol[lv])
		}
	}
	height := float64(svgMargin+46) + float64(rows)*(svgNodeH+svgRowGap) + svgMargin

	var b strings.Builder
	fmt.Fprintf(&b, `<?xml version="1.0" encoding="UTF-8"?>`+"\n")
	fmt.Fprintf(&b, `<svg xmlns="http://www.w3.org/2000/svg" width="%.0f" height="%.0f" `+
		`viewBox="0 0 %.0f %.0f" font-family="monospace" font-size="12">`+"\n",
		width, height, width, height)
	b.WriteString("<defs>\n")
	b.WriteString(`<marker id="arrow" viewBox="0 0 10 10" refX="9" refY="5" ` +
		`markerWidth="7" markerHeight="7" orient="auto-start-reverse">` +
		`<path d="M 0 0 L 10 5 L 0 10 z" fill="#44506a"/></marker>` + "\n")
	b.WriteString("</defs>\n")
	fmt.Fprintf(&b, `<rect x="0" y="0" width="%.0f" height="%.0f" fill="#ffffff"/>`+"\n", width, height)
	fmt.Fprintf(&b, `<text x="%d" y="%d" font-size="15" font-weight="bold" fill="#1d2333">%s</text>`+"\n",
		svgMargin, svgMargin+14, escapeXML(title))
	fmt.Fprintf(&b, `<text x="%d" y="%d" fill="#5a6478">arrow points the way the connection travels: `+
		`jump host -&gt; destination</text>`+"\n", svgMargin, svgMargin+32)

	// Edges first so the boxes sit on top of them.
	for _, e := range g.Edges {
		from, okF := pos[strings.ToLower(e.From)]
		via, okV := pos[strings.ToLower(e.Via)]
		if !okF || !okV {
			continue
		}
		x1 := via[0] + via[2]
		y1 := via[1] + via[3]/2
		x2 := from[0]
		y2 := from[1] + from[3]/2
		mid := (x1 + x2) / 2
		fmt.Fprintf(&b, `<path d="M %.1f %.1f C %.1f %.1f, %.1f %.1f, %.1f %.1f" `+
			`fill="none" stroke="#44506a" stroke-width="1.4" marker-end="url(#arrow)"/>`+"\n",
			x1, y1, mid, y1, mid, y2, x2, y2)
	}

	for _, n := range g.Nodes {
		p, ok := pos[strings.ToLower(n.Name)]
		if !ok {
			continue
		}
		fill, stroke := "#eef2fb", "#44506a"
		switch {
		case n.InCycle:
			fill, stroke = "#fdecec", "#b3261e"
		case !n.Declared:
			fill, stroke = "#fff7e6", "#8a6100"
		}
		fmt.Fprintf(&b, `<rect x="%.1f" y="%.1f" width="%.1f" height="%.1f" rx="6" `+
			`fill="%s" stroke="%s" stroke-width="1.4"/>`+"\n",
			p[0], p[1], p[2], p[3], fill, stroke)
		fmt.Fprintf(&b, `<text x="%.1f" y="%.1f" fill="#1d2333">%s</text>`+"\n",
			p[0]+svgPadding, p[1]+p[3]/2+4, escapeXML(nodeLabel(n)))
	}

	b.WriteString(`<text x="24" y="` + fmt.Sprintf("%.0f", height-8) + `" fill="#5a6478">` +
		escapeXML(fmt.Sprintf("%d hosts, %d jumps", len(g.Nodes), len(g.Edges))) +
		`</text>` + "\n")
	b.WriteString("</svg>\n")
	return b.String()
}

func nodeLabel(n GraphNode) string {
	label := n.Name
	if n.HostName != "" && !strings.EqualFold(n.HostName, n.Name) {
		label += " (" + n.HostName + ")"
	}
	if !n.Declared {
		label += " [no Host block]"
	}
	if n.InCycle {
		label += " [cycle]"
	}
	return label
}

func escapeXML(s string) string {
	var b strings.Builder
	if err := xml.EscapeText(&b, []byte(s)); err != nil {
		return ""
	}
	return b.String()
}