// Copyright © 2016 Asteris, LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package graphviz

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/asteris-llc/converge/graph"
	pp "github.com/asteris-llc/converge/prettyprinters"
	"github.com/asteris-llc/converge/rpc/pb"
)

// SubgraphMarkerKey is a type alias for an integer and represents the state
// values for subgraph tracking.
type SubgraphMarkerKey int

const (
	// SubgraphMarkerStart specifies that the current node is the beginning of a
	// subgraph.
	SubgraphMarkerStart SubgraphMarkerKey = iota

	// SubgraphMarkerEnd pecifies that the current node is the end of a subgraph.
	SubgraphMarkerEnd

	// SubgraphMarkerNOP pecifies that the current node should not change the
	// subgraph state.
	SubgraphMarkerNOP
)

// PropertySet is a map of graphviz compatible node or edge options and their
// values.  See http://www.graphviz.org/doc/info/attrs.html for a list of
// supported attributes.  Node that graph and subgraph attributes are not
// currently supported.
type PropertySet map[string]string

// GraphEntity defines the value at a given vertex, providing both it's
// canonical name and the untyped value associated with the vertex.
type GraphEntity struct {
	Name  string      // The canonical node ID
	Value interface{} // The value at the node
}

// PrintProvider allows specific serializable types to be rendered as a
// Graphviz document.
type PrintProvider interface {

	// Given a graph entity, this function shall return a unique ID that will be
	// used for the vertex.  Note that this is not used for display, see
	// VertexGetLabel() for information on vertex displays.
	VertexGetID(GraphEntity) (pp.VisibleRenderable, error)

	// Defines the label associated with the vertex.  This will be the name
	// applied to the vertex when it is drawn.  Labels are stored in the vertex
	// attribute list automatically.  The 'label' parameter should therefore not
	// be returned as part of VertexGetProperties().
	VertexGetLabel(GraphEntity) (pp.VisibleRenderable, error)

	// VertexGetProperties allows the PrintProvider to provide additional
	// attributes to a given vertex.  Note that the 'label' attribute is special
	// and should not be returned as part of this call.
	VertexGetProperties(GraphEntity) PropertySet

	// Defines the label associated with the edge formed between two graph
	// entities.  This will be the name applied to the edge when it is drawn.
	EdgeGetLabel(GraphEntity, GraphEntity) (pp.VisibleRenderable, error)

	// EdgeGetProperties allows the PrintProvider to provide additional
	// attributes to a given edge.  Note that the 'label' attribute is special
	// and should not be returned as part of this call.
	EdgeGetProperties(GraphEntity, GraphEntity) PropertySet

	// Allows the underlying provider to identify vertices that are the beginning
	// of a subgraph.  The function should return one of:
	//   SubgraphMarkerStart
	//   SubgraphMarkerEnd
	//   SubgraphMarkerNOP
	// in order to specify the vertexs' relation to subgraphs.
	SubgraphMarker(GraphEntity) SubgraphMarkerKey
}

// Options specifies global graph options that can be configured for output.
// Arbitrary graphviz options are not supported.
type Options struct {
	// Specifies the way that connections between verices are drawn.  The default
	// is 'spline', other options include: 'line', 'orth', and 'none'.  See
	// http://www.graphviz.org/doc/info/attrs.html#d:splines
	Splines string

	// Specifies the direction of the graph.  See
	// http://www.graphviz.org/doc/info/attrs.html#d:rankdir for more information.
	Rankdir string
}

// DefaultOptions returns an Options struct with default values set.
func DefaultOptions() Options {
	return Options{}
}

// Printer is a DigraphPrettyPrinter implementation for drawing graphviz
// compatible DOT source code from a digraph.
type Printer struct {
	options       Options
	printProvider PrintProvider
	clusterIndex  int
}

// New will create a new graphviz.Printer with the options and print provider
// specified.
func New(opts Options, provider PrintProvider) *Printer {
	return &Printer{
		options:       opts,
		printProvider: provider,
		clusterIndex:  0,
	}
}

// MarkNode will call SubgraphMarker() on the print provider to determine
// whether the current node is the beginning of a subgraph.
func (p *Printer) MarkNode(g *graph.Graph, id string) *pp.SubgraphMarker {
	var val interface{}
	if meta, ok := g.Get(id); ok {
		val = meta.Value()
	}

	entity := GraphEntity{Name: id, Value: val}

	sgState := p.printProvider.SubgraphMarker(entity)
	subgraphID := p.clusterIndex
	switch sgState {
	case SubgraphMarkerStart:
		p.clusterIndex++
		return &pp.SubgraphMarker{
			SubgraphID: subgraphID,
			Start:      true,
		}
	case SubgraphMarkerEnd:
		return &pp.SubgraphMarker{
			SubgraphID: subgraphID,
			Start:      false,
		}
	case SubgraphMarkerNOP:
		return nil
	}
	return nil
}

// DrawNode prints node data by calling VertexGetID(), VertexGetLabel() and
// VertexGetProperties() on the associated print provider.  It will return a
// visible Renderable IFF the underlying VertexID is renderable.
func (p *Printer) DrawNode(g *graph.Graph, id string) (pp.Renderable, error) {
	var val interface{}
	if meta, ok := g.Get(id); ok {
		val = meta.Value()
	}

	graphEntity := GraphEntity{id, val}
	vertexID, err := p.printProvider.VertexGetID(graphEntity)
	if err != nil || !vertexID.Visible() {
		return pp.HiddenString(), err
	}

	vertexLabel, err := p.printProvider.VertexGetLabel(graphEntity)
	if err != nil {
		return pp.HiddenString(), err
	}

	attributes := p.printProvider.VertexGetProperties(graphEntity)
	attributes = maybeSetProperty(attributes, "label", escapeNewline(vertexLabel))
	attributeStr := buildAttributeString(attributes)
	vertexID = escapeNewline(vertexID)

	return pp.SprintfRenderable(
		true,
		"\"%s\" %s;\n",
		vertexID,
		attributeStr,
	), nil
}

// DrawEdge prints edge data in a fashion similar to DrawNode.  It will return a
// visible Renderable IFF the source and destination vertices are visible.
func (p *Printer) DrawEdge(g *graph.Graph, id1, id2 string) (pp.Renderable, error) {
	var srcVal, destVal interface{}

	if src, ok := g.Get(id1); ok {
		srcVal = src.Value()
	}

	if dest, ok := g.Get(id2); ok {
		destVal = dest.Value()
	}

	sourceEntity := GraphEntity{id1, srcVal}
	destEntity := GraphEntity{id2, destVal}

	sourceVertex, err := p.printProvider.VertexGetID(sourceEntity)
	if err != nil {
		return pp.HiddenString(), err
	}
	destVertex, err := p.printProvider.VertexGetID(destEntity)
	if err != nil {
		return pp.HiddenString(), err
	}
	label, err := p.printProvider.EdgeGetLabel(sourceEntity, destEntity)
	if err != nil {
		return pp.HiddenString(), err
	}
	attributes := p.printProvider.EdgeGetProperties(sourceEntity, destEntity)
	maybeSetProperty(attributes, "label", escapeNewline(label))

	srcVert, sok := srcVal.(*pb.GraphComponent_Vertex)
	destVert, dok := destVal.(*pb.GraphComponent_Vertex)
	if sok && dok && srcVert.Kind == "module" && destVert.Kind != "module" {
		return pp.HiddenString(), nil
	}

	edgeStr := fmt.Sprintf("\"%s\" -> \"%s\" %s;\n",
		escapeNewline(destVertex),
		escapeNewline(sourceVertex),
		buildAttributeString(attributes),
	)
	visible := sourceVertex.Visible() && destVertex.Visible()

	return pp.RenderableString(edgeStr, visible), nil
}

// StartSubgraph returns a string with the beginning of the subgraph cluster
func (p *Printer) StartSubgraph(g *graph.Graph, startNode string, subgraphID pp.SubgraphID) (pp.Renderable, error) {
	clusterStart := fmt.Sprintf("subgraph cluster_%d {\n", subgraphID.(int))
	return pp.VisibleString(clusterStart), nil
}

// FinishSubgraph provides the closing '}' for a subgraph
func (*Printer) FinishSubgraph(*graph.Graph, pp.SubgraphID) (pp.Renderable, error) {
	return pp.VisibleString("}\n"), nil
}

// StartPP begins the DOT output as an unnamed digraph.
func (p *Printer) StartPP(*graph.Graph) (pp.Renderable, error) {
	attrs := p.GraphAttributes()
	return pp.VisibleString(fmt.Sprintf("digraph {\n%s\n", attrs)), nil
}

// GraphAttributes returns a string containing all of the global graph
// attributes specified in Options.
func (p *Printer) GraphAttributes() string {
	var buffer bytes.Buffer
	if p.options.Splines != "" {
		_, _ = buffer.WriteString(fmt.Sprintf("splines = \"%s\";\n", p.options.Splines))
	}
	if p.options.Rankdir != "" {
		_, _ = buffer.WriteString(fmt.Sprintf("rankdir = \"%s\";\n", p.options.Rankdir))
	}
	return buffer.String()
}

// FinishPP returns the closing '}' to end the DOT file.
func (*Printer) FinishPP(*graph.Graph) (pp.Renderable, error) {
	return pp.VisibleString("}"), nil
}

// buildAttributeString takes a set of attribute keys and values and generates a
// DOT format attribute set in the form of [key1="value1",key2="value2"...]
func buildAttributeString(p PropertySet) string {
	var attrs []string
	if len(p) == 0 {
		return ""
	}
	for k, v := range p {
		attribute := fmt.Sprintf(` %s="%s" `, k, escapeNewline(pp.VisibleString(v)))
		attrs = append(attrs, attribute)
	}
	return fmt.Sprintf("[%s]", strings.Join(attrs, ","))
}

// DefaultProvider returns an empty BasicProvider as a convenience
func DefaultProvider() PrintProvider {
	return BasicProvider{}
}

// A GraphIDProvider is a basic PrintProvider for Graphviz that uses the Node ID
// from the digraph as the Vertex ID and label.
type GraphIDProvider struct {
	BasicProvider
}

// IDProvider is a convenience function to generate a GraphIDProvider.
func IDProvider() PrintProvider {
	return GraphIDProvider{}
}

// A BasicProvider is a basic PrintProvider for Graphviz that uses the value
// in the default format (%v) as the node label and id.
type BasicProvider struct{}

// VertexGetID provides a basic implementation that returns the %v quoted value
// of the node.
func (p BasicProvider) VertexGetID(e GraphEntity) (pp.VisibleRenderable, error) {
	return pp.VisibleString(fmt.Sprintf("%v", e.Value)), nil
}

// VertexGetLabel provides a basic implementation that returns the %v quoted
// value of the node (as with VertexGetID).
func (p BasicProvider) VertexGetLabel(e GraphEntity) (pp.VisibleRenderable, error) {
	return pp.VisibleString(fmt.Sprintf("%v", e.Value)), nil
}

// VertexGetProperties provides a basic implementation that returns an empty
// property set.
func (p BasicProvider) VertexGetProperties(GraphEntity) PropertySet {
	return PropertySet{}
}

// EdgeGetLabel provides a basic implementation leaves the edge unlabeled.
func (p BasicProvider) EdgeGetLabel(GraphEntity, GraphEntity) (pp.VisibleRenderable, error) {
	return pp.HiddenString(), nil
}

// EdgeGetProperties provides a basic implementation that returns an empty
// property set.
func (p BasicProvider) EdgeGetProperties(GraphEntity, GraphEntity) PropertySet {
	return PropertySet{}
}

// SubgraphMarker provides a basic implementation that returns a NOP.
func (p BasicProvider) SubgraphMarker(GraphEntity) SubgraphMarkerKey {
	return SubgraphMarkerNOP
}

// VertexGetID provides a basic implementation that uses the ID from the graph
// to generate the VertexID.
func (p GraphIDProvider) VertexGetID(e GraphEntity) (pp.VisibleRenderable, error) {
	return pp.VisibleString(e.Name), nil
}

// VertexGetLabel provides a basic implementation that uses the ID from the
// graph to generate the Vertex Label.
func (p GraphIDProvider) VertexGetLabel(e GraphEntity) (pp.VisibleRenderable, error) {
	return pp.VisibleString(e.Name), nil
}

// Replace embedded newlines with their escaped form.
func escapeNewline(r pp.Renderable) pp.VisibleRenderable {
	return pp.VisibleString(strings.Replace(r.String(), "\n", "\\n", -1))
}

// maybeSetProperty will add an attribute to the property set iff the value is
// renderable.
func maybeSetProperty(p PropertySet, key string, r pp.VisibleRenderable) PropertySet {
	if r.Visible() {
		p[key] = r.String()
	}
	return p
}
