package policy

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"time"

	"go.starlark.net/starlark"
)

// ExecModule provides the ability to execute external binary commands from Starlark policies.
type ExecModule struct{}

func (m *ExecModule) String() string       { return "<module exec>" }
func (m *ExecModule) Type() string         { return "module" }
func (m *ExecModule) Freeze()              {}
func (m *ExecModule) Truth() starlark.Bool { return true }
func (m *ExecModule) Hash() (uint32, error) {
	return 0, errors.New("unhashable")
}

func (m *ExecModule) Attr(name string) (starlark.Value, error) {
	switch name {
	case "pipe_to_process":
		return starlark.NewBuiltin("pipe_to_process", m.pipeToProcess), nil
	}
	return nil, nil
}

func (m *ExecModule) AttrNames() []string {
	return []string{"pipe_to_process"}
}

// pipeToProcess executes an external binary, piping input_data to stdin,
// and returns an ExecResult with code, stdout, stderr, and timed_out fields.
func (m *ExecModule) pipeToProcess(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
	var command string
	var inputData starlark.Value = starlark.None
	var argList starlark.Value = starlark.None
	var timeoutVal starlark.Value = starlark.None
	var envDict starlark.Value = starlark.None

	if err := starlark.UnpackArgs(b.Name(), args, kwargs,
		"command", &command,
		"input_data?", &inputData,
		"args?", &argList,
		"timeout?", &timeoutVal,
		"environment?", &envDict,
	); err != nil {
		return nil, err
	}

	var inputBytes []byte
	switch v := inputData.(type) {
	case starlark.String:
		inputBytes = []byte(v)
	case starlark.Bytes:
		inputBytes = []byte(v)
	case starlark.NoneType:
		inputBytes = nil
	default:
		return nil, fmt.Errorf("input_data must be a string or bytes, got %s", inputData.Type())
	}

	var goArgs []string
	if argList != starlark.None {
		lst, ok := argList.(*starlark.List)
		if !ok {
			return nil, fmt.Errorf("args must be a list, got %s", argList.Type())
		}
		for it := range lst.Elements() {
			s, ok := it.(starlark.String)
			if !ok {
				return nil, fmt.Errorf("args must be a list of strings, got %s", it.Type())
			}
			goArgs = append(goArgs, string(s))
		}
	}

	var timeoutDur = 30 * time.Second
	if timeoutVal != starlark.None {
		switch v := timeoutVal.(type) {
		case starlark.Int:
			if i, ok := v.Int64(); ok {
				timeoutDur = time.Duration(i) * time.Second
			} else {
				return nil, errors.New("timeout must be a valid integer")
			}
		case starlark.Float:
			f, _ := starlark.AsFloat(v)
			timeoutDur = time.Duration(f) * time.Second
		default:
			return nil, fmt.Errorf("timeout must be a number, got %s", timeoutVal.Type())
		}
	}

	hasPath := false
	var env []string
	if envDict != starlark.None {
		dict, ok := envDict.(*starlark.Dict)
		if !ok {
			return nil, fmt.Errorf("environment must be a dict, got %s", envDict.Type())
		}
		for _, item := range dict.Items() {
			key, ok := item[0].(starlark.String)
			if !ok {
				return nil, fmt.Errorf("environment key must be a string, got %s", item[0].Type())
			}
			val, ok := item[1].(starlark.String)
			if !ok {
				return nil, fmt.Errorf("environment value must be a string, got %s", item[1].Type())
			}
			env = append(env, string(key)+"="+string(val))
			if key == "PATH" && len(val) > 0 {
				hasPath = true
			}
		}
	}

	if !hasPath {
		if path := os.Getenv("PATH"); path != "" {
			env = append(env, "PATH="+path)
		}
	}

	ctx, ok := thread.Local("context").(context.Context)
	if !ok || ctx == nil {
		ctx = context.Background()
	}
	ctx, cancel := context.WithTimeout(ctx, timeoutDur)
	defer cancel()

	cmd := exec.CommandContext(ctx, command, goArgs...)
	cmd.Env = env
	cmd.Stdin = bytes.NewReader(inputBytes)

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

	runErr := cmd.Run()

	timedOut := errors.Is(ctx.Err(), context.DeadlineExceeded)

	code := 0
	if runErr != nil {
		if timedOut {
			code = -1
		} else if exitErr, ok := errors.AsType[*exec.ExitError](runErr); ok {
			code = exitErr.ExitCode()
		} else {
			return nil, fmt.Errorf("exec failed: %w", runErr)
		}
	}

	return &ExecResult{
		code:     code,
		stdout:   stdout.String(),
		stderr:   stderr.String(),
		timedOut: timedOut,
	}, nil
}

// ExecResult represents the result of an external process execution.
type ExecResult struct {
	code     int
	stdout   string
	stderr   string
	timedOut bool
}

func (r *ExecResult) String() string {
	return fmt.Sprintf("<ExecResult code=%d>", r.code)
}
func (r *ExecResult) Type() string { return "ExecResult" }
func (r *ExecResult) Freeze()      {}
func (r *ExecResult) Truth() starlark.Bool {
	return starlark.Bool(r.code == 0 && !r.timedOut)
}
func (r *ExecResult) Hash() (uint32, error) {
	return 0, errors.New("unhashable")
}

func (r *ExecResult) Attr(name string) (starlark.Value, error) {
	switch name {
	case "code":
		return starlark.MakeInt(r.code), nil
	case "stdout":
		return starlark.String(r.stdout), nil
	case "stderr":
		return starlark.String(r.stderr), nil
	case "timed_out":
		return starlark.Bool(r.timedOut), nil
	}
	return nil, nil
}

func (r *ExecResult) AttrNames() []string {
	return []string{"code", "stdout", "stderr", "timed_out"}
}