package agentworkspace

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"sync"

	"iop/packages/go/agenttask"
)

// ErrConfinementUnavailable means this host cannot install the filesystem
// policy required for an unattended overlay child. Callers must fail closed.
var ErrConfinementUnavailable = errors.New(
	"agentworkspace: executable filesystem confinement is unavailable",
)

const confinementPolicySchemaVersion uint32 = 1

type confinementPolicy struct {
	SchemaVersion uint32   `json:"schema_version"`
	Revision      string   `json:"revision"`
	Backend       string   `json:"backend"`
	WritableRoots []string `json:"writable_roots"`
}

// ConfinementProof is an opaque launcher issued only after the platform
// confinement implementation and every bound filesystem identity validate.
type ConfinementProof struct {
	binding agenttask.ConfinementBinding
	policy  confinementPolicy
}

var _ agenttask.InvocationConfinement = (*ConfinementProof)(nil)

type startedConfinement struct {
	child  *exec.Cmd
	stdin  io.WriteCloser
	stdout io.ReadCloser
	stderr io.ReadCloser

	abortOnce sync.Once
	abortErr  error
}

var _ agenttask.StartedConfinement = (*startedConfinement)(nil)

func (started *startedConfinement) Child() *exec.Cmd {
	if started == nil {
		return nil
	}
	return started.child
}

func (started *startedConfinement) Stdin() io.WriteCloser {
	if started == nil {
		return nil
	}
	return started.stdin
}

func (started *startedConfinement) Stdout() io.ReadCloser {
	if started == nil {
		return nil
	}
	return started.stdout
}

func (started *startedConfinement) Stderr() io.ReadCloser {
	if started == nil {
		return nil
	}
	return started.stderr
}

func (started *startedConfinement) Abort() error {
	if started == nil {
		return nil
	}
	started.abortOnce.Do(func() {
		var cleanupErrors []error
		for _, endpoint := range []io.Closer{
			started.stdin,
			started.stdout,
			started.stderr,
		} {
			if endpoint == nil {
				continue
			}
			if err := endpoint.Close(); err != nil && !errors.Is(err, os.ErrClosed) {
				cleanupErrors = append(cleanupErrors, err)
			}
		}
		if started.child != nil &&
			started.child.Process != nil &&
			started.child.ProcessState == nil {
			if err := started.child.Process.Kill(); err != nil &&
				!errors.Is(err, os.ErrProcessDone) {
				cleanupErrors = append(cleanupErrors, err)
			}
			if err := started.child.Wait(); err != nil {
				var exitError *exec.ExitError
				if !errors.As(err, &exitError) {
					cleanupErrors = append(cleanupErrors, err)
				}
			}
		}
		started.abortErr = errors.Join(cleanupErrors...)
	})
	return started.abortErr
}

func confinementBinding(
	runtimeRoot string,
	record OverlayRecord,
) agenttask.ConfinementBinding {
	return agenttask.ConfinementBinding{
		Revision:           record.ConfinementRevision,
		IsolationID:        record.IsolationID,
		IsolationRevision:  record.Revision,
		PinnedBaseRevision: record.SnapshotRevision,
		ConfigRevision:     record.ConfigRevision,
		GrantRevision:      record.GrantRevision,
		ProfileRevision:    record.ProfileRevision,
		BaseRoot:           record.CanonicalRoot,
		RuntimeRoot:        runtimeRoot,
		SnapshotRoot:       record.Locator.SnapshotRoot,
		TaskRoot:           filepath.Dir(record.Locator.OverlayRecord),
		WorkingDir:         record.Locator.ViewRoot,
		WritableRoots: []string{
			record.Locator.ViewRoot,
			record.Locator.TempRoot,
			record.Locator.CacheRoot,
		},
	}
}

func newConfinementProof(
	binding agenttask.ConfinementBinding,
) (*ConfinementProof, error) {
	if err := validateConfinementBinding(binding); err != nil {
		return nil, err
	}
	revision, err := resolveConfinementRevision(binding)
	if err != nil {
		return nil, err
	}
	if binding.Revision != revision {
		return nil, errors.New(
			"agentworkspace: overlay record has a mismatched confinement revision",
		)
	}
	return &ConfinementProof{
		binding: cloneConfinementBinding(binding),
		policy: confinementPolicy{
			SchemaVersion: confinementPolicySchemaVersion,
			Revision:      revision,
			WritableRoots: append([]string(nil), binding.WritableRoots...),
		},
	}, nil
}

// Revision returns the exact overlay/profile/platform policy revision.
func (proof *ConfinementProof) Revision() string {
	if proof == nil {
		return ""
	}
	return proof.binding.Revision
}

// Binding returns a defensive copy of the immutable proof inputs.
func (proof *ConfinementProof) Binding() agenttask.ConfinementBinding {
	if proof == nil {
		return agenttask.ConfinementBinding{}
	}
	return cloneConfinementBinding(proof.binding)
}

// Validate rejects omission, rebinding, or in-memory corruption before a
// provider invoker receives the proof.
func (proof *ConfinementProof) Validate(
	expected agenttask.ConfinementBinding,
) error {
	if proof == nil {
		return errors.New("confinement proof is missing")
	}
	if err := validateConfinementBinding(proof.binding); err != nil {
		return err
	}
	revision, err := resolveConfinementRevision(proof.binding)
	if err != nil {
		return err
	}
	if proof.binding.Revision != revision ||
		proof.policy.SchemaVersion != confinementPolicySchemaVersion ||
		proof.policy.Revision != revision ||
		!reflect.DeepEqual(proof.policy.WritableRoots, proof.binding.WritableRoots) {
		return errors.New("confinement proof integrity check failed")
	}
	if !reflect.DeepEqual(proof.binding, expected) {
		return errors.New("confinement proof does not match the dispatch identity")
	}
	return nil
}

// Start installs the proof's OS policy and starts the child before returning,
// so the caller cannot replace the sandbox helper path or arguments.
func (proof *ConfinementProof) Start(
	ctx context.Context,
	spec agenttask.ConfinementCommand,
) (agenttask.StartedConfinement, error) {
	if proof == nil {
		return nil, errors.New("agentworkspace: confinement proof is missing")
	}
	if err := proof.Validate(proof.Binding()); err != nil {
		return nil, fmt.Errorf("agentworkspace: validate confinement proof: %w", err)
	}
	if strings.TrimSpace(spec.Name) == "" {
		return nil, errors.New("agentworkspace: confined command is required")
	}
	command, err := platformConfinementCommand(
		ctx,
		proof.binding,
		proof.policy,
		spec.Name,
		spec.Args,
	)
	if err != nil {
		return nil, err
	}
	command.Dir = proof.binding.WorkingDir
	if spec.Env != nil {
		command.Env = append([]string(nil), spec.Env...)
	}
	started, err := startConfinementCommand(command)
	if err != nil {
		return nil, err
	}
	return started, nil
}

type confinementPipeFactory func() (*os.File, *os.File, error)

func startConfinementCommand(
	command *exec.Cmd,
) (agenttask.StartedConfinement, error) {
	return startConfinementCommandWithPipes(command, os.Pipe)
}

func startConfinementCommandWithPipes(
	command *exec.Cmd,
	openPipe confinementPipeFactory,
) (agenttask.StartedConfinement, error) {
	if command == nil {
		return nil, errors.New("agentworkspace: confined child command is missing")
	}
	if openPipe == nil {
		return nil, errors.New("agentworkspace: confinement pipe factory is missing")
	}
	if command.Stdin != nil || command.Stdout != nil || command.Stderr != nil {
		return nil, errors.New("agentworkspace: confined child I/O must be proof-owned")
	}

	var endpoints []*os.File
	closeEndpoints := func() {
		for _, endpoint := range endpoints {
			if endpoint != nil {
				_ = endpoint.Close()
			}
		}
	}
	stdinChild, stdinParent, err := openPipe()
	if err != nil {
		return nil, fmt.Errorf("agentworkspace: create confined stdin pipe: %w", err)
	}
	endpoints = append(endpoints, stdinChild, stdinParent)
	stdoutParent, stdoutChild, err := openPipe()
	if err != nil {
		closeEndpoints()
		return nil, fmt.Errorf("agentworkspace: create confined stdout pipe: %w", err)
	}
	endpoints = append(endpoints, stdoutParent, stdoutChild)
	stderrParent, stderrChild, err := openPipe()
	if err != nil {
		closeEndpoints()
		return nil, fmt.Errorf("agentworkspace: create confined stderr pipe: %w", err)
	}
	endpoints = append(endpoints, stderrParent, stderrChild)

	command.Stdin = stdinChild
	command.Stdout = stdoutChild
	command.Stderr = stderrChild
	if err := command.Start(); err != nil {
		closeEndpoints()
		return nil, fmt.Errorf("agentworkspace: start confined child: %w", err)
	}
	started := &startedConfinement{
		child:  command,
		stdin:  stdinParent,
		stdout: stdoutParent,
		stderr: stderrParent,
	}
	var closeErrors []error
	for _, childEndpoint := range []*os.File{stdinChild, stdoutChild, stderrChild} {
		if err := childEndpoint.Close(); err != nil && !errors.Is(err, os.ErrClosed) {
			closeErrors = append(closeErrors, err)
		}
	}
	if err := errors.Join(closeErrors...); err != nil {
		_ = started.Abort()
		return nil, fmt.Errorf("agentworkspace: release confined child pipe endpoints: %w", err)
	}
	return started, nil
}

func resolveConfinementRevision(
	binding agenttask.ConfinementBinding,
) (string, error) {
	platformRevision, err := platformConfinementRevision()
	if err != nil {
		return "", err
	}
	hashParts := []string{
		fmt.Sprintf("%d", confinementPolicySchemaVersion),
		platformRevision,
		binding.IsolationID,
		binding.IsolationRevision,
		binding.PinnedBaseRevision,
		binding.ConfigRevision,
		binding.GrantRevision,
		binding.ProfileRevision,
		binding.BaseRoot,
		binding.RuntimeRoot,
		binding.SnapshotRoot,
		binding.TaskRoot,
		binding.WorkingDir,
	}
	for _, root := range binding.WritableRoots {
		hashParts = append(hashParts, root)
	}
	hash := sha256.New()
	for _, part := range hashParts {
		writeDigestPart(hash, part)
	}
	return "sha256:" + hex.EncodeToString(hash.Sum(nil)), nil
}

func validateConfinementBinding(binding agenttask.ConfinementBinding) error {
	for name, value := range map[string]string{
		"isolation ID":       binding.IsolationID,
		"isolation revision": binding.IsolationRevision,
		"base revision":      binding.PinnedBaseRevision,
		"config revision":    binding.ConfigRevision,
		"grant revision":     binding.GrantRevision,
		"profile revision":   binding.ProfileRevision,
		"canonical root":     binding.BaseRoot,
		"runtime root":       binding.RuntimeRoot,
		"snapshot root":      binding.SnapshotRoot,
		"task root":          binding.TaskRoot,
		"working directory":  binding.WorkingDir,
	} {
		if strings.TrimSpace(value) == "" {
			return fmt.Errorf("agentworkspace: confinement %s is required", name)
		}
	}
	if len(binding.WritableRoots) != 3 ||
		binding.WritableRoots[0] != binding.WorkingDir {
		return errors.New(
			"agentworkspace: confinement requires exact view, temp, and cache writable roots",
		)
	}
	paths := []string{
		binding.BaseRoot,
		binding.RuntimeRoot,
		binding.SnapshotRoot,
		binding.TaskRoot,
		binding.WorkingDir,
	}
	paths = append(paths, binding.WritableRoots...)
	for _, path := range paths {
		if !filepath.IsAbs(path) || filepath.Clean(path) != path {
			return errors.New(
				"agentworkspace: confinement paths must be absolute and clean",
			)
		}
		resolved, err := filepath.EvalSymlinks(path)
		if err != nil || filepath.Clean(resolved) != path {
			return errors.New(
				"agentworkspace: confinement paths must be existing canonical paths",
			)
		}
	}
	if binding.BaseRoot == binding.TaskRoot ||
		pathContains(binding.BaseRoot, binding.RuntimeRoot) ||
		pathContains(binding.RuntimeRoot, binding.BaseRoot) {
		return errors.New(
			"agentworkspace: canonical and runtime roots must be separate",
		)
	}
	if !pathContains(binding.RuntimeRoot, binding.SnapshotRoot) ||
		!pathContains(binding.RuntimeRoot, binding.TaskRoot) ||
		pathContains(binding.TaskRoot, binding.SnapshotRoot) {
		return errors.New(
			"agentworkspace: snapshot and task roots must have the strict runtime layout",
		)
	}
	seen := make(map[string]struct{}, len(binding.WritableRoots))
	for _, root := range binding.WritableRoots {
		if !pathContains(binding.TaskRoot, root) || root == binding.TaskRoot {
			return errors.New(
				"agentworkspace: writable roots must remain below the task root",
			)
		}
		if _, duplicate := seen[root]; duplicate {
			return errors.New("agentworkspace: writable roots must be unique")
		}
		seen[root] = struct{}{}
	}
	if !pathContains(binding.TaskRoot, binding.WorkingDir) {
		return errors.New(
			"agentworkspace: working directory must remain below the task root",
		)
	}
	return nil
}

func cloneConfinementBinding(
	binding agenttask.ConfinementBinding,
) agenttask.ConfinementBinding {
	binding.WritableRoots = append([]string(nil), binding.WritableRoots...)
	return binding
}