package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os/exec"
	"strings"
)

const (
	hierarchyGraphQueryID       = "hierarchy"
	rustHierarchyLayoutEngineID = "rust_radial_sugiyama"
)

type hierarchyLayoutResult struct {
	Nodes         []Node
	Edges         []Edge
	RouteSegments []RouteSegment
}

type hierarchyLayoutPrepared struct {
	Request         hierarchyLayoutRequest
	NormalizedEdges []Edge
}

type hierarchyLayoutRequest struct {
	RootIRI string                       `json:"root_iri"`
	Nodes   []hierarchyLayoutRequestNode `json:"nodes"`
	Edges   []hierarchyLayoutRequestEdge `json:"edges"`
}

type hierarchyLayoutRequestNode struct {
	NodeID uint32 `json:"node_id"`
	IRI    string `json:"iri"`
}

type hierarchyLayoutRequestEdge struct {
	EdgeIndex    int     `json:"edge_index"`
	ParentID     uint32  `json:"parent_id"`
	ChildID      uint32  `json:"child_id"`
	PredicateIRI *string `json:"predicate_iri,omitempty"`
}

type hierarchyLayoutResponse struct {
	Nodes         []hierarchyLayoutResponseNode         `json:"nodes"`
	RouteSegments []hierarchyLayoutResponseRouteSegment `json:"route_segments"`
}

type hierarchyLayoutResponseNode struct {
	NodeID uint32  `json:"node_id"`
	X      float64 `json:"x"`
	Y      float64 `json:"y"`
	Level  int     `json:"level"`
}

type hierarchyLayoutResponseRouteSegment struct {
	EdgeIndex int                               `json:"edge_index"`
	Kind      string                            `json:"kind"`
	Points    []hierarchyLayoutResponseRoutePoint `json:"points"`
}

type hierarchyLayoutResponseRoutePoint struct {
	X float64 `json:"x"`
	Y float64 `json:"y"`
}

type hierarchyEdgeKey struct {
	ParentID uint32
	ChildID  uint32
}

func shouldUseRustHierarchyLayout(cfg Config, graphQueryID string) bool {
	return cfg.HierarchyLayoutEngine == "rust" && graphQueryID == hierarchyGraphQueryID
}

func prepareHierarchyLayoutRequest(
	rootIRI string,
	nodes []Node,
	edges []Edge,
	preds *PredicateDict,
) hierarchyLayoutPrepared {
	requestNodes := make([]hierarchyLayoutRequestNode, 0, len(nodes))
	for _, node := range nodes {
		requestNodes = append(requestNodes, hierarchyLayoutRequestNode{
			NodeID: node.ID,
			IRI:    node.IRI,
		})
	}

	predicateIRIs := []string(nil)
	if preds != nil {
		predicateIRIs = preds.IRIs()
	}

	seenEdges := make(map[hierarchyEdgeKey]struct{}, len(edges))
	normalizedEdges := make([]Edge, 0, len(edges))
	requestEdges := make([]hierarchyLayoutRequestEdge, 0, len(edges))
	for _, edge := range edges {
		parentID := edge.Target
		childID := edge.Source
		if parentID == childID {
			continue
		}

		key := hierarchyEdgeKey{ParentID: parentID, ChildID: childID}
		if _, ok := seenEdges[key]; ok {
			continue
		}
		seenEdges[key] = struct{}{}

		normalizedEdges = append(normalizedEdges, edge)

		var predicateIRI *string
		if int(edge.PredicateID) >= 0 && int(edge.PredicateID) < len(predicateIRIs) {
			value := predicateIRIs[edge.PredicateID]
			if strings.TrimSpace(value) != "" {
				predicateIRI = &value
			}
		}

		requestEdges = append(requestEdges, hierarchyLayoutRequestEdge{
			EdgeIndex:    len(normalizedEdges) - 1,
			ParentID:     parentID,
			ChildID:      childID,
			PredicateIRI: predicateIRI,
		})
	}

	return hierarchyLayoutPrepared{
		Request: hierarchyLayoutRequest{
			RootIRI: rootIRI,
			Nodes:   requestNodes,
			Edges:   requestEdges,
		},
		NormalizedEdges: normalizedEdges,
	}
}

func applyHierarchyLayoutResponse(
	nodes []Node,
	normalizedEdges []Edge,
	response hierarchyLayoutResponse,
) (hierarchyLayoutResult, error) {
	positionByID := make(map[uint32]hierarchyLayoutResponseNode, len(response.Nodes))
	for _, node := range response.Nodes {
		if _, ok := positionByID[node.NodeID]; ok {
			return hierarchyLayoutResult{}, fmt.Errorf("hierarchy layout bridge returned duplicate node_id %d", node.NodeID)
		}
		positionByID[node.NodeID] = node
	}

	filteredNodes := make([]Node, 0, len(response.Nodes))
	keptNodeIDs := make(map[uint32]struct{}, len(response.Nodes))
	for _, node := range nodes {
		position, ok := positionByID[node.ID]
		if !ok {
			continue
		}
		node.X = position.X
		node.Y = position.Y
		filteredNodes = append(filteredNodes, node)
		keptNodeIDs[node.ID] = struct{}{}
	}
	if len(filteredNodes) != len(response.Nodes) {
		return hierarchyLayoutResult{}, fmt.Errorf("hierarchy layout bridge returned unknown node ids")
	}

	filteredEdges := make([]Edge, 0, len(normalizedEdges))
	normalizedToFilteredEdge := make(map[int]int, len(normalizedEdges))
	for normalizedIndex, edge := range normalizedEdges {
		if _, ok := keptNodeIDs[edge.Source]; !ok {
			continue
		}
		if _, ok := keptNodeIDs[edge.Target]; !ok {
			continue
		}
		normalizedToFilteredEdge[normalizedIndex] = len(filteredEdges)
		filteredEdges = append(filteredEdges, edge)
	}

	routeSegments := make([]RouteSegment, 0, len(response.RouteSegments))
	for _, segment := range response.RouteSegments {
		filteredEdgeIndex, ok := normalizedToFilteredEdge[segment.EdgeIndex]
		if !ok {
			return hierarchyLayoutResult{}, fmt.Errorf("hierarchy layout bridge returned route for unknown edge_index %d", segment.EdgeIndex)
		}
		points := make([]RoutePoint, 0, len(segment.Points))
		for _, point := range segment.Points {
			points = append(points, RoutePoint{
				X: point.X,
				Y: point.Y,
			})
		}
		routeSegments = append(routeSegments, RouteSegment{
			EdgeIndex: filteredEdgeIndex,
			Kind:      segment.Kind,
			Points:    points,
		})
	}

	return hierarchyLayoutResult{
		Nodes:         filteredNodes,
		Edges:         filteredEdges,
		RouteSegments: routeSegments,
	}, nil
}

func runHierarchyLayoutBridge(
	ctx context.Context,
	cfg Config,
	request hierarchyLayoutRequest,
) (hierarchyLayoutResponse, error) {
	input, err := json.Marshal(request)
	if err != nil {
		return hierarchyLayoutResponse{}, fmt.Errorf("marshal hierarchy layout request failed: %w", err)
	}

	bridgeCtx, cancel := context.WithTimeout(ctx, cfg.HierarchyLayoutTimeout)
	defer cancel()

	cmd := exec.CommandContext(bridgeCtx, cfg.HierarchyLayoutBridgeBin)
	cmd.Dir = cfg.HierarchyLayoutBridgeWorkdir
	cmd.Stdin = bytes.NewReader(input)

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		if bridgeCtx.Err() != nil {
			return hierarchyLayoutResponse{}, fmt.Errorf("hierarchy layout bridge timed out after %s", cfg.HierarchyLayoutTimeout)
		}
		detail := strings.TrimSpace(stderr.String())
		if detail == "" {
			detail = err.Error()
		}
		return hierarchyLayoutResponse{}, fmt.Errorf("hierarchy layout bridge failed: %s", detail)
	}

	var response hierarchyLayoutResponse
	if err := json.Unmarshal(stdout.Bytes(), &response); err != nil {
		detail := strings.TrimSpace(stderr.String())
		if detail != "" {
			return hierarchyLayoutResponse{}, fmt.Errorf("parse hierarchy layout bridge response failed: %v (stderr: %s)", err, detail)
		}
		return hierarchyLayoutResponse{}, fmt.Errorf("parse hierarchy layout bridge response failed: %w", err)
	}

	return response, nil
}

func layoutHierarchyWithRust(
	ctx context.Context,
	cfg Config,
	nodes []Node,
	edges []Edge,
	preds *PredicateDict,
) (hierarchyLayoutResult, error) {
	prepared := prepareHierarchyLayoutRequest(cfg.HierarchyLayoutRootIRI, nodes, edges, preds)
	response, err := runHierarchyLayoutBridge(ctx, cfg, prepared.Request)
	if err != nil {
		return hierarchyLayoutResult{}, err
	}
	return applyHierarchyLayoutResponse(nodes, prepared.NormalizedEdges, response)
}