package main

import (
	"flag"
	"fmt"
	"os"
	"reflect"
	"sort"
	"strings"

	"github.com/tendermint/tendermint/blockchain"
	"github.com/tendermint/tendermint/mempool"
	"github.com/tendermint/tendermint/state"

	"github.com/tendermint/tendermint/consensus"

	cfg "github.com/tendermint/tendermint/config"
	"github.com/tendermint/tendermint/node"
)

// depNode is a node in a dependency graph.
type depNode struct {
	rt        reflect.Type
	name      string
	pkg       string
	kind      reflect.Kind
	level     int
	dependsOn map[string]*depNode
}

var supportedStructs = map[string]reflect.Type{
	"blockchain.BlockchainReactor": reflect.TypeOf((*blockchain.BlockchainReactor)(nil)).Elem(),
	"consensus.ConsensusReactor":   reflect.TypeOf((*consensus.ConsensusReactor)(nil)).Elem(),
	"config.Config":                reflect.TypeOf((*cfg.Config)(nil)).Elem(),
	"mempool.Mempool":              reflect.TypeOf((*mempool.Mempool)(nil)).Elem(),
	"mempool.MempoolReactor":       reflect.TypeOf((*mempool.MempoolReactor)(nil)).Elem(),
	"node.Node":                    reflect.TypeOf((*node.Node)(nil)).Elem(),
	"state.BlockExecutor":          reflect.TypeOf((*state.BlockExecutor)(nil)).Elem(),
}

var rootCmd *flag.FlagSet

var (
	flagList    bool
	flagStruct  string
	flagDepth   int
	flagExclude string
)

func init() {
	rootCmd = flag.NewFlagSet("root", flag.ExitOnError)
	rootCmd.Usage = func() {
		fmt.Println(`Object graph visualization for Tendermint data structures.

Usage:
  tm-obj-viz [flags] <struct> <outputFile>

Flags:`)
		rootCmd.PrintDefaults()
		fmt.Println("")
	}
	rootCmd.IntVar(&flagDepth, "depth", 3, "how many levels deep to recursively analyze the object graph")
	rootCmd.BoolVar(&flagList, "list", false, "specify this flag to list the available structs for visualization")
	rootCmd.StringVar(&flagStruct, "struct", "node.Node", "the struct whose object graph is to be printed out")
	rootCmd.StringVar(&flagExclude, "exclude", "", "a comma-separated list of types to exclude when rendering the graph")
}

func showSupportedStructs() {
	structs := make([]string, 0)
	for structName := range supportedStructs {
		structs = append(structs, structName)
	}
	sort.SliceStable(structs[:], func(i, j int) bool {
		return strings.Compare(structs[i], structs[j]) < 0
	})
	fmt.Println("Supported structs:")
	for _, structName := range structs {
		fmt.Printf("  %s\n", structName)
	}
}

func longNodeID(t reflect.Type) string {
	return fmt.Sprintf("%s.%s", t.PkgPath(), t.Name())
}

func (n *depNode) longPath() string {
	return longNodeID(n.rt)
}

func shortNodeID(t reflect.Type) string {
	return strings.Replace(t.PkgPath(), "github.com/tendermint/tendermint/", "", 1) + "." + t.Name()
}

func (n *depNode) shortPath() string {
	return shortNodeID(n.rt)
}

func (n *depNode) isTendermint() bool {
	tmparts := strings.Split(n.pkg, "github.com/tendermint/tendermint")
	return (len(tmparts) > 1) && (len(strings.Split(tmparts[1], "vendor/")) < 2)
}

func (n *depNode) attr() string {
	attrParts := make([]string, 0)
	if n.isTendermint() {
		switch n.kind {
		case reflect.Struct:
			attrParts = append(attrParts, "style=filled", "fillcolor=lightblue")
		case reflect.Interface:
			attrParts = append(attrParts, "style=filled", "fillcolor=\"#eeeeee\"")
		default:
			attrParts = append(attrParts, "style=filled", "fillcolor=\"#ccffcc\"")
		}
	} else {
		attrParts = append(attrParts, "style=dashed")
	}
	return strings.Join(attrParts, ",")
}

func buildDepGraph(n *depNode, maxLevels, curLevel int, seenDeps map[string]*depNode, exclude map[string]interface{}) int {
	if maxLevels == curLevel {
		return curLevel
	}
	highestLevel := curLevel
	if n.kind == reflect.Struct {
		for i := 0; i < n.rt.NumField(); i++ {
			subFieldType := n.rt.Field(i).Type
			for subFieldType.Kind() == reflect.Ptr {
				subFieldType = subFieldType.Elem()
			}
			if len(subFieldType.Name()) > 0 && len(subFieldType.PkgPath()) > 0 {
				longID := longNodeID(subFieldType)
				shortID := shortNodeID(subFieldType)
				// if we haven't explicitly excluded this particular type
				if _, ok := exclude[shortID]; !ok {
					subNode, ok := seenDeps[longID]
					// if we haven't seen this type before
					if !ok {
						subNode = &depNode{
							rt:        subFieldType,
							name:      subFieldType.Name(),
							pkg:       subFieldType.PkgPath(),
							kind:      subFieldType.Kind(),
							level:     n.level + 1,
							dependsOn: make(map[string]*depNode),
						}
						seenDeps[longID] = subNode
						if subNode.kind == reflect.Struct && subNode.isTendermint() {
							highestLevel = buildDepGraph(subNode, maxLevels, curLevel+1, seenDeps, exclude)
						}
					}
					n.dependsOn[longID] = subNode
				}
			}
		}
	}
	return highestLevel
}

// printDepNode recursively prints the given dependency node's sub-dependencies.
func printDepNode(n *depNode, curLevel int, seenEdges map[string]interface{}, rootIndent string) {
	indent := rootIndent + strings.Repeat("  ", curLevel+1)
	for _, subNode := range n.dependsOn {
		edge := fmt.Sprintf("\"%s\"", n.shortPath()) + " -> " + fmt.Sprintf("\"%s\"", subNode.shortPath())
		if _, ok := seenEdges[edge]; !ok {
			seenEdges[edge] = nil
			fmt.Println(indent + edge + ";")
			printDepNode(subNode, curLevel+1, seenEdges, rootIndent)
		}
	}
}

func printGraphFormatting(n *depNode, seenNodes map[string]*depNode) {
	attrs := n.attr()
	if len(attrs) > 0 {
		fmt.Println(fmt.Sprintf("  \"%s\" [%s];", n.shortPath(), n.attr()))
	}
	for id, subNode := range n.dependsOn {
		if _, ok := seenNodes[id]; !ok {
			seenNodes[id] = subNode
			printGraphFormatting(subNode, seenNodes)
		}
	}
}

func printLegend() {
	fmt.Println("  subgraph cluster_legend {")
	fmt.Println("    fontname = \"helvetica\";")
	fmt.Println("    fontsize = 18;")
	fmt.Println("    label = \"Legend\";")
	fmt.Println("    color = lightgray;")
	fmt.Println("    \"Tendermint struct\" [style=filled,fillcolor=lightblue];")
	fmt.Println("    \"Tendermint interface\" [style=filled,fillcolor=\"#eeeeee\"];")
	fmt.Println("    \"Tendermint type\" [style=filled,fillcolor=\"#ccffcc\"];")
	fmt.Println("    \"Non-Tendermint type\" [style=dashed];")
	fmt.Println("  }")
}

func printExcluded(exclude map[string]interface{}) {
	if len(exclude) > 0 {
		fmt.Println("  subgraph cluster_exclude {")
		fmt.Println("    fontname = \"helvetica\";")
		fmt.Println("    fontsize = 18;")
		fmt.Println("    label = \"Excluded\";")
		fmt.Println("    color = lightgray;")
		for e := range exclude {
			fmt.Println("    \"" + e + "\" [style=filled,fillcolor=\"#ff8888\"];")
		}
		fmt.Println("  }")
	}
}

func printDepGraph(s reflect.Type, maxLevels int, exclude map[string]interface{}) {
	for s.Kind() == reflect.Ptr {
		s = s.Elem()
	}
	rootNode := &depNode{
		rt:        s,
		name:      s.Name(),
		pkg:       s.PkgPath(),
		kind:      s.Kind(),
		level:     0,
		dependsOn: make(map[string]*depNode),
	}
	seenDeps := make(map[string]*depNode)
	buildDepGraph(rootNode, maxLevels, 0, seenDeps, exclude)

	fmt.Println("digraph G {")
	fmt.Println("  rankdir=\"LR\";")
	fmt.Println("  ranksep=\"2 equally\";")
	fmt.Println("  splines=spline;")
	fmt.Println("  ratio=fill;")
	fmt.Println("  node [ fontname=\"helvetica\" ];")

	// then print all the edges
	seenEdges := make(map[string]interface{})
	fmt.Println("  subgraph cluster_graph {")
	fmt.Println("    style=invis;")
	printDepNode(rootNode, 0, seenEdges, "  ")
	fmt.Println("  }")

	// format the different kinds of nodes
	seenDeps = make(map[string]*depNode)
	printGraphFormatting(rootNode, seenDeps)

	// print out which classes have been excluded from the graph
	printExcluded(exclude)

	// print out the legend
	printLegend()

	fmt.Println("}")
}

func main() {
	if err := rootCmd.Parse(os.Args[1:]); err != nil {
		fmt.Println("Failed to parse command line arguments:", err)
		os.Exit(1)
	}

	if flagList {
		showSupportedStructs()
		os.Exit(0)
	}

	t, ok := supportedStructs[flagStruct]
	if !ok {
		fmt.Println("Unrecognized struct:", flagStruct)
		os.Exit(1)
	}
	exclude := make(map[string]interface{})
	for _, e := range strings.Split(flagExclude, ",") {
		if len(e) > 0 {
			exclude[e] = nil
		}
	}
	printDepGraph(t, flagDepth, exclude)
}