package synthetic

import (
	"errors"
	"fmt"

	"github.com/uopi/uopi/internal/broker"
	"github.com/uopi/uopi/internal/dsp"
)

// runtimeGraph is the executable form of a synthetic signal's DAG. Nodes are
// held in topological order so a single forward pass computes every value with
// each node's inputs already resolved. Op-node state maps persist across
// evaluations (for stateful nodes like moving_average / lua).
type runtimeGraph struct {
	order    []*rtNode    // topological order (sources first, output last)
	sources  []rtSource   // source nodes, in topological order
	outputID string       // id of the output node
	outType  dsp.ValType  // best-effort output type (scalar/array/unknown)
}

type rtNode struct {
	id     string
	kind   string         // source | op | output
	op     dsp.Node       // set for kind==op
	state  map[string]any // persistent per-node state (op only)
	inputs []string       // upstream node ids, in input order
}

type rtSource struct {
	id  string
	ref broker.SignalRef
}

// sourceRefs returns the broker references for every source node, in a stable
// order matching rg.sources.
func (rg *runtimeGraph) sourceRefs() []broker.SignalRef {
	refs := make([]broker.SignalRef, len(rg.sources))
	for i, s := range rg.sources {
		refs[i] = s.ref
	}
	return refs
}

// evalSample computes the output Sample (scalar or array) given the latest
// value for each source node (keyed by source node id). Nodes are visited in
// topological order so every input is present by the time a node is processed.
//
// Op dispatch:
//   - ArrayNode ops (reductions/producers) run natively on Samples.
//   - stateless elementwise ops broadcast over array inputs.
//   - stateful ops (filters) and lua are scalar-only; an array input errors,
//     since their per-evaluation state cannot be split across array lanes.
func (rg *runtimeGraph) evalSample(sourceVals map[string]dsp.Sample) (dsp.Sample, error) {
	vals, err := rg.evalSampleTrace(sourceVals)
	if err != nil {
		return dsp.Sample{}, err
	}
	return vals[rg.outputID], nil
}

// evalSampleTrace runs a full forward pass like evalSample but returns the value
// computed for *every* node (sources, ops and the output), keyed by node id. It
// is used by the editor's live/debug trace to show each node's current value.
// On an op error it returns the values computed so far together with the error,
// so partial results can still be displayed.
func (rg *runtimeGraph) evalSampleTrace(sourceVals map[string]dsp.Sample) (map[string]dsp.Sample, error) {
	vals := make(map[string]dsp.Sample, len(rg.order)+len(sourceVals))
	for id, v := range sourceVals {
		vals[id] = v
	}
	for _, n := range rg.order {
		switch n.kind {
		case "op":
			in := make([]dsp.Sample, len(n.inputs))
			for i, id := range n.inputs {
				in[i] = vals[id]
			}
			r, err := evalOp(n, in)
			if err != nil {
				return vals, fmt.Errorf("node %s (%s): %w", n.id, n.op.Type(), err)
			}
			vals[n.id] = r
		case "output":
			if len(n.inputs) > 0 {
				vals[n.id] = vals[n.inputs[0]]
			}
		}
	}
	return vals, nil
}

// evalOp runs a single op node over its Sample inputs, choosing the right
// execution path for the node type.
func evalOp(n *rtNode, in []dsp.Sample) (dsp.Sample, error) {
	if an, ok := n.op.(dsp.ArrayNode); ok {
		return an.ProcessSample(in, n.state)
	}
	if dsp.StatelessElementwise(n.op.Type()) {
		return dsp.ApplyElementwise(n.op, in, n.state)
	}
	// Stateful / lua: scalar-only.
	row := make([]float64, len(in))
	for i, s := range in {
		if s.IsArray {
			return dsp.Sample{}, fmt.Errorf("does not accept an array input")
		}
		row[i] = s.F
	}
	r, err := n.op.Process(row, n.state)
	if err != nil {
		return dsp.Sample{}, err
	}
	return dsp.Scalar(r), nil
}

// eval is the scalar wrapper around evalSample, kept so callers and tests that
// deal purely in float64 (legacy linear graphs, scalar sources) are unchanged.
func (rg *runtimeGraph) eval(sourceVals map[string]float64) (float64, error) {
	sv := make(map[string]dsp.Sample, len(sourceVals))
	for id, v := range sourceVals {
		sv[id] = dsp.Scalar(v)
	}
	out, err := rg.evalSample(sv)
	if err != nil {
		return 0, err
	}
	return out.F, nil
}

// compileGraph converts a SignalDef into an executable runtimeGraph. When the
// def carries an explicit Graph it is used directly; otherwise the legacy
// Inputs+Pipeline form is converted to an equivalent linear graph (see toGraph).
func compileGraph(def SignalDef) (*runtimeGraph, error) {
	g := toGraph(def)
	if g == nil || len(g.Nodes) == 0 {
		return &runtimeGraph{}, nil
	}
	order, err := topoOrder(g)
	if err != nil {
		return nil, err
	}
	rg := &runtimeGraph{outputID: g.Output}
	// nodeType tracks each node's best-effort output type for static
	// propagation. Sources are unknown at compile time (their real type is
	// only known once data flows), so type errors here are advisory; runtime
	// Sample typing is authoritative.
	nodeType := make(map[string]dsp.ValType, len(order))
	for _, gn := range order {
		switch gn.Kind {
		case "source":
			rg.sources = append(rg.sources, rtSource{id: gn.ID, ref: broker.SignalRef{DS: gn.DS, Name: gn.Signal}})
			nodeType[gn.ID] = dsp.ValUnknown
		case "op":
			node, err := buildNode(NodeDef{Type: gn.Op, Params: gn.Params})
			if err != nil {
				return nil, fmt.Errorf("node %q: %w", gn.ID, err)
			}
			inTypes := make([]dsp.ValType, len(gn.Inputs))
			for i, id := range gn.Inputs {
				inTypes[i] = nodeType[id]
			}
			ot, terr := dsp.OpOutputType(gn.Op, inTypes)
			if terr != nil {
				return nil, fmt.Errorf("node %q: %w", gn.ID, terr)
			}
			nodeType[gn.ID] = ot
			rg.order = append(rg.order, &rtNode{id: gn.ID, kind: "op", op: node, state: map[string]any{}, inputs: gn.Inputs})
		case "output":
			rg.outputID = gn.ID
			if len(gn.Inputs) > 0 {
				nodeType[gn.ID] = nodeType[gn.Inputs[0]]
			}
			rg.order = append(rg.order, &rtNode{id: gn.ID, kind: "output", inputs: gn.Inputs})
		default:
			return nil, fmt.Errorf("node %q: unknown kind %q", gn.ID, gn.Kind)
		}
	}
	rg.outType = nodeType[rg.outputID]
	return rg, nil
}

// topoOrder returns the graph's nodes in a topological (dependency-first) order,
// treating each node's Inputs as its predecessors. It errors on dangling input
// references or cycles.
func topoOrder(g *Graph) ([]GraphNode, error) {
	byID := make(map[string]GraphNode, len(g.Nodes))
	for _, n := range g.Nodes {
		byID[n.ID] = n
	}
	indeg := make(map[string]int, len(g.Nodes))
	succ := make(map[string][]string, len(g.Nodes))
	for _, n := range g.Nodes {
		if _, ok := indeg[n.ID]; !ok {
			indeg[n.ID] = 0
		}
		for _, in := range n.Inputs {
			if _, ok := byID[in]; !ok {
				return nil, fmt.Errorf("node %q references unknown input %q", n.ID, in)
			}
			indeg[n.ID]++
			succ[in] = append(succ[in], n.ID)
		}
	}
	// Seed the queue with roots, preserving the node slice order for determinism.
	queue := make([]string, 0, len(g.Nodes))
	for _, n := range g.Nodes {
		if indeg[n.ID] == 0 {
			queue = append(queue, n.ID)
		}
	}
	order := make([]GraphNode, 0, len(g.Nodes))
	for len(queue) > 0 {
		id := queue[0]
		queue = queue[1:]
		order = append(order, byID[id])
		for _, s := range succ[id] {
			indeg[s]--
			if indeg[s] == 0 {
				queue = append(queue, s)
			}
		}
	}
	if len(order) != len(g.Nodes) {
		return nil, errors.New("graph contains a cycle")
	}
	return order, nil
}

// toGraph returns the DAG for a SignalDef. If def.Graph is set it is returned
// as-is. Otherwise the legacy linear form is converted: each input signal
// becomes a source node, the pipeline becomes a chain of op nodes (the first op
// receiving every source, each later op the previous op's output), terminated by
// an output node. With no pipeline the output takes the first source directly,
// matching the old runPipeline behaviour.
func toGraph(def SignalDef) *Graph {
	if def.Graph != nil && len(def.Graph.Nodes) > 0 {
		return def.Graph
	}

	inputs := def.Inputs
	if len(inputs) == 0 && def.DS != "" && def.Signal != "" {
		inputs = []InputRef{{DS: def.DS, Signal: def.Signal}}
	}

	nodes := make([]GraphNode, 0, len(inputs)+len(def.Pipeline)+1)
	srcIDs := make([]string, 0, len(inputs))
	for i, inp := range inputs {
		id := fmt.Sprintf("s%d", i)
		nodes = append(nodes, GraphNode{ID: id, Kind: "source", DS: inp.DS, Signal: inp.Signal})
		srcIDs = append(srcIDs, id)
	}

	opIDs := make([]string, 0, len(def.Pipeline))
	for i, nd := range def.Pipeline {
		id := fmt.Sprintf("p%d", i)
		var ins []string
		if i == 0 {
			ins = srcIDs
		} else {
			ins = []string{opIDs[i-1]}
		}
		nodes = append(nodes, GraphNode{ID: id, Kind: "op", Op: nd.Type, Params: nd.Params, Inputs: ins})
		opIDs = append(opIDs, id)
	}

	var outInputs []string
	if len(opIDs) > 0 {
		outInputs = []string{opIDs[len(opIDs)-1]}
	} else if len(srcIDs) > 0 {
		outInputs = []string{srcIDs[0]}
	}
	nodes = append(nodes, GraphNode{ID: "out", Kind: "output", Inputs: outInputs})

	return &Graph{Nodes: nodes, Output: "out"}
}