package organ

import (
	"context"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
)

// Registry manages loaded organs for one lattice instance.
// It owns the wazero runtime and handles module lifecycle.
type Registry struct {
	mu      sync.RWMutex
	organs  map[OrganID]*loadedOrgan
	byType  map[OrganType][]*loadedOrgan
	runtime wazero.Runtime
	ctx     context.Context
}

// loadedOrgan extends Organ with the wazero module instance.
type loadedOrgan struct {
	Organ
	compiled wazero.CompiledModule
	instance api.Module
}

// NewRegistry creates a registry with a fresh wazero runtime.
func NewRegistry(ctx context.Context) (*Registry, error) {
	r := wazero.NewRuntime(ctx)

	// Instantiate WASI for organs that use it.
	if _, err := wasi_snapshot_preview1.Instantiate(ctx, r); err != nil {
		r.Close(ctx)
		return nil, fmt.Errorf("wasi init: %w", err)
	}

	return &Registry{
		organs:  make(map[OrganID]*loadedOrgan),
		byType:  make(map[OrganType][]*loadedOrgan),
		runtime: r,
		ctx:     ctx,
	}, nil
}

// Load compiles and instantiates a WASM module.
// The organ must export: alloc, dealloc, organ_type, organ_version, process.
func (reg *Registry) Load(wasmBytes []byte, manifest Manifest) (*Organ, error) {
	reg.mu.Lock()
	defer reg.mu.Unlock()

	id := ComputeID(wasmBytes)
	manifest.ID = id
	manifest.Size = len(wasmBytes)

	// Already loaded?
	if existing, ok := reg.organs[id]; ok {
		return &existing.Organ, nil
	}

	// Compile.
	compiled, err := reg.runtime.CompileModule(reg.ctx, wasmBytes)
	if err != nil {
		return nil, fmt.Errorf("compile: %w", err)
	}

	// Instantiate as a WASI reactor (calls _initialize, not _start).
	config := wazero.NewModuleConfig().
		WithName(fmt.Sprintf("organ-%s", id)).
		WithStartFunctions("_initialize")

	instance, err := reg.runtime.InstantiateModule(reg.ctx, compiled, config)
	if err != nil {
		return nil, fmt.Errorf("instantiate: %w", err)
	}

	// Verify required exports exist.
	for _, name := range []string{"alloc", "dealloc", "process"} {
		if instance.ExportedFunction(name) == nil {
			instance.Close(reg.ctx)
			return nil, fmt.Errorf("organ missing required export: %s", name)
		}
	}

	lo := &loadedOrgan{
		Organ: Organ{
			Manifest:  manifest,
			WasmBytes: wasmBytes,
			Status:    StatusLoaded,
			Loaded:    time.Now(),
		},
		compiled: compiled,
		instance: instance,
	}

	reg.organs[id] = lo
	reg.byType[manifest.Type] = append(reg.byType[manifest.Type], lo)

	return &lo.Organ, nil
}

// Call invokes an organ's process function with the given opcode and input.
// The opcode byte is prepended to the input before passing to the WASM module.
func (reg *Registry) Call(id OrganID, opcode byte, input []byte) ([]byte, error) {
	reg.mu.RLock()
	lo, ok := reg.organs[id]
	reg.mu.RUnlock()

	if !ok {
		return nil, fmt.Errorf("organ not found: %s", id)
	}
	if lo.Status != StatusLoaded {
		return nil, fmt.Errorf("organ not loaded: %s (status %d)", id, lo.Status)
	}

	mod := lo.instance

	// Prepend opcode.
	payload := make([]byte, 1+len(input))
	payload[0] = opcode
	copy(payload[1:], input)

	// Allocate in guest memory.
	allocFn := mod.ExportedFunction("alloc")
	results, err := allocFn.Call(reg.ctx, uint64(len(payload)))
	if err != nil {
		return nil, fmt.Errorf("alloc: %w", err)
	}
	ptr := uint32(results[0])

	// Write input to guest memory.
	if !mod.Memory().Write(ptr, payload) {
		return nil, fmt.Errorf("memory write failed at ptr=%d len=%d", ptr, len(payload))
	}

	// Call process.
	processFn := mod.ExportedFunction("process")
	results, err = processFn.Call(reg.ctx, uint64(ptr), uint64(len(payload)))
	if err != nil {
		return nil, fmt.Errorf("process: %w", err)
	}

	// Unpack result: high 32 bits = pointer, low 32 bits = length.
	packed := results[0]
	outPtr := uint32(packed >> 32)
	outLen := uint32(packed & 0xFFFFFFFF)

	// Read output from guest memory.
	output, ok := mod.Memory().Read(outPtr, outLen)
	if !ok {
		return nil, fmt.Errorf("memory read failed at ptr=%d len=%d", outPtr, outLen)
	}

	// Make a copy — wazero memory may be invalidated.
	result := make([]byte, len(output))
	copy(result, output)

	// Free both buffers in guest.
	deallocFn := mod.ExportedFunction("dealloc")
	deallocFn.Call(reg.ctx, uint64(ptr), uint64(len(payload)))
	deallocFn.Call(reg.ctx, uint64(outPtr), uint64(outLen))

	lo.CallCount.Add(1)
	lo.LastCall = time.Now()

	return result, nil
}

// Evict destroys an organ: closes the wazero module, releases memory,
// removes from registry.
func (reg *Registry) Evict(id OrganID) error {
	reg.mu.Lock()
	defer reg.mu.Unlock()

	lo, ok := reg.organs[id]
	if !ok {
		return fmt.Errorf("organ not found: %s", id)
	}

	// Close the WASM instance.
	if lo.instance != nil {
		lo.instance.Close(reg.ctx)
	}

	lo.Status = StatusEvicted
	lo.WasmBytes = nil // release for GC

	// Remove from type index.
	organs := reg.byType[lo.Manifest.Type]
	for i, o := range organs {
		if o.Manifest.ID == id {
			reg.byType[lo.Manifest.Type] = append(organs[:i], organs[i+1:]...)
			break
		}
	}

	delete(reg.organs, id)
	return nil
}

// Best returns the highest-fitness organ of the given type, or nil.
func (reg *Registry) Best(t OrganType) *Organ {
	reg.mu.RLock()
	defer reg.mu.RUnlock()

	organs := reg.byType[t]
	if len(organs) == 0 {
		return nil
	}

	var best *loadedOrgan
	for _, o := range organs {
		if o.Status != StatusLoaded {
			continue
		}
		if best == nil || best.Manifest.Fitness.Less(o.Manifest.Fitness) {
			best = o
		}
	}

	if best == nil {
		return nil
	}
	return &best.Organ
}

// All returns all loaded organs of the given type, sorted by fitness descending.
func (reg *Registry) All(t OrganType) []*Organ {
	reg.mu.RLock()
	defer reg.mu.RUnlock()

	organs := reg.byType[t]
	result := make([]*Organ, 0, len(organs))
	for _, o := range organs {
		if o.Status == StatusLoaded {
			result = append(result, &o.Organ)
		}
	}

	sort.Slice(result, func(i, j int) bool {
		return result[j].Manifest.Fitness.Less(result[i].Manifest.Fitness)
	})

	return result
}

// Get returns an organ by ID, or nil if not found.
func (reg *Registry) Get(id OrganID) *Organ {
	reg.mu.RLock()
	defer reg.mu.RUnlock()

	if lo, ok := reg.organs[id]; ok {
		return &lo.Organ
	}
	return nil
}

// Close shuts down the registry and all loaded organs.
func (reg *Registry) Close() error {
	reg.mu.Lock()
	defer reg.mu.Unlock()

	// wazero runtime close destroys all modules.
	return reg.runtime.Close(reg.ctx)
}