package transport

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"sync"
	"time"
)

// DefaultShutdownTimeout is the default time to wait for a subprocess to
// exit gracefully before escalating to SIGKILL.
const DefaultShutdownTimeout = 2 * time.Second

// Stdio returns os.Stdin and os.Stdout as an io.Reader/io.Writer pair.
// This is used by agent processes to communicate over their own stdin/stdout.
func Stdio() (io.Reader, io.Writer) {
	return os.Stdin, os.Stdout
}

// Subprocess manages a spawned agent or client process with stdin/stdout pipes.
type Subprocess struct {
	// Stdin writes to the subprocess's stdin.
	Stdin io.WriteCloser
	// Stdout reads from the subprocess's stdout.
	Stdout io.Reader
	// Stderr reads from the subprocess's stderr (may be nil if not captured).
	Stderr io.Reader

	cmd             *exec.Cmd
	shutdownTimeout time.Duration
	mu              sync.Mutex
	waited          bool
}

// SpawnConfig holds configuration for spawning a subprocess.
type SpawnConfig struct {
	// Command is the executable to run (required).
	Command string
	// Args are arguments to pass to the command.
	Args []string
	// Env is additional environment variables (merged onto a clean base env).
	Env map[string]string
	// Cwd is the working directory for the subprocess.
	Cwd string
	// CaptureStderr controls whether stderr is captured. Defaults to true.
	CaptureStderr bool
	// ShutdownTimeout is how long to wait for graceful shutdown. Defaults to DefaultShutdownTimeout.
	ShutdownTimeout time.Duration
}

// SpawnOption is a functional option for SpawnConfig.
type SpawnOption func(*SpawnConfig)

// WithArgs sets the arguments for the subprocess.
func WithArgs(args ...string) SpawnOption {
	return func(c *SpawnConfig) { c.Args = args }
}

// WithEnv sets additional environment variables for the subprocess.
func WithEnv(env map[string]string) SpawnOption {
	return func(c *SpawnConfig) { c.Env = env }
}

// WithCwd sets the working directory for the subprocess.
func WithCwd(cwd string) SpawnOption {
	return func(c *SpawnConfig) { c.Cwd = cwd }
}

// WithCaptureStderr controls stderr capture.
func WithCaptureStderr(capture bool) SpawnOption {
	return func(c *SpawnConfig) { c.CaptureStderr = capture }
}

// WithShutdownTimeout sets the shutdown timeout.
func WithShutdownTimeout(d time.Duration) SpawnOption {
	return func(c *SpawnConfig) { c.ShutdownTimeout = d }
}

// Spawn starts a subprocess with the given command and returns a Subprocess
// with connected stdin/stdout pipes.
// Use Subprocess.Close() to shut down the subprocess gracefully.
func Spawn(command string, opts ...SpawnOption) (*Subprocess, error) {
	cfg := &SpawnConfig{
		Command:         command,
		CaptureStderr:   true,
		ShutdownTimeout: DefaultShutdownTimeout,
	}
	for _, opt := range opts {
		opt(cfg)
	}

	cmd := exec.Command(cfg.Command, cfg.Args...)
	cmd.Stdin = nil // will be replaced with a pipe

	// Build clean environment
	cmd.Env = buildEnv(cfg.Env)

	if cfg.Cwd != "" {
		cmd.Dir = cfg.Cwd
	}

	stdinPipe, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("spawn: stdin pipe: %w", err)
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("spawn: stdout pipe: %w", err)
	}

	var stderrPipe io.Reader
	if cfg.CaptureStderr {
		stderrPipe, err = cmd.StderrPipe()
		if err != nil {
			return nil, fmt.Errorf("spawn: stderr pipe: %w", err)
		}
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("spawn: start: %w", err)
	}

	return &Subprocess{
		Stdin:           stdinPipe,
		Stdout:          stdoutPipe,
		Stderr:          stderrPipe,
		cmd:             cmd,
		shutdownTimeout: cfg.ShutdownTimeout,
	}, nil
}

// Close shuts down the subprocess gracefully: closes stdin, waits for exit,
// then escalates to kill if necessary.
func (s *Subprocess) Close() error {
	timeout := s.shutdownTimeout
	if timeout <= 0 {
		timeout = DefaultShutdownTimeout
	}
	return s.closeWithTimeout(timeout)
}

// CloseWithTimeout shuts down the subprocess with a custom timeout.
func (s *Subprocess) CloseWithTimeout(timeout time.Duration) error {
	return s.closeWithTimeout(timeout)
}

func (s *Subprocess) closeWithTimeout(timeout time.Duration) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.waited {
		return nil
	}

	// Close stdin to signal the process
	if s.Stdin != nil {
		s.Stdin.Close()
	}

	// Wait for the process to exit
	done := make(chan error, 1)
	go func() {
		done <- s.cmd.Wait()
	}()

	select {
	case err := <-done:
		s.waited = true
		return err
	case <-time.After(timeout):
		// Escalate: send SIGKILL
		s.cmd.Process.Kill()
		s.waited = true
		return <-done
	}
}

// ProcessID returns the OS process ID of the subprocess.
func (s *Subprocess) ProcessID() int {
	if s.cmd != nil && s.cmd.Process != nil {
		return s.cmd.Process.Pid
	}
	return 0
}

// buildEnv returns a clean environment with extra vars merged in.
func buildEnv(extra map[string]string) []string {
	// Start with essential vars from the current environment
	keep := map[string]bool{
		"HOME":            true,
		"PATH":            true,
		"SHELL":           true,
		"USER":            true,
		"LANG":            true,
		"TERM":            true,
		"TZ":              true,
		"XDG_CONFIG_HOME": true,
		"APPDATA":         true, // Windows
		"SYSTEMROOT":      true, // Windows
	}

	env := make([]string, 0)
	for _, e := range os.Environ() {
		for k := range keep {
			if len(e) > len(k) && e[:len(k)+1] == k+"=" {
				env = append(env, e)
				break
			}
		}
	}

	// Merge extra vars
	for k, v := range extra {
		env = append(env, k+"="+v)
	}

	return env
}