package tmux

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"
)

// RealPane talks to a live tmux pane that runs an LLM CLI. It implements
// PaneIO by:
//   - sending input via `tmux send-keys -l … Enter`
//   - capturing output via `tmux pipe-pane` writing to a log file
//   - tailing that file for the harness-out fenced block
//
// One RealPane = one tmux pane = one CLI session = one agent.
type RealPane struct {
	SocketName  string // tmux -L <socket>; default "harness"
	Session     string // tmux session name; one per agent
	Pane        string // tmux pane target (usually "<session>:0.0")
	LogPath     string // pipe-pane output file
	CLICmd      []string // command to launch in the pane, e.g. ["claude"]
	IdleTimeout time.Duration // how long after last output we consider the turn done
	HardTimeout time.Duration // absolute max wait per Receive
	PromptRegex *regexp.Regexp // optional: detects the CLI's prompt to short-circuit idle wait

	// RateLimitRegex matches output indicating the CLI is rate-limited.
	// When non-nil and Receive sees a match, returns ErrRateLimited.
	RateLimitRegex *regexp.Regexp

	// Env are extra environment variables exported into the new tmux
	// session so the CLI subprocess sees them. Used by sidecar-channel
	// wiring to pass HARNESS_SEND_SOCKET / HARNESS_MAILBOX_PATH down to
	// the model's tool environment. Applied as `tmux new-session -e K=V`.
	Env map[string]string

	mu       sync.Mutex
	readPos  int64       // bytes consumed from LogPath so far
	started  bool
}

// ErrRateLimited is returned by Receive when the CLI emitted a rate-limit
// signature. Caller (the runtime/agent) should park the agent.
var ErrRateLimited = errors.New("tmux: cli rate limited")

// Open creates the tmux session, starts the CLI in it, sets up pipe-pane to
// LogPath, waits for the CLI's first prompt (or IdleTimeout), and returns the
// RealPane ready to Send.
//
// Holding p.mu across the banner-drain would deadlock — drain calls readMore
// which also takes p.mu. So we acquire the mutex only for the state-mutating
// setup, release it, then drain.
func (p *RealPane) Open(ctx context.Context) error {
	p.mu.Lock()
	if p.started {
		p.mu.Unlock()
		return nil
	}
	if len(p.CLICmd) == 0 {
		p.mu.Unlock()
		return errors.New("tmux: CLICmd required")
	}
	if p.SocketName == "" {
		p.SocketName = "harness"
	}
	if p.IdleTimeout == 0 {
		p.IdleTimeout = 3 * time.Second
	}
	if p.HardTimeout == 0 {
		p.HardTimeout = 5 * time.Minute
	}
	if p.LogPath == "" {
		p.mu.Unlock()
		return errors.New("tmux: LogPath required")
	}

	// Ensure parent dir.
	if err := os.MkdirAll(filepath.Dir(p.LogPath), 0o755); err != nil {
		p.mu.Unlock()
		return fmt.Errorf("tmux: mkdir %s: %w", filepath.Dir(p.LogPath), err)
	}
	// Truncate any previous log so readPos starts at 0.
	if err := os.WriteFile(p.LogPath, nil, 0o644); err != nil {
		p.mu.Unlock()
		return fmt.Errorf("tmux: prepare log %s: %w", p.LogPath, err)
	}

	cliShell := shellJoin(p.CLICmd)

	// tmux new-session -d -s <session> -L <socket> [-e K=V ...] "<cli>"
	args := []string{"-L", p.SocketName, "new-session", "-d", "-s", p.Session}
	for _, k := range sortedKeys(p.Env) {
		args = append(args, "-e", k+"="+p.Env[k])
	}
	args = append(args, cliShell)
	if out, err := exec.CommandContext(ctx, "tmux", args...).CombinedOutput(); err != nil {
		p.mu.Unlock()
		return fmt.Errorf("tmux new-session: %w: %s", err, string(out))
	}
	p.Pane = p.Session + ":0.0"

	// Hook pipe-pane to capture every byte the pane writes.
	if out, err := exec.CommandContext(ctx, "tmux",
		"-L", p.SocketName,
		"pipe-pane", "-t", p.Pane, "-o",
		fmt.Sprintf(`cat >> %q`, p.LogPath),
	).CombinedOutput(); err != nil {
		p.mu.Unlock()
		return fmt.Errorf("tmux pipe-pane: %w: %s", err, string(out))
	}
	p.started = true
	p.mu.Unlock()

	// Banner drain runs WITHOUT the mutex — read/readMore re-acquire it.
	bannerCtx, cancel := context.WithTimeout(ctx, p.IdleTimeout*2)
	defer cancel()
	_, _ = p.read(bannerCtx)
	return nil
}

// Close kills the tmux session.
func (p *RealPane) Close() error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if !p.started {
		return nil
	}
	out, err := exec.Command("tmux", "-L", p.SocketName, "kill-session", "-t", p.Session).CombinedOutput()
	p.started = false
	if err != nil && !strings.Contains(string(out), "can't find session") {
		return fmt.Errorf("tmux kill-session: %w: %s", err, string(out))
	}
	return nil
}

// Send delivers the payload to the pane via send-keys. The payload is a
// freeform prompt (already rendered by the per-CLI adapter); we don't add a
// trailing newline because send-keys' final `Enter` arg does that.
func (p *RealPane) Send(ctx context.Context, payload string) error {
	p.mu.Lock()
	if !p.started {
		p.mu.Unlock()
		return errors.New("tmux: pane not open")
	}
	p.mu.Unlock()

	// For multi-line payloads, send-keys -l is safer; it sends literal chars.
	// We send the text, then Enter.
	args := []string{"-L", p.SocketName, "send-keys", "-t", p.Pane, "-l", payload}
	if out, err := exec.CommandContext(ctx, "tmux", args...).CombinedOutput(); err != nil {
		return fmt.Errorf("tmux send-keys -l: %w: %s", err, string(out))
	}
	if out, err := exec.CommandContext(ctx, "tmux",
		"-L", p.SocketName, "send-keys", "-t", p.Pane, "Enter",
	).CombinedOutput(); err != nil {
		return fmt.Errorf("tmux send-keys Enter: %w: %s", err, string(out))
	}
	return nil
}

// Receive tails the log file until the pane goes idle (no new bytes for
// IdleTimeout) AND either (a) the optional PromptRegex matches OR (b)
// a `harness-out` fenced block has appeared in the new data.
//
// Returns ErrRateLimited if RateLimitRegex matches.
func (p *RealPane) Receive(ctx context.Context) (string, error) {
	hard, cancel := context.WithTimeout(ctx, p.HardTimeout)
	defer cancel()

	var buf bytes.Buffer
	lastChange := time.Now()
	ticker := time.NewTicker(150 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-hard.Done():
			if buf.Len() == 0 {
				return "", fmt.Errorf("tmux: hard timeout (no output)")
			}
			return buf.String(), nil
		case <-ticker.C:
		}

		chunk, err := p.readMore()
		if err != nil {
			return "", err
		}
		if len(chunk) > 0 {
			buf.Write(chunk)
			lastChange = time.Now()

			if p.RateLimitRegex != nil && p.RateLimitRegex.MatchString(buf.String()) {
				return buf.String(), ErrRateLimited
			}
			// Early exit: if we already see a complete harness-out block,
			// give the pane one more idle window and return.
			if _, ok := ExtractBlock(buf.String()); ok {
				if time.Since(lastChange) >= p.IdleTimeout/2 {
					return buf.String(), nil
				}
			}
		}

		// Idle + (prompt seen OR harness-out present) → done.
		if time.Since(lastChange) >= p.IdleTimeout {
			s := buf.String()
			if _, ok := ExtractBlock(s); ok {
				return s, nil
			}
			if p.PromptRegex != nil && p.PromptRegex.MatchString(s) {
				return s, nil
			}
			// Still no clear signal — return what we have. The sidecar
			// parser will reformat-retry if it can't extract a block.
			return s, nil
		}
	}
}

// read is the initial banner drain after Open. Same idle loop without the
// block-detection short-circuit.
func (p *RealPane) read(ctx context.Context) (string, error) {
	var buf bytes.Buffer
	lastChange := time.Now()
	ticker := time.NewTicker(150 * time.Millisecond)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return buf.String(), nil
		case <-ticker.C:
		}
		chunk, err := p.readMore()
		if err != nil {
			return "", err
		}
		if len(chunk) > 0 {
			buf.Write(chunk)
			lastChange = time.Now()
		}
		if time.Since(lastChange) >= p.IdleTimeout {
			return buf.String(), nil
		}
	}
}

// readMore reads any new bytes appended to LogPath since the last call.
func (p *RealPane) readMore() ([]byte, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	f, err := os.Open(p.LogPath)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	st, err := f.Stat()
	if err != nil {
		return nil, err
	}
	size := st.Size()
	if size <= p.readPos {
		return nil, nil
	}
	if _, err := f.Seek(p.readPos, 0); err != nil {
		return nil, err
	}
	buf := make([]byte, size-p.readPos)
	n, err := f.Read(buf)
	if err != nil {
		return nil, err
	}
	p.readPos += int64(n)
	return buf[:n], nil
}

// shellJoin builds a single shell-quotable command string.
func shellJoin(parts []string) string {
	q := make([]string, len(parts))
	for i, p := range parts {
		if strings.ContainsAny(p, " \t'\"$\\") {
			q[i] = "'" + strings.ReplaceAll(p, "'", `'\''`) + "'"
		} else {
			q[i] = p
		}
	}
	return strings.Join(q, " ")
}

// TmuxAvailable returns true if `tmux -V` exits cleanly. Spawners use this to
// short-circuit before attempting Open.
func TmuxAvailable() bool {
	return exec.Command("tmux", "-V").Run() == nil
}

// sortedKeys returns m's keys in deterministic order. Used to make
// `tmux new-session -e K=V` argv stable across runs (helps log diffs
// and test reproduction).
func sortedKeys(m map[string]string) []string {
	if len(m) == 0 {
		return nil
	}
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}