// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package diagram

import (
	"fmt"
	"slices"
	"sort"
	"strings"
)

// Layout constants, in SVG units.
const (
	margin       = 32
	rankGap      = 64
	laneGap      = 48
	nodeMinWidth = 96
	textPad      = 28
)

type laidNode struct {
	x, y, w, h int // centre and size in the TB orientation
}

type flowLayout struct {
	laid   []laidNode // parallel to d.nodes
	rank   []int      // parallel to d.nodes
	width  int
	height int
}

func renderFlowchart(lines []string, dir string) ([]byte, error) {
	d, err := parseFlowchart(lines, dir)
	if err != nil {
		return nil, err
	}
	lay := layoutFlowchart(d)
	return writeFlowchart(d, lay), nil
}

// layoutFlowchart assigns ranks by longest path, orders each rank by
// barycentre sweeps with the order of appearance as the tie breaker, and
// packs the rows. Every step walks nodes and edges by index, so the
// layout never depends on map iteration.
func layoutFlowchart(d *flowDiagram) *flowLayout {
	n := len(d.nodes)
	lay := &flowLayout{laid: make([]laidNode, n), rank: make([]int, n)}

	for i, node := range d.nodes {
		lay.laid[i].w, lay.laid[i].h = nodeSize(node)
	}

	// Ranks: relax along the edges, in index order, one pass per node at
	// most. A cycle cannot push the ranks past that bound.
	for pass := 0; pass <= len(d.nodes); pass++ {
		changed := false
		for _, e := range d.edges {
			if lay.rank[e.to] < lay.rank[e.from]+1 {
				lay.rank[e.to] = lay.rank[e.from] + 1
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	maxRank := 0
	for _, r := range lay.rank {
		if r > maxRank {
			maxRank = r
		}
	}
	rows := make([][]int, maxRank+1)
	for i, r := range lay.rank {
		rows[r] = append(rows[r], i)
	}

	// Four alternating barycentre sweeps.
	for pass := range 4 {
		down := pass%2 == 0
		ranks := make([]int, 0, len(rows))
		if down {
			for r := range rows {
				ranks = append(ranks, r)
			}
		} else {
			for r := range slices.Backward(rows) {
				ranks = append(ranks, r)
			}
		}
		for _, r := range ranks {
			if down && r == 0 || !down && r == len(rows)-1 {
				continue
			}
			neighbour := r - 1
			if !down {
				neighbour = r + 1
			}
			position := map[int]int{}
			for p, i := range rows[neighbour] {
				position[i] = p
			}
			bary := map[int]int{}
			sum := map[int]int{}
			for _, e := range d.edges {
				from, to := e.from, e.to
				if !down {
					from, to = to, from
				}
				if lay.rank[from] == neighbour && lay.rank[to] == r {
					if p, ok := position[from]; ok {
						sum[to] += p
						bary[to]++
					}
				}
			}
			sort.SliceStable(rows[r], func(a, b int) bool {
				ia, ib := rows[r][a], rows[r][b]
				switch {
				case bary[ia] == 0 && bary[ib] == 0:
					return ia < ib
				case bary[ia] == 0:
					return false
				case bary[ib] == 0:
					return true
				}
				ma := sum[ia] / bary[ia]
				mb := sum[ib] / bary[ib]
				if ma != mb {
					return ma < mb
				}
				return ia < ib
			})
		}
	}

	// Rank heights, row widths and positions.
	rowHeight := make([]int, len(rows))
	rowWidth := make([]int, len(rows))
	for r, row := range rows {
		for _, i := range row {
			if lay.laid[i].h > rowHeight[r] {
				rowHeight[r] = lay.laid[i].h
			}
			rowWidth[r] += lay.laid[i].w + laneGap
		}
		rowWidth[r] -= laneGap
	}
	maxWidth := 0
	y := margin
	rowY := make([]int, len(rows))
	for r := range rows {
		rowY[r] = y + rowHeight[r]/2
		y += rowHeight[r] + rankGap
		if rowWidth[r] > maxWidth {
			maxWidth = rowWidth[r]
		}
	}
	for r, row := range rows {
		x := margin + (maxWidth-rowWidth[r])/2 + laneGap/2
		for _, i := range row {
			x += lay.laid[i].w / 2
			lay.laid[i].x = x
			lay.laid[i].y = rowY[r]
			x += lay.laid[i].w/2 + laneGap/2
		}
	}
	lay.width = maxWidth + 2*margin
	lay.height = y - rankGap + margin

	// A subgraph with its own direction lays its members out inside the
	// box they occupy, in that direction.
	for _, sg := range d.subgraphs {
		if sg.dir == "" || sg.dir == d.dir || len(sg.nodes) < 2 {
			continue
		}
		reLayoutSubgraph(d, lay, sg)
	}
	return lay
}

func reLayoutSubgraph(d *flowDiagram, lay *flowLayout, sg *flowSubgraph) {
	minX, minY, maxX, maxY := boundingBox(lay, sg.nodes)
	horizontal := sg.dir == "LR" || sg.dir == "RL"
	members := map[int]bool{}
	for _, i := range sg.nodes {
		members[i] = true
	}
	rank := map[int]int{}
	for pass := 0; pass <= len(sg.nodes); pass++ {
		changed := false
		for _, e := range d.edges {
			if members[e.from] && members[e.to] && rank[e.to] < rank[e.from]+1 {
				rank[e.to] = rank[e.from] + 1
				changed = true
			}
		}
		if !changed {
			break
		}
	}
	maxRank := 0
	for _, i := range sg.nodes {
		if rank[i] > maxRank {
			maxRank = rank[i]
		}
	}
	rows := make([][]int, maxRank+1)
	for _, i := range sg.nodes {
		rows[rank[i]] = append(rows[rank[i]], i)
	}
	centreS := (minX + maxX) / 2
	centreP := (minY + maxY) / 2
	if horizontal {
		centreS, centreP = centreP, centreS
	}
	spanP := (maxY - minY) - laneGap
	if horizontal {
		spanP = (maxX - minX) - laneGap
	}
	for r, row := range rows {
		p := centreP - spanP/2 + (spanP*(2*r+1))/(2*(maxRank+1))
		var sSize int
		for _, i := range row {
			if horizontal {
				sSize += lay.laid[i].h
			} else {
				sSize += lay.laid[i].w
			}
		}
		sSize += laneGap / 2 * (len(row) - 1)
		s := centreS - sSize/2
		for _, i := range row {
			if horizontal {
				s += lay.laid[i].h / 2
				lay.laid[i].x = p
				lay.laid[i].y = s
				s += lay.laid[i].h/2 + laneGap/2
			} else {
				s += lay.laid[i].w / 2
				lay.laid[i].y = p
				lay.laid[i].x = s
				s += lay.laid[i].w/2 + laneGap/2
			}
		}
	}
	// Keep every member inside the box.
	for _, i := range sg.nodes {
		l := &lay.laid[i]
		l.x = clamp(l.x, minX+l.w/2, maxX-l.w/2)
		l.y = clamp(l.y, minY+l.h/2, maxY-l.h/2)
	}
}

func clamp(v, low, high int) int {
	if v < low {
		return low
	}
	if v > high {
		return high
	}
	return v
}

func boundingBox(lay *flowLayout, nodes []int) (int, int, int, int) {
	minX, minY := 1<<30, 1<<30
	maxX, maxY := -1<<30, -1<<30
	for _, i := range nodes {
		l := lay.laid[i]
		minX = min(minX, l.x-l.w/2)
		maxX = max(maxX, l.x+l.w/2)
		minY = min(minY, l.y-l.h/2)
		maxY = max(maxY, l.y+l.h/2)
	}
	return minX, minY, maxX, maxY
}

func nodeSize(node *flowNode) (int, int) {
	lines := labelLines(node.label)
	width := 0
	for _, l := range lines {
		width = max(width, textWidth(l, 14))
	}
	h := 26 + 18*len(lines)
	w := max(width+textPad, nodeMinWidth)
	switch node.shape {
	case "diamond":
		w = max(width*2+textPad*2, 150)
		h = max(30+26*len(lines), w/2)
	case "circle":
		d := max(max(width+40, h), 68)
		w, h = d, d
	}
	return w, h
}

// mapPoint maps a TB-space point into the final orientation.
func mapPoint(dir string, p point, width, height int) point {
	switch dir {
	case "BT":
		return point{p.x, height - p.y}
	case "LR":
		return point{p.y, p.x}
	case "RL":
		return point{height - p.y, p.x}
	}
	return p
}

// edgeStroke gives the path attributes for an edge kind, with the styles
// of the linkStyle declarations appended.
func edgeStroke(kind string, styles []stylePair) string {
	attrs := ` fill="none" stroke="#555" stroke-width="2"`
	switch {
	case strings.HasPrefix(kind, "thick"):
		attrs = ` fill="none" stroke="#555" stroke-width="3.5"`
	case strings.HasPrefix(kind, "dotted"):
		attrs = ` fill="none" stroke="#555" stroke-width="2" stroke-dasharray="6 5"`
	}
	return attrs + styleString(styles)
}

// nodeStyles gathers the inline styles and the class declarations of a
// node, declaration order preserved.
func (d *flowDiagram) nodeStyles(n *flowNode) []stylePair {
	pairs := append([]stylePair{}, n.styles...)
	for _, class := range n.classes {
		pairs = append(pairs, d.classes[class]...)
	}
	return pairs
}

func writeFlowchart(d *flowDiagram, lay *flowLayout) []byte {
	width, height := lay.width, lay.height
	if d.dir == "LR" || d.dir == "RL" {
		width, height = height, width
	}
	svg := newSVGBuilder(width, height)

	// Subgraph boxes, outer before inner, so the parents frame their
	// children.
	for _, sg := range slices.Backward(d.subgraphs) {

		if len(sg.nodes) == 0 {
			continue
		}
		minX, minY, maxX, maxY := boundingBox(lay, sg.nodes)
		a := mapPoint(d.dir, point{minX - 20, minY - 40}, width, height)
		b := mapPoint(d.dir, point{maxX + 20, maxY + 18}, width, height)
		x0, y0 := min(a.x, b.x), min(a.y, b.y)
		x1, y1 := max(a.x, b.x), max(a.y, b.y)
		svg.rect(x0, y0, x1-x0, y1-y0, 8, ` fill="#f5f5f5" fill-opacity="0.7" stroke="#999"`)
		svg.text(point{x0 + 10, y0 + 18}, sg.title, "start", ` font-size="14" font-weight="bold"`)
	}

	// Edges.
	type drawn struct {
		p0, c1, c2, p3 point
		arrow          bool
		label          string
		attrs          string
	}
	var drawnEdges []drawn
	for i, e := range d.edges {
		a, b := lay.laid[e.from], lay.laid[e.to]
		var p0, p3, c1, c2 point
		switch {
		case e.from == e.to:
			p0 = point{a.x + a.w/2, a.y - 8}
			p3 = point{a.x + a.w/2, a.y + 8}
			c1 = point{a.x + a.w/2 + 46, a.y - 28}
			c2 = point{a.x + a.w/2 + 46, a.y + 28}
		case lay.rank[e.to] > lay.rank[e.from]:
			p0 = point{a.x, a.y + a.h/2}
			p3 = point{b.x, b.y - b.h/2}
			mid := max((p3.y-p0.y)/2, 24)
			c1 = point{p0.x, p0.y + mid}
			c2 = point{p3.x, p3.y - mid}
		case lay.rank[e.to] < lay.rank[e.from]:
			p0 = point{a.x, a.y - a.h/2}
			p3 = point{b.x, b.y + b.h/2}
			mid := max((p0.y-p3.y)/2, 24)
			c1 = point{p0.x, p0.y - mid}
			c2 = point{p3.x, p3.y + mid}
		default:
			if b.x >= a.x {
				p0 = point{a.x + a.w/2, a.y}
				p3 = point{b.x - b.w/2, b.y}
			} else {
				p0 = point{a.x - a.w/2, a.y}
				p3 = point{b.x + b.w/2, b.y}
			}
			c1 = point{p0.x + 42, p0.y}
			c2 = point{p3.x - 42, p3.y}
		}
		styles := append([]stylePair{}, d.linkDefault...)
		styles = append(styles, d.linkByIndex[i]...)
		drawnEdges = append(drawnEdges, drawn{
			p0:    mapPoint(d.dir, p0, width, height),
			c1:    mapPoint(d.dir, c1, width, height),
			c2:    mapPoint(d.dir, c2, width, height),
			p3:    mapPoint(d.dir, p3, width, height),
			arrow: strings.HasSuffix(e.kind, "-arrow"),
			label: e.label,
			attrs: edgeStroke(e.kind, styles),
		})
	}
	for _, e := range drawnEdges {
		svg.path(fmt.Sprintf("M %d %d C %d %d, %d %d, %d %d", e.p0.x, e.p0.y, e.c1.x, e.c1.y, e.c2.x, e.c2.y, e.p3.x, e.p3.y), e.attrs)
	}
	for _, e := range drawnEdges {
		if !e.arrow {
			continue
		}
		svg.polygon(arrowHead(e.p3, e.c2), ` fill="#555"`)
	}

	// Nodes.
	for i, node := range d.nodes {
		l := lay.laid[i]
		c := mapPoint(d.dir, point{l.x, l.y}, width, height)
		w, h := l.w, l.h
		if d.dir == "LR" || d.dir == "RL" {
			w, h = h, w
		}
		attrs := ` fill="#ffffff" stroke="#333" stroke-width="1.5"` + styleString(d.nodeStyles(node))
		switch node.shape {
		case "round":
			svg.rect(c.x-w/2, c.y-h/2, w, h, 10, attrs)
		case "stadium":
			svg.rect(c.x-w/2, c.y-h/2, w, h, min(w, h)/2, attrs)
		case "circle":
			svg.rect(c.x-w/2, c.y-h/2, w, h, w/2, attrs)
		case "diamond":
			svg.polygon([]point{{c.x, c.y - h/2}, {c.x + w/2, c.y}, {c.x, c.y + h/2}, {c.x - w/2, c.y}}, attrs)
		case "hex":
			cut := min(20, w/4)
			svg.polygon([]point{{c.x - w/2 + cut, c.y - h/2}, {c.x + w/2 - cut, c.y - h/2}, {c.x + w/2, c.y}, {c.x + w/2 - cut, c.y + h/2}, {c.x - w/2 + cut, c.y + h/2}, {c.x - w/2, c.y}}, attrs)
		case "asym":
			svg.polygon([]point{{c.x - w/2, c.y - h/2}, {c.x + w/2 - 18, c.y - h/2}, {c.x + w/2, c.y}, {c.x + w/2 - 18, c.y + h/2}, {c.x - w/2, c.y + h/2}}, attrs)
		case "sub":
			svg.rect(c.x-w/2, c.y-h/2, w, h, 0, attrs)
			if d.dir == "LR" || d.dir == "RL" {
				svg.line(c.x, c.y-h/2+5, c.x, c.y+h/2-5, ` stroke="#333" stroke-width="1.5"`)
				svg.line(c.x, c.y-h/2+10, c.x, c.y+h/2-10, ` stroke="#333" stroke-width="1.5"`)
			} else {
				svg.line(c.x-w/2+5, c.y, c.x+w/2-5, c.y, ` stroke="#333" stroke-width="1.5"`)
			}
		default:
			svg.rect(c.x-w/2, c.y-h/2, w, h, 0, attrs)
		}
		lines := labelLines(node.label)
		for k, ln := range lines {
			y := c.y + 5 + (k-(len(lines)-1)/2)*18
			if len(lines)%2 == 0 {
				y = c.y - 4 + k*18
			}
			svg.text(point{c.x, y}, ln, "middle", ` font-size="14"`)
		}
	}

	// Edge labels on top.
	for _, e := range drawnEdges {
		if e.label == "" {
			continue
		}
		mid := point{(e.p0.x + e.p3.x) / 2, (e.p0.y+e.p3.y)/2 - 7}
		svg.text(mid, e.label, "middle", ` font-size="12"`)
	}
	return svg.finish()
}

// arrowHead builds a filled triangle at tip pointing from the control
// point towards the tip.
func arrowHead(tip, ctrl point) []point {
	dx, dy := tip.x-ctrl.x, tip.y-ctrl.y
	n := max(abs(dx)+abs(dy), 1)
	ux, uy := dx*1000/n, dy*1000/n
	base := point{tip.x - ux*11/1000, tip.y - uy*11/1000}
	return []point{
		tip,
		{base.x - uy*5/1000, base.y + ux*5/1000},
		{base.x + uy*5/1000, base.y - ux*5/1000},
	}
}

func abs(v int) int {
	if v < 0 {
		return -v
	}
	return v
}