package executor

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"time"

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

// WasmMagic is the magic number that every valid wasm binary starts with.
var WasmMagic = []byte{0x00, 0x61, 0x73, 0x6d}

// Config controls the sandbox limits for wasm execution.
type Config struct {
	// MaxExecTime is the wall-clock timeout for a single execution.
	MaxExecTime time.Duration

	// MaxCompileTime is the wall-clock timeout for compiling a module.
	MaxCompileTime time.Duration

	// MaxMemoryMB is the maximum memory in megabytes a module may use.
	MaxMemoryMB uint32

	// MaxWasmSize is the maximum size in bytes of a wasm binary.
	MaxWasmSize int64

	// MaxOutputBytes is the maximum bytes captured from stdout/stderr each.
	MaxOutputBytes int64
}

// DefaultConfig returns sensible defaults for running untrusted wasm.
func DefaultConfig() Config {
	return Config{
		MaxExecTime:    5 * time.Second,
		MaxCompileTime: 10 * time.Second,
		MaxMemoryMB:    16,
		MaxWasmSize:    10 * 1024 * 1024, // 10 MB
		MaxOutputBytes: 10_000,
	}
}

// Result holds the output of a wasm execution.
type Result struct {
	Status     string
	Stdout     string
	Stderr     string
	DurationMs int64
	ExitCode   uint32
}

const (
	StatusSuccess = "success"
	StatusError   = "error"
	StatusTimeout = "timeout"
)

// Validate checks that the wasm binary looks valid and is within size limits.
func Validate(wasm []byte, maxSize int64) error {
	if int64(len(wasm)) > maxSize {
		return fmt.Errorf("wasm binary too large: %d bytes (max %d)", len(wasm), maxSize)
	}
	if len(wasm) < 8 {
		return errors.New("wasm binary too small to be valid")
	}
	if !bytes.HasPrefix(wasm, WasmMagic) {
		return errors.New("invalid wasm magic number")
	}
	return nil
}

// Execute runs a wasm binary in a sandboxed wazero runtime and returns the result.
// The provided context can be used for cancellation independent of the execution timeout.
//
// Sandbox guarantees:
//   - No filesystem access (all path operations return ENOSYS)
//   - No network access (no host functions exposed)
//   - No environment variables
//   - No command-line arguments
//   - No real clock (walltime is fake, starting at 2022-01-01)
//   - No real sleep (nanosleep is a no-op)
//   - Deterministic random source (not crypto-secure)
//   - Stdin returns immediate EOF
//   - Stdout/stderr captured into bounded buffers
//   - Memory hard-capped at MaxMemoryMB
//   - Wall-clock execution timeout via context cancellation
//   - Separate compilation timeout
//   - Fresh runtime per execution (no state leaks between runs)
func Execute(ctx context.Context, cfg Config, wasm []byte) (result *Result, err error) {
	// Recover from any panic in wazero to prevent crashing the daemon.
	defer func() {
		if r := recover(); r != nil {
			result = &Result{
				Status: StatusError,
				Stderr: fmt.Sprintf("runtime panic: %v", r),
			}
			err = nil
		}
	}()

	if err := Validate(wasm, cfg.MaxWasmSize); err != nil {
		return nil, err
	}

	memPages := cfg.MaxMemoryMB * 16 // 1 MB = 16 pages (each page is 64 KB)

	rtCfg := wazero.NewRuntimeConfig().
		WithMemoryLimitPages(memPages).
		WithCloseOnContextDone(true).
		// Disable debug info to reduce compilation time and memory for
		// untrusted modules. We don't need DWARF stack traces.
		WithDebugInfoEnabled(false)

	rt := wazero.NewRuntimeWithConfig(ctx, rtCfg)
	defer func() { _ = rt.Close(ctx) }()

	if _, err := wasi_snapshot_preview1.Instantiate(ctx, rt); err != nil {
		return nil, fmt.Errorf("instantiating wasi: %w", err)
	}

	// Compile with a separate timeout so a malicious module can't hang
	// the daemon during compilation.
	compileCtx, compileCancel := context.WithTimeout(ctx, cfg.MaxCompileTime)
	defer compileCancel()

	compiled, err := rt.CompileModule(compileCtx, wasm)
	if err != nil {
		return nil, fmt.Errorf("compiling module: %w", err)
	}

	stdout := &limitedWriter{max: cfg.MaxOutputBytes}
	stderr := &limitedWriter{max: cfg.MaxOutputBytes}

	// Explicitly configure every capability to be denied. While wazero
	// defaults are already sandboxed, being explicit provides defense in
	// depth — if a future wazero version changes defaults, we're still safe.
	modCfg := wazero.NewModuleConfig().
		WithName("guest").
		// I/O: stdin is EOF, stdout/stderr captured into bounded buffers.
		WithStdin(nil).
		WithStdout(stdout).
		WithStderr(stderr).
		// No filesystem access. This is the default, but be explicit.
		WithFSConfig(wazero.NewFSConfig()).
		// No environment variables.
		// No command-line arguments.
		// (wazero defaults to empty for both, but calling WithArgs/WithEnv
		// with values would add them — we simply don't.)
		//
		// Clock: use wazero's fake deterministic clocks (default).
		// - Walltime starts at 2022-01-01T00:00:00Z, advances 1ms per read.
		// - Nanotime starts at 0, advances 1ms per read.
		// - Nanosleep is a no-op (returns immediately).
		// - Random source is deterministic (seed 42).
		// We do NOT call WithWalltime, WithNanotime, WithSysNanosleep, or
		// WithRandSource, which would install the real host implementations.
		//
		// Start function: call _start (WASI entry point).
		WithStartFunctions("_start")

	execCtx, execCancel := context.WithTimeout(ctx, cfg.MaxExecTime)
	defer execCancel()

	start := time.Now()
	_, err = rt.InstantiateModule(execCtx, compiled, modCfg)
	duration := time.Since(start)

	result = &Result{
		Stdout:     stdout.String(),
		Stderr:     stderr.String(),
		DurationMs: duration.Milliseconds(),
	}

	if err != nil {
		var exitErr *sys.ExitError
		if errors.As(err, &exitErr) {
			result.ExitCode = exitErr.ExitCode()
			switch result.ExitCode {
			case 0:
				result.Status = StatusSuccess
			case sys.ExitCodeDeadlineExceeded:
				result.Status = StatusTimeout
			case sys.ExitCodeContextCanceled:
				result.Status = StatusTimeout
			default:
				result.Status = StatusError
			}
			return result, nil
		}
		return nil, fmt.Errorf("instantiating module: %w", err)
	}

	result.Status = StatusSuccess
	return result, nil
}

// limitedWriter captures output up to a maximum number of bytes.
type limitedWriter struct {
	buf bytes.Buffer
	max int64
}

func (w *limitedWriter) Write(p []byte) (int, error) {
	remaining := w.max - int64(w.buf.Len())
	if remaining <= 0 {
		// Silently discard but report success to the guest so it doesn't error.
		return len(p), nil
	}
	if int64(len(p)) > remaining {
		p = p[:remaining]
	}
	return w.buf.Write(p)
}

func (w *limitedWriter) String() string {
	return w.buf.String()
}

// WriteTo implements io.WriterTo for efficient copying.
func (w *limitedWriter) WriteTo(dst io.Writer) (int64, error) {
	return w.buf.WriteTo(dst)
}