package infra

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"sync"
	"syscall"
	"time"

	"github.com/creack/pty"
)

// PTYCaptureIdle launches a command inside a pseudo-terminal (PTY) and captures
// all output until the terminal has been idle for at least idle duration, or
// maxDur has elapsed. Before sending inputs it waits warmup to let the process
// initialize. Between each input step it waits stepDelay.
//
// The returned string is the raw PTY output, ANSI escape sequences included.
// To turn it into plain text: use vt_render_go_tui to reconstruct the 2D screen
// layout for TUIs with absolute cursor positioning (claude, htop), or
// strip_ansi_go_core for sequential, log-like output.
func PTYCaptureIdle(
	ctx context.Context,
	name string,
	args []string,
	warmup time.Duration,
	inputs []string,
	stepDelay time.Duration,
	idle time.Duration,
	maxDur time.Duration,
) (string, error) {
	cmd := exec.CommandContext(ctx, name, args...)

	ptmx, err := pty.Start(cmd)
	if err != nil {
		return "", fmt.Errorf("pty_capture_idle: pty.Start: %w", err)
	}

	// Set a reasonable terminal size so TUIs render without truncating.
	if szErr := pty.Setsize(ptmx, &pty.Winsize{Rows: 40, Cols: 120}); szErr != nil {
		// Non-fatal: continue even if resize fails.
		_ = szErr
	}

	var (
		mu       sync.Mutex
		buf      bytes.Buffer
		lastByte = time.Now()
	)

	// Reader goroutine: copy PTY output into buf and track last-byte time.
	readDone := make(chan struct{})
	go func() {
		defer close(readDone)
		tmp := make([]byte, 4096)
		for {
			n, rerr := ptmx.Read(tmp)
			if n > 0 {
				mu.Lock()
				buf.Write(tmp[:n])
				lastByte = time.Now()
				mu.Unlock()
			}
			if rerr != nil {
				// EIO/EOF is normal on Linux when the PTY master is closed
				// after the child exits. Not a real error.
				return
			}
		}
	}()

	start := time.Now()

	// Wait for warmup so the TUI/CLI has time to initialize.
	select {
	case <-time.After(warmup):
	case <-ctx.Done():
		_ = ptmx.Close()
		<-readDone
		mu.Lock()
		out := buf.String()
		mu.Unlock()
		return out, fmt.Errorf("pty_capture_idle: context cancelled during warmup: %w", ctx.Err())
	}

	// Send inputs one by one with stepDelay between them.
	for _, in := range inputs {
		if _, werr := ptmx.Write([]byte(in)); werr != nil {
			// PTY may have closed already; stop sending.
			break
		}
		select {
		case <-time.After(stepDelay):
		case <-ctx.Done():
			goto done
		}
	}

done:
	// Poll until idle or maxDur.
	ticker := time.NewTicker(100 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			mu.Lock()
			sinceLastByte := time.Since(lastByte)
			mu.Unlock()

			elapsed := time.Since(start)
			if sinceLastByte >= idle || elapsed >= maxDur {
				goto shutdown
			}
		case <-ctx.Done():
			goto shutdown
		}
	}

shutdown:
	// Close the PTY master. This signals EOF to the reader goroutine.
	_ = ptmx.Close()

	// Graceful shutdown: SIGTERM first, then SIGKILL after 2s.
	if cmd.Process != nil {
		_ = cmd.Process.Signal(syscall.SIGTERM)
		killTimer := time.NewTimer(2 * time.Second)
		waitCh := make(chan error, 1)
		go func() { waitCh <- cmd.Wait() }()
		select {
		case <-waitCh:
			// Process exited cleanly.
		case <-killTimer.C:
			_ = cmd.Process.Kill()
			<-waitCh
		}
		killTimer.Stop()
	}

	<-readDone

	mu.Lock()
	out := buf.String()
	mu.Unlock()

	return out, nil
}