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

package diagram

import (
	"fmt"
	"strings"
)

// stylePair is one key:value declaration of a classDef, style or
// linkStyle statement.
type stylePair struct {
	key   string
	value string
}

type flowNode struct {
	id      string
	label   string
	shape   string
	classes []string
	styles  []stylePair
}

type flowEdge struct {
	from, to int
	kind     string
	label    string
	styles   []stylePair
}

type flowSubgraph struct {
	id, title, dir string
	nodes          []int
	depth          int
}

type flowDiagram struct {
	dir         string
	nodes       []*flowNode
	index       map[string]int
	edges       []*flowEdge
	subgraphs   []*flowSubgraph
	classes     map[string][]stylePair
	linkDefault []stylePair
	linkByIndex map[int][]stylePair
}

func parseFlowchart(lines []string, dir string) (*flowDiagram, error) {
	d := &flowDiagram{
		dir:         dir,
		index:       map[string]int{},
		classes:     map[string][]stylePair{},
		linkByIndex: map[int][]stylePair{},
	}
	var stack []*flowSubgraph
	for n, line := range lines {
		if err := d.parseFlowStatement(line, &stack, n+2); err != nil {
			return nil, err
		}
	}
	if len(stack) > 0 {
		return nil, fmt.Errorf("diagram: flowchart line %d: subgraph %q has no end", len(lines)+1, stack[len(stack)-1].id)
	}
	return d, nil
}

func (d *flowDiagram) parseFlowStatement(line string, stack *[]*flowSubgraph, lineNo int) error {
	sc := &scanner{src: []rune(line)}
	word := sc.word()
	switch strings.ToLower(word) {
	case "subgraph":
		return d.parseSubgraph(sc, stack, lineNo)
	case "end":
		if len(*stack) == 0 {
			return fmt.Errorf("diagram: flowchart line %d: end without subgraph", lineNo)
		}
		*stack = (*stack)[:len(*stack)-1]
		return nil
	case "direction":
		sc.skipSpaces()
		dir := strings.ToUpper(sc.rest())
		if !flowDirections[dir] {
			return fmt.Errorf("diagram: flowchart line %d: unknown flow direction %q", lineNo, dir)
		}
		if len(*stack) > 0 {
			(*stack)[len(*stack)-1].dir = dir
		} else {
			d.dir = dir
		}
		return nil
	case "classdef":
		return d.parseClassDef(sc, lineNo)
	case "class":
		return d.parseClass(sc, lineNo)
	case "style":
		return d.parseStyle(sc, lineNo)
	case "linkstyle":
		return d.parseLinkStyle(sc, lineNo)
	case "click":
		return fmt.Errorf("diagram: flowchart line %d: click is unsupported, the SVG carries no interactivity", lineNo)
	}
	sc.pos = 0
	return d.parseChain(sc, *stack, lineNo)
}

// node returns the index of a node, creating it when it first appears.
// The innermost open subgraph claims a node at its first appearance.
func (d *flowDiagram) node(id, label, shape string, stack []*flowSubgraph) int {
	if i, ok := d.index[id]; ok {
		if shape != "" {
			d.nodes[i].shape = shape
			d.nodes[i].label = label
		}
		return i
	}
	if label == "" {
		label = id
	}
	if shape == "" {
		shape = "rect"
	}
	n := &flowNode{id: id, label: label, shape: shape}
	d.nodes = append(d.nodes, n)
	i := len(d.nodes) - 1
	d.index[id] = i
	if len(stack) > 0 {
		sg := stack[len(stack)-1]
		sg.nodes = append(sg.nodes, i)
	}
	return i
}

func (d *flowDiagram) parseChain(sc *scanner, stack []*flowSubgraph, lineNo int) error {
	left, err := d.parseNodeList(sc, stack, lineNo)
	if err != nil {
		return err
	}
	for {
		sc.skipSpaces()
		if !sc.atEdgeChar() {
			return nil
		}
		kind, label, err := sc.parseEdge(lineNo)
		if err != nil {
			return err
		}
		right, err := d.parseNodeList(sc, stack, lineNo)
		if err != nil {
			return err
		}
		for _, from := range left {
			for _, to := range right {
				d.edges = append(d.edges, &flowEdge{from: from, to: to, kind: kind, label: label})
			}
		}
		left = right
	}
}

// parseNodeList reads one or more nodes separated by &. A missing node is
// an error: an edge must land somewhere.
func (d *flowDiagram) parseNodeList(sc *scanner, stack []*flowSubgraph, lineNo int) ([]int, error) {
	var list []int
	for {
		i, err := d.parseNode(sc, stack, lineNo)
		if err != nil {
			return nil, err
		}
		list = append(list, i)
		sc.skipSpaces()
		if sc.peek() == '&' {
			sc.pos++
			continue
		}
		return list, nil
	}
}

func (d *flowDiagram) parseNode(sc *scanner, stack []*flowSubgraph, lineNo int) (int, error) {
	sc.skipSpaces()
	var id strings.Builder
	for !sc.eof() && isIDRune(sc.peek()) {
		id.WriteRune(sc.peek())
		sc.pos++
	}
	if id.Len() == 0 {
		return 0, fmt.Errorf("diagram: flowchart line %d: expected a node", lineNo)
	}
	label, shape := sc.parseShape()
	return d.node(id.String(), label, shape, stack), nil
}

// parseShape reads an optional shape with its label.
func (sc *scanner) parseShape() (string, string) {
	switch {
	case sc.hasPrefix("[["):
		return sc.bracketed("[[", "]]"), "sub"
	case sc.hasPrefix("(["):
		return sc.bracketed("([", "])"), "stadium"
	case sc.hasPrefix("(("):
		return sc.bracketed("((", "))"), "circle"
	case sc.hasPrefix("["):
		return sc.bracketed("[", "]"), "rect"
	case sc.hasPrefix("{{"):
		return sc.bracketed("{{", "}}"), "hex"
	case sc.hasPrefix("{"):
		return sc.bracketed("{", "}"), "diamond"
	case sc.hasPrefix("("):
		return sc.bracketed("(", ")"), "round"
	case sc.hasPrefix(">"):
		return sc.bracketed(">", "]"), "asym"
	}
	return "", ""
}

// bracketed consumes the opener and reads until the closer, honouring
// double quotes. A label that is exactly a quoted string loses its
// quotes.
func (sc *scanner) bracketed(open, close string) string {
	sc.pos += len(open)
	start := sc.pos
	for !sc.eof() {
		if sc.peek() == '"' {
			sc.pos++
			for !sc.eof() && sc.peek() != '"' {
				sc.pos++
			}
			sc.pos++
			continue
		}
		if sc.hasPrefix(close) {
			label := string(sc.src[start:sc.pos])
			sc.pos += len(close)
			return unquote(strings.TrimSpace(label))
		}
		sc.pos++
	}
	return unquote(strings.TrimSpace(string(sc.src[start:])))
}

func unquote(s string) string {
	if len(s) >= 2 && strings.HasPrefix(s, `"`) && strings.HasSuffix(s, `"`) {
		return s[1 : len(s)-1]
	}
	return s
}

func (d *flowDiagram) parseSubgraph(sc *scanner, stack *[]*flowSubgraph, lineNo int) error {
	sc.skipSpaces()
	id := sc.word()
	title := id
	sc.skipSpaces()
	switch {
	case id == "":
		id = fmt.Sprintf("subgraph-%d", len(d.subgraphs))
		title = id
	case sc.hasPrefix("["):
		title = sc.bracketed("[", "]")
	default:
		rest := strings.TrimSpace(sc.rest())
		if rest != "" {
			id = id + " " + rest
			title = id
		}
	}
	sg := &flowSubgraph{id: id, title: title, depth: len(*stack)}
	d.subgraphs = append(d.subgraphs, sg)
	*stack = append(*stack, sg)
	return nil
}

func (d *flowDiagram) parseClassDef(sc *scanner, lineNo int) error {
	sc.skipSpaces()
	name := sc.word()
	if name == "" {
		return fmt.Errorf("diagram: flowchart line %d: classDef needs a name", lineNo)
	}
	pairs, err := parseStylePairs(sc.rest())
	if err != nil {
		return fmt.Errorf("diagram: flowchart line %d: %v", lineNo, err)
	}
	d.classes[name] = append(d.classes[name], pairs...)
	return nil
}

func (d *flowDiagram) parseClass(sc *scanner, lineNo int) error {
	sc.skipSpaces()
	ids := strings.Split(sc.word(), ",")
	sc.skipSpaces()
	names := strings.Split(strings.TrimSpace(sc.rest()), ",")
	if len(ids) == 0 || ids[0] == "" || len(names) == 0 || names[0] == "" {
		return fmt.Errorf("diagram: flowchart line %d: class needs nodes and a class name", lineNo)
	}
	for _, id := range ids {
		id = strings.TrimSpace(id)
		if id == "" {
			continue
		}
		i, ok := d.index[id]
		if !ok {
			i = d.node(id, "", "", nil)
		}
		for _, name := range names {
			name = strings.TrimSpace(name)
			if name != "" {
				d.nodes[i].classes = append(d.nodes[i].classes, name)
			}
		}
	}
	return nil
}

func (d *flowDiagram) parseStyle(sc *scanner, lineNo int) error {
	sc.skipSpaces()
	id := sc.word()
	if id == "" {
		return fmt.Errorf("diagram: flowchart line %d: style needs a node", lineNo)
	}
	pairs, err := parseStylePairs(sc.rest())
	if err != nil {
		return fmt.Errorf("diagram: flowchart line %d: %v", lineNo, err)
	}
	i, ok := d.index[id]
	if !ok {
		i = d.node(id, "", "", nil)
	}
	d.nodes[i].styles = append(d.nodes[i].styles, pairs...)
	return nil
}

func (d *flowDiagram) parseLinkStyle(sc *scanner, lineNo int) error {
	sc.skipSpaces()
	target := sc.word()
	rest := sc.rest()
	pairs, err := parseStylePairs(rest)
	if err != nil {
		return fmt.Errorf("diagram: flowchart line %d: %v", lineNo, err)
	}
	if strings.ToLower(target) == "default" {
		d.linkDefault = append(d.linkDefault, pairs...)
		return nil
	}
	for part := range strings.SplitSeq(target, ",") {
		var n int
		if _, err := fmt.Sscanf(strings.TrimSpace(part), "%d", &n); err != nil {
			return fmt.Errorf("diagram: flowchart line %d: linkStyle needs an index or default", lineNo)
		}
		d.linkByIndex[n] = append(d.linkByIndex[n], pairs...)
	}
	return nil
}

// parseStylePairs reads comma separated key:value declarations.
func parseStylePairs(s string) ([]stylePair, error) {
	var pairs []stylePair
	for part := range strings.SplitSeq(s, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		key, value, found := strings.Cut(part, ":")
		if !found {
			return nil, fmt.Errorf("expected key:value, got %q", part)
		}
		pairs = append(pairs, stylePair{key: strings.TrimSpace(key), value: unquote(strings.TrimSpace(value))})
	}
	if len(pairs) == 0 {
		return nil, fmt.Errorf("expected style declarations")
	}
	return pairs, nil
}

func isIDRune(r rune) bool {
	return r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '_'
}

// scanner walks one statement.
type scanner struct {
	src []rune
	pos int
}

func (sc *scanner) eof() bool { return sc.pos >= len(sc.src) }

func (sc *scanner) peek() rune {
	if sc.eof() {
		return 0
	}
	return sc.src[sc.pos]
}

func (sc *scanner) skipSpaces() {
	for !sc.eof() && (sc.peek() == ' ' || sc.peek() == '\t') {
		sc.pos++
	}
}

func (sc *scanner) rest() string { return string(sc.src[sc.pos:]) }

func (sc *scanner) word() string {
	start := sc.pos
	for !sc.eof() && sc.peek() != ' ' && sc.peek() != '\t' {
		sc.pos++
	}
	return string(sc.src[start:sc.pos])
}

func (sc *scanner) hasPrefix(s string) bool {
	runes := []rune(s)
	if sc.pos+len(runes) > len(sc.src) {
		return false
	}
	for i, r := range runes {
		if sc.src[sc.pos+i] != r {
			return false
		}
	}
	return true
}

func (sc *scanner) atEdgeChar() bool {
	c := sc.peek()
	return c == '-' || c == '.' || c == '='
}

func isEdgeRune(r rune) bool { return r == '-' || r == '.' || r == '=' }

// edgeRun consumes a run of edge characters and returns it.
func (sc *scanner) edgeRun() string {
	start := sc.pos
	for !sc.eof() && isEdgeRune(sc.peek()) {
		sc.pos++
	}
	return string(sc.src[start:sc.pos])
}

// parseEdge reads one edge token: a run of -, . and = characters, an
// optional arrow head, and a label either in pipes or as text between two
// runs.
func (sc *scanner) parseEdge(lineNo int) (kind, label string, err error) {
	runs := sc.edgeRun()
	arrow := false
	switch {
	case sc.peek() == '|':
		sc.pos++
		start := sc.pos
		for !sc.eof() && sc.peek() != '|' {
			sc.pos++
		}
		label = unquote(strings.TrimSpace(string(sc.src[start:sc.pos])))
		if !sc.eof() {
			sc.pos++
		}
	case sc.peek() == '>':
		sc.pos++
		arrow = true
		if sc.peek() == '|' {
			// A pipe label behind the arrow head: -->|text|
			sc.pos++
			start := sc.pos
			for !sc.eof() && sc.peek() != '|' {
				sc.pos++
			}
			label = unquote(strings.TrimSpace(string(sc.src[start:sc.pos])))
			if !sc.eof() {
				sc.pos++
			}
		}
	default:
		save := sc.pos
		sc.skipSpaces()
		if !sc.eof() && !isEdgeRune(sc.peek()) {
			start := sc.pos
			for !sc.eof() && !isEdgeRune(sc.peek()) {
				sc.pos++
			}
			second := sc.edgeRun()
			if second == "" {
				// No closing run: the words ahead are the next node, not a
				// label.
				sc.pos = save
			} else {
				label = unquote(strings.TrimSpace(string(sc.src[start : sc.pos-len(second)])))
				runs += second
				if sc.peek() == '>' {
					sc.pos++
					arrow = true
				}
			}
		} else {
			sc.pos = save
		}
	}
	if len(runs) < 2 || !arrow && len(runs) < 3 {
		return "", "", fmt.Errorf("diagram: flowchart line %d: unfinished edge token", lineNo)
	}
	switch {
	case strings.Contains(runs, "="):
		kind = "thick"
	case strings.Contains(runs, "."):
		kind = "dotted"
	default:
		kind = "solid"
	}
	if arrow {
		kind += "-arrow"
	}
	return kind, label, nil
}