package drawing

import (
	"bytes"
	"fmt"
	"io"
	"os"
)

type Digraph struct {
	Title     string
	Subgraphs []Subgraph
}

type Subgraph interface {
	Print()
	Fprint(w io.Writer)
}

type DotNode struct {
	Data       string
	Children   []*DotNode
	Attributes map[string]string
}

type DotTree struct {
	Root       *DotNode
	Nodes      map[string]*DotNode
	Attributes map[string]string
}

func (d *Digraph) String() string {
	buf := &bytes.Buffer{}
	d.Fprint(buf)
	return buf.String()
}

func (d *Digraph) Print() {
	d.Fprint(os.Stdout)
}

func (d *Digraph) Fprint(w io.Writer) {
	fmt.Fprintf(w, "strict digraph \"%v\" {\n", d.Title)

	for _, sub := range d.Subgraphs {
		sub.Fprint(w)
	}

	fmt.Fprintln(w, "}")
}

func (t *DotTree) Print() {
	t.Fprint(os.Stdout)
}

func (t *DotTree) Fprint(w io.Writer) {
	fmt.Fprintf(w, "subgraph \"cluster_%v\" {\n", t.Root.Data)
	for attr, val := range t.Attributes {
		fmt.Fprintf(w, `%s%s=%s%s`, "\t", attr, val, "\n")
	}
	t.Root.print(w)
	fmt.Fprintln(w, "}")
}

func (node *DotNode) print(w io.Writer) {
	// for edges from node N to child C
	for _, child := range node.Children {
		// write N
		nodeAttrs := ""
		if len(node.Attributes) > 0 {
			nodeAttrs += " ["
			for attr, val := range node.Attributes {
				nodeAttrs += fmt.Sprintf(`%s=%s `, attr, val)
			}
			nodeAttrs += "]"
		}
		fmt.Fprintf(w, `"%s"%s%s`, node.Data, nodeAttrs, "\n")

		// write C
		childAttrs := ""
		if len(child.Attributes) > 0 {
			childAttrs += " ["
			for attr, val := range child.Attributes {
				childAttrs += fmt.Sprintf(`%s=%s `, attr, val)
			}
			childAttrs += "]"
		}
		fmt.Fprintf(w, `"%s"%s%s`, child.Data, childAttrs, "\n")

		// write N->C
		fmt.Fprintf(w, `"%s" -> "%s"%s`, node.Data, child.Data, "\n")

		// repeat for children of C
		child.print(w)
	}
}