package agents

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"sync/atomic"
	"time"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/mcp"
	"g.fc.run/ai/byre/internal/terminal"
)

const (
	typeQuiet   = 500 * time.Millisecond
	typeMaxWait = 10 * time.Second
	typeKeep    = 256 << 10
)

var agentTerminalSize = api.TerminalSize{Cols: 120, Rows: 40}

func (m *Manager) terminalReadTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "terminal_read",
		Description: "Read your task's terminal: the shell in this task's worktree that the human uses from the Terminal tab, and that terminal_type types into. " +
			"Returns its recent output as plain text, escape sequences removed, at most the last 64 KiB. " +
			"Full-screen programs (vim, top, less) come out garbled, since this is the output stream, not a picture of the screen.",
		InputSchema: json.RawMessage(`{"type": "object", "properties": {}}`),
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			term := m.terminals.Get(l.taskID)
			if term == nil {
				return "No shell is running in this task's terminal.", nil
			}
			out := terminal.Text(term.Backlog())
			exit, exited := term.ExitStatus()
			if exited {
				out += "\n[shell ended: " + exit + "]"
			}
			return out, nil
		},
	}
}

type terminalTypeArgs struct {
	Text string `json:"text"`
}

func (m *Manager) terminalTypeTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "terminal_type",
		Description: "Type into your task's terminal: the shell in this task's worktree that the human uses from the Terminal tab, starting it when none is running and waiting for its prompt before typing. " +
			"text goes to the shell exactly as keystrokes, so end a command with \"\\n\" to run it, and send control characters for keys (\"\\u0003\" is Ctrl-C). " +
			"The human sees everything you type and its output. " +
			"To hand the human a command to run (one needing their password, say), type it without the newline and tell them to press Enter in the Terminal tab. In a zsh terminal, the reply then says byre will message you when the shell next returns to its prompt, with the exit status and output, starting a turn if you are idle, so end your turn rather than asking whether they ran it. " +
			"Returns the output that followed, as terminal_read does, once the terminal has been quiet for half a second or after 10 seconds; read the rest of a longer command's output with terminal_read.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"text": {"type": "string", "description": "The keystrokes to send."}
			},
			"required": ["text"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			args := terminalTypeArgs{}
			err := json.Unmarshal(raw, &args)
			if err != nil {
				return "", fmt.Errorf("malformed terminal_type call: %w", err)
			}
			task, err := m.store.GetTask(ctx, l.taskID)
			if err != nil {
				return "", err
			}
			if task.Status != api.TaskStatusActive {
				return "", fmt.Errorf("task %s is %s", task.Id, task.Status)
			}
			term, err := m.terminals.Attach(task, agentTerminalSize)
			if err != nil {
				return "", err
			}
			if args.Text == "" || strings.HasSuffix(args.Text, "\n") || strings.HasSuffix(args.Text, "\r") {
				return typeAndWait(ctx, term, args.Text, nil)
			}
			if !term.Prompts {
				out, err := typeAndWait(ctx, term, args.Text, nil)
				return out + "\n[this shell isn't zsh, so byre can't tell when the line is run: ask the human to tell you once they have run it]", err
			}
			out, err := typeAndWait(ctx, term, args.Text, m.wakeOnPrompt(l))
			return out + "\n[byre messages you when the shell is next back at its prompt]", err
		},
	}
}

func (m *Manager) wakeOnPrompt(l *live) func(string, []byte) {
	return func(exit string, out []byte) {
		text := fmt.Sprintf("The line you left unsent in the terminal was run, and the shell is back at its prompt with exit status %s (130 means it was cancelled with Ctrl-C). The terminal's output since you typed it:\n\n%s", exit, terminal.Text(out))
		_, _, err := m.Message(context.Background(), l.id, "byre", text, "")
		if err != nil {
			slog.Error("wake agent on terminal prompt", "session", l.id, "task", l.taskID, "err", err)
			return
		}
		slog.Info("woke agent on terminal prompt", "session", l.id, "task", l.taskID, "exit", exit)
	}
}

func typeAndWait(ctx context.Context, term *terminal.Terminal, text string, wake func(string, []byte)) (string, error) {
	term.CancelWake()
	backlog, out := term.Subscribe()
	if out == nil {
		exit, _ := term.ExitStatus()
		return "", fmt.Errorf("the shell ended (%s); type again to start a new one", exit)
	}
	defer term.Unsubscribe(out)

	got := []byte{}
	if promptWait(term) > 0 {
		var err error
		got, err = awaitPrompt(ctx, term, out, keepTail(backlog))
		if err != nil {
			return "", err
		}
	}

	stop := atomic.Bool{}
	defer stop.Store(true)
	typed := atomic.Int64{}
	written := make(chan error, 1)
	go func() {
		// One byte per write: a program that stops reading blocks a write, and the PTY can't be given a deadline, so this keeps what arrives after we give up to one byte.
		for i := 0; i < len(text) && !stop.Load(); i++ {
			err := term.Write([]byte{text[i]})
			if err != nil {
				written <- err
				return
			}
			typed.Add(1)
		}
		if wake != nil && typed.Load() == int64(len(text)) {
			term.WakeOnPrompt(wake)
		}
		written <- nil
	}()

	result := func() string {
		res := terminal.Text(got)
		exit, exited := term.ExitStatus()
		if exited {
			res += "\n[shell ended: " + exit + "]"
		}
		if n := typed.Load(); n < int64(len(text)) {
			res += fmt.Sprintf("\n[only %d of %d bytes were typed: the terminal stopped reading input]", n, len(text))
		}
		return res
	}
	deadline := time.After(typeMaxWait)
	quiet := time.NewTimer(typeQuiet)
	defer quiet.Stop()
	for {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case err := <-written:
			if err != nil {
				return "", err
			}
			written = nil
		case <-deadline:
			return result(), nil
		case <-quiet.C:
			return result(), nil
		case b, open := <-out:
			if !open {
				return result(), nil
			}
			got = keepTail(append(got, b...))
			quiet.Reset(typeQuiet)
		}
	}
}

func keepTail(b []byte) []byte {
	if len(b) > 2*typeKeep {
		return bytes.Clone(b[len(b)-typeKeep:])
	}
	return b
}

func promptWait(term *terminal.Terminal) time.Duration {
	started, last := term.Activity()
	wait := time.Until(started.Add(typeMaxWait))
	if !last.IsZero() {
		wait = min(wait, time.Until(last.Add(typeQuiet)))
	}
	return wait
}

func awaitPrompt(ctx context.Context, term *terminal.Terminal, out chan []byte, got []byte) ([]byte, error) {
	for {
		wait := promptWait(term)
		if wait <= 0 {
			return got, nil
		}
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			return nil, ctx.Err()
		case <-timer.C:
		case b, open := <-out:
			timer.Stop()
			if !open {
				return got, nil
			}
			got = keepTail(append(got, b...))
		}
	}
}