package visualization

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestAddNode(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	
	assert.NotNil(t, dag.nodes["T1"])
	assert.Equal(t, "completed", dag.nodes["T1"].Status)
}

func TestAddEdge(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	dag.AddNode("T2", "completed", 200)
	
	err := dag.AddEdge("T1", "T2")
	assert.NoError(t, err)
	assert.Equal(t, 1, len(dag.edges["T1"]))
}

func TestAddEdgeNotFound(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	
	err := dag.AddEdge("T1", "T2")
	assert.Error(t, err)
}

func TestMarkCriticalPath(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	dag.AddNode("T2", "completed", 200)
	
	err := dag.MarkCriticalPath([]string{"T1", "T2"})
	assert.NoError(t, err)
	assert.True(t, dag.nodes["T1"].Critical)
	assert.True(t, dag.nodes["T2"].Critical)
}

func TestRenderDOT(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	dag.AddNode("T2", "running", 200)
	dag.AddEdge("T1", "T2")
	
	dot := dag.RenderDOT()
	assert.True(t, strings.Contains(dot, "digraph WorkflowDAG"))
	assert.True(t, strings.Contains(dot, "T1"))
	assert.True(t, strings.Contains(dot, "T2"))
	assert.True(t, strings.Contains(dot, "->"))
}

func TestRenderHTML(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	
	html := dag.RenderHTML()
	assert.True(t, strings.Contains(html, "<html>"))
	assert.True(t, strings.Contains(html, "Workflow DAG"))
	assert.True(t, strings.Contains(html, "T1"))
}

func TestGetTopologicalSort(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	dag.AddNode("T2", "completed", 200)
	dag.AddNode("T3", "completed", 150)
	
	dag.AddEdge("T1", "T2")
	dag.AddEdge("T2", "T3")
	
	sorted, err := dag.GetTopologicalSort()
	assert.NoError(t, err)
	assert.Equal(t, 3, len(sorted))
}

func TestGetParallel(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	dag.AddNode("T2", "completed", 200)
	dag.AddNode("T3", "completed", 150)
	
	dag.AddEdge("T1", "T3")
	
	parallel := dag.GetParallel()
	assert.True(t, len(parallel) > 0)
}

func TestGetStats(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	dag.AddNode("T2", "completed", 200)
	dag.AddEdge("T1", "T2")
	
	stats := dag.GetStats()
	assert.Equal(t, 2, stats["node_count"])
	assert.Equal(t, 1, stats["edge_count"])
	assert.Equal(t, 300.0, stats["total_duration"])
}

func TestCriticalPathHighlighting(t *testing.T) {
	dag := NewDAGRenderer()
	dag.AddNode("T1", "completed", 100)
	dag.AddNode("T2", "completed", 200)
	dag.MarkCriticalPath([]string{"T1", "T2"})
	
	dot := dag.RenderDOT()
	assert.True(t, strings.Contains(dot, "fillcolor=red"))
}

func TestComplexDAG(t *testing.T) {
	dag := NewDAGRenderer()
	
	// Create a diamond-shaped DAG
	dag.AddNode("Start", "completed", 50)
	dag.AddNode("A", "completed", 100)
	dag.AddNode("B", "completed", 150)
	dag.AddNode("End", "completed", 50)
	
	dag.AddEdge("Start", "A")
	dag.AddEdge("Start", "B")
	dag.AddEdge("A", "End")
	dag.AddEdge("B", "End")
	
	stats := dag.GetStats()
	assert.Equal(t, 4, stats["node_count"])
	assert.Equal(t, 4, stats["edge_count"])
}
