package dsp

import "fmt"

// ValType is the data type of a Sample: a scalar float, a float array
// (waveform), or — at graph-compile time, before a source's real type is
// known — unknown.
type ValType uint8

const (
	ValUnknown ValType = iota
	ValScalar
	ValArray
)

func (t ValType) String() string {
	switch t {
	case ValScalar:
		return "scalar"
	case ValArray:
		return "array"
	default:
		return "unknown"
	}
}

// Sample is a value flowing through the synthetic DSP graph: either a scalar
// float64 or a float64 array (waveform). It is the array-aware counterpart of
// the bare float64 the legacy scalar Node interface uses.
type Sample struct {
	F       float64
	Arr     []float64
	IsArray bool
}

// Scalar wraps a float64 as a scalar Sample.
func Scalar(f float64) Sample { return Sample{F: f} }

// Array wraps a []float64 as an array Sample.
func Array(a []float64) Sample { return Sample{Arr: a, IsArray: true} }

// Type reports whether the sample is a scalar or an array.
func (s Sample) Type() ValType {
	if s.IsArray {
		return ValArray
	}
	return ValScalar
}

// AsAny returns the value in the form datasource.Value.Data expects: a
// []float64 for arrays, a float64 for scalars.
func (s Sample) AsAny() any {
	if s.IsArray {
		return s.Arr
	}
	return s.F
}

// AsArray returns the sample's data as a slice: the array itself, or a
// single-element slice for a scalar. Used by reductions that accept either.
func (s Sample) AsArray() []float64 {
	if s.IsArray {
		return s.Arr
	}
	return []float64{s.F}
}

// ArrayNode is an optional extension of Node implemented by ops that operate
// natively on Samples (reductions array→scalar, producers array→array, etc.).
// eval prefers ProcessSample when a node implements it.
type ArrayNode interface {
	Node
	ProcessSample(inputs []Sample, state map[string]any) (Sample, error)
}

// statelessElementwise lists scalar ops that are safe to broadcast element-wise
// over array inputs: they hold no per-evaluation state, so running the legacy
// Process once per array lane is well-defined. Stateful ops (moving_average,
// rms, lowpass, derivative, integrate) and lua are excluded — a single shared
// state map cannot be meaningfully split across lanes.
var statelessElementwise = map[string]bool{
	"gain": true, "offset": true, "add": true, "subtract": true,
	"multiply": true, "divide": true, "clamp": true, "threshold": true,
	"expr": true,
}

// StatelessElementwise reports whether a scalar op type may be broadcast over
// array inputs via ApplyElementwise.
func StatelessElementwise(nodeType string) bool { return statelessElementwise[nodeType] }

// scalarInputs wraps a legacy float64 input slice as scalar Samples.
func scalarInputs(in []float64) []Sample {
	out := make([]Sample, len(in))
	for i, v := range in {
		out[i] = Scalar(v)
	}
	return out
}

// ApplyElementwise runs a stateless scalar Node over Sample inputs. If every
// input is scalar it calls Process once and wraps the result. If any input is
// an array it broadcasts: scalar inputs act as constants, all array inputs must
// share a common length (else an error), and Process is invoked once per index.
//
// The node MUST be stateless (see StatelessElementwise) — a shared state map
// cannot be split across array lanes.
func ApplyElementwise(n Node, inputs []Sample, state map[string]any) (Sample, error) {
	// Determine the array length, if any input is an array.
	length := -1
	for _, s := range inputs {
		if !s.IsArray {
			continue
		}
		if length == -1 {
			length = len(s.Arr)
		} else if len(s.Arr) != length {
			return Sample{}, fmt.Errorf("%s: array length mismatch (%d vs %d)", n.Type(), length, len(s.Arr))
		}
	}

	if length == -1 {
		// All scalar — single legacy call.
		row := make([]float64, len(inputs))
		for i, s := range inputs {
			row[i] = s.F
		}
		r, err := n.Process(row, state)
		if err != nil {
			return Sample{}, err
		}
		return Scalar(r), nil
	}

	out := make([]float64, length)
	row := make([]float64, len(inputs))
	for i := 0; i < length; i++ {
		for j, s := range inputs {
			if s.IsArray {
				row[j] = s.Arr[i]
			} else {
				row[j] = s.F
			}
		}
		r, err := n.Process(row, state)
		if err != nil {
			return Sample{}, err
		}
		out[i] = r
	}
	return Array(out), nil
}