package visualization

import (
	"fmt"
	"strings"
)

// TaskNode represents a task in the DAG
type TaskNode struct {
	ID       string
	Status   string // pending, running, completed, failed
	Duration float64
	Critical bool
}

// DAGRenderer renders workflow dependency graphs
type DAGRenderer struct {
	nodes map[string]*TaskNode
	edges map[string][]string
}

// NewDAGRenderer creates a new DAG renderer
func NewDAGRenderer() *DAGRenderer {
	return &DAGRenderer{
		nodes: make(map[string]*TaskNode),
		edges: make(map[string][]string),
	}
}

// AddNode adds a task node
func (dr *DAGRenderer) AddNode(id, status string, duration float64) {
	dr.nodes[id] = &TaskNode{
		ID:       id,
		Status:   status,
		Duration: duration,
	}
}

// AddEdge adds a dependency edge
func (dr *DAGRenderer) AddEdge(from, to string) error {
	if _, exists := dr.nodes[from]; !exists {
		return fmt.Errorf("source node not found: %s", from)
	}
	if _, exists := dr.nodes[to]; !exists {
		return fmt.Errorf("target node not found: %s", to)
	}

	dr.edges[from] = append(dr.edges[from], to)
	return nil
}

// MarkCriticalPath marks nodes on the critical path
func (dr *DAGRenderer) MarkCriticalPath(nodes []string) error {
	for _, nodeID := range nodes {
		if node, exists := dr.nodes[nodeID]; exists {
			node.Critical = true
		} else {
			return fmt.Errorf("node not found: %s", nodeID)
		}
	}
	return nil
}

// RenderDOT generates DOT format for Graphviz
func (dr *DAGRenderer) RenderDOT() string {
	var buf strings.Builder

	buf.WriteString("digraph WorkflowDAG {\n")
	buf.WriteString("  rankdir=LR;\n")
	buf.WriteString("  node [shape=box];\n\n")

	// Render nodes
	for _, node := range dr.nodes {
		color := "lightgray"
		if node.Critical {
			color = "red"
		} else if node.Status == "completed" {
			color = "lightgreen"
		} else if node.Status == "failed" {
			color = "lightcoral"
		} else if node.Status == "running" {
			color = "lightyellow"
		}

		label := fmt.Sprintf("%s\\n%.0fms", node.ID, node.Duration)
		buf.WriteString(fmt.Sprintf("  \"%s\" [label=\"%s\", fillcolor=%s, style=filled];\n",
			node.ID, label, color))
	}

	buf.WriteString("\n")

	// Render edges
	for from, tos := range dr.edges {
		for _, to := range tos {
			buf.WriteString(fmt.Sprintf("  \"%s\" -> \"%s\";\n", from, to))
		}
	}

	buf.WriteString("}\n")
	return buf.String()
}

// RenderHTML generates a simple HTML visualization
func (dr *DAGRenderer) RenderHTML() string {
	var buf strings.Builder

	buf.WriteString("<html><body>\n")
	buf.WriteString("<h1>Workflow DAG</h1>\n")
	buf.WriteString("<table border='1'>\n")
	buf.WriteString("<tr><th>Task ID</th><th>Status</th><th>Duration (ms)</th><th>Critical Path</th></tr>\n")

	for _, node := range dr.nodes {
		critical := "No"
		if node.Critical {
			critical = "Yes"
		}
		buf.WriteString(fmt.Sprintf("<tr><td>%s</td><td>%s</td><td>%.0f</td><td>%s</td></tr>\n",
			node.ID, node.Status, node.Duration, critical))
	}

	buf.WriteString("</table>\n")
	buf.WriteString("</body></html>\n")
	return buf.String()
}

// GetTopologicalSort returns tasks in topological order
func (dr *DAGRenderer) GetTopologicalSort() ([]string, error) {
	// Simple topological sort using DFS
	visited := make(map[string]bool)
	result := make([]string, 0)

	var visit func(string) error
	visit = func(nodeID string) error {
		if visited[nodeID] {
			return nil
		}
		visited[nodeID] = true

		// Visit dependencies first
		for _, dep := range dr.edges[nodeID] {
			if err := visit(dep); err != nil {
				return err
			}
		}

		result = append(result, nodeID)
		return nil
	}

	for nodeID := range dr.nodes {
		if err := visit(nodeID); err != nil {
			return nil, err
		}
	}

	return result, nil
}

// GetParallel returns groups of tasks that can run in parallel
func (dr *DAGRenderer) GetParallel() map[int][]string {
	levels := make(map[int][]string)
	inDegree := make(map[string]int)

	// Calculate in-degree
	for _, node := range dr.nodes {
		inDegree[node.ID] = 0
	}

	for _, tos := range dr.edges {
		for _, to := range tos {
			inDegree[to]++
		}
	}

	// Find nodes by level
	processed := make(map[string]bool)
	level := 0

	for len(processed) < len(dr.nodes) {
		var current []string

		for _, node := range dr.nodes {
			if !processed[node.ID] && inDegree[node.ID] == 0 {
				current = append(current, node.ID)
			}
		}

		if len(current) == 0 {
			break
		}

		levels[level] = current

		// Update in-degrees
		for _, nodeID := range current {
			processed[nodeID] = true
			for _, to := range dr.edges[nodeID] {
				inDegree[to]--
			}
		}

		level++
	}

	return levels
}

// GetStats returns statistics about the DAG
func (dr *DAGRenderer) GetStats() map[string]interface{} {
	totalDuration := 0.0
	maxDuration := 0.0
	criticalCount := 0
	edgeCount := 0

	for _, node := range dr.nodes {
		totalDuration += node.Duration
		if node.Duration > maxDuration {
			maxDuration = node.Duration
		}
		if node.Critical {
			criticalCount++
		}
	}

	// Count total edges
	for _, tos := range dr.edges {
		edgeCount += len(tos)
	}

	return map[string]interface{}{
		"node_count":      len(dr.nodes),
		"edge_count":      edgeCount,
		"total_duration":  totalDuration,
		"max_duration":    maxDuration,
		"critical_count":  criticalCount,
	}
}
