package tools

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"
)

// maxBackgroundOutput is how much of one background process's output is kept.
// A dev server logs every request, so the buffer has to be bounded or a
// session that leaves one running for an hour ends up holding its whole log.
// The oldest bytes go first: what a model wants from a server it started is
// the most recent lines, not the banner.
const maxBackgroundOutput = 256 << 10

// Processes owns the commands started with background: true. They outlive the
// call that started them, which is the point — a dev server has to still be
// running when the browser asks for a page — so something has to own them and
// kill them at the end of the session. Nothing else in the harness has a
// lifetime long enough to do that.
type Processes struct {
	mu     sync.Mutex
	next   int
	byID   map[string]*process
	closed bool
}

func NewProcesses() *Processes {
	return &Processes{byID: map[string]*process{}}
}

// process is one background command and everything the model can ask about it.
type process struct {
	id      string
	command string
	out     *ring
	cmd     *exec.Cmd
	started time.Time

	mu       sync.Mutex
	done     bool
	exitCode int
	runErr   error
	ended    time.Time
}

// Start launches command in dir and returns its handle. The command is staged
// through a file for the same reason the foreground path does it: with
// `bash -c` the whole script sits in the shell's own argv, where a pkill
// pattern in the script matches the shell running it.
func (p *Processes) Start(dir, command string) (string, error) {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return "", fmt.Errorf("the session is shutting down; no new background processes")
	}
	p.next++
	id := fmt.Sprintf("bg_%d", p.next)
	p.mu.Unlock()

	script, err := stageScript(command)
	if err != nil {
		return "", err
	}

	proc := &process{
		id:      id,
		command: command,
		out:     &ring{},
		started: time.Now(),
	}
	// No context here on purpose. A background process is meant to outlive the
	// tool call that started it; cancelling it with that call's context would
	// kill the dev server the moment the model's turn ended.
	cmd := exec.Command("bash", script)
	cmd.Dir = resolveBase(dir, "")
	cmd.Stdout = proc.out
	cmd.Stderr = proc.out
	// Its own process group, so the whole tree can be killed later. A server
	// started from a shell is a child of that shell, and killing the shell
	// alone leaves the port held by an orphan nobody can see.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	if err := cmd.Start(); err != nil {
		os.Remove(script)
		return "", fmt.Errorf("cannot start: %w", err)
	}
	proc.cmd = cmd

	p.mu.Lock()
	p.byID[id] = proc
	p.mu.Unlock()

	go func() {
		err := cmd.Wait()
		os.Remove(script)
		proc.mu.Lock()
		proc.done, proc.runErr, proc.ended = true, err, time.Now()
		proc.exitCode = -1
		if cmd.ProcessState != nil {
			proc.exitCode = cmd.ProcessState.ExitCode()
		}
		proc.mu.Unlock()
	}()
	return id, nil
}

// Report returns the status line for one process and whatever output has
// appeared since it was last asked. Reads are incremental because polling is
// the normal way to use this: re-sending the whole log every turn would cost
// more context than the process is worth.
func (p *Processes) Report(id string) (string, error) {
	proc, err := p.get(id)
	if err != nil {
		return "", err
	}
	fresh, dropped := proc.out.next()

	var b strings.Builder
	b.WriteString(proc.status())
	if dropped > 0 {
		fmt.Fprintf(&b, "\n(%d bytes of older output were dropped: the buffer keeps the most recent %s)",
			dropped, humanBytes(maxBackgroundOutput))
	}
	if fresh == "" {
		b.WriteString("\n(no new output)")
	} else {
		fmt.Fprintf(&b, "\n%s", headTail(strings.TrimRight(fresh, "\n"), bashHeadLines, bashTailLines))
	}
	return b.String(), nil
}

// List is what the model gets when it asks for output without naming one,
// which is how it recovers from having lost track of an id.
func (p *Processes) List() string {
	p.mu.Lock()
	procs := make([]*process, 0, len(p.byID))
	for _, proc := range p.byID {
		procs = append(procs, proc)
	}
	p.mu.Unlock()
	if len(procs) == 0 {
		return "No background processes."
	}
	sort.Slice(procs, func(i, j int) bool { return procs[i].id < procs[j].id })

	var b strings.Builder
	b.WriteString("Background processes:")
	for _, proc := range procs {
		fmt.Fprintf(&b, "\n  %s", proc.status())
	}
	return b.String()
}

// Kill stops one process and its children. Killing something that has already
// exited is a no-op rather than an error: a model tidying up at the end of a
// task should not have to check first, and an error there costs a turn.
func (p *Processes) Kill(id string) (string, error) {
	proc, err := p.get(id)
	if err != nil {
		return "", err
	}
	if proc.finished() {
		return fmt.Sprintf("%s had already exited.", id), nil
	}
	proc.kill()
	// Give Wait a moment so the status the model reads back is the real one.
	for i := 0; i < 100 && !proc.finished(); i++ {
		time.Sleep(5 * time.Millisecond)
	}
	return fmt.Sprintf("Killed %s (%s).", id, oneLineCommand(proc.command)), nil
}

// Close kills everything still running. It is called when the session ends:
// a process that outlives the harness holds its port with nothing left to
// explain why.
func (p *Processes) Close() {
	p.mu.Lock()
	p.closed = true
	procs := make([]*process, 0, len(p.byID))
	for _, proc := range p.byID {
		procs = append(procs, proc)
	}
	p.mu.Unlock()

	for _, proc := range procs {
		if !proc.finished() {
			proc.kill()
		}
	}
	// Reaping is best-effort: the group has had SIGKILL, and a shutdown path
	// that can block is worse than a process that takes a moment to die.
	deadline := time.Now().Add(time.Second)
	for _, proc := range procs {
		for !proc.finished() && time.Now().Before(deadline) {
			time.Sleep(5 * time.Millisecond)
		}
	}
}

// Running reports how many are still alive, for a caller that wants to say so
// before exiting.
func (p *Processes) Running() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	var n int
	for _, proc := range p.byID {
		if !proc.finished() {
			n++
		}
	}
	return n
}

// buffered is the buffer size, for the test that pins the cap down.
func (p *Processes) buffered(id string) int {
	proc, err := p.get(id)
	if err != nil {
		return 0
	}
	proc.out.mu.Lock()
	defer proc.out.mu.Unlock()
	return len(proc.out.buf)
}

// get resolves an id, and when it misses says which ids exist. A model that
// guessed an id recovers in the same turn instead of guessing again.
func (p *Processes) get(id string) (*process, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if proc, ok := p.byID[id]; ok {
		return proc, nil
	}
	if len(p.byID) == 0 {
		return nil, fmt.Errorf("no such background process %q; none have been started", id)
	}
	ids := make([]string, 0, len(p.byID))
	for k := range p.byID {
		ids = append(ids, k)
	}
	sort.Strings(ids)
	return nil, fmt.Errorf("no such background process %q. The ones that exist are: %s",
		id, strings.Join(ids, ", "))
}

func (pr *process) finished() bool {
	pr.mu.Lock()
	defer pr.mu.Unlock()
	return pr.done
}

func (pr *process) kill() {
	if pr.cmd.Process == nil {
		return
	}
	// Negative pid is the process group, which is why Setpgid was set.
	syscall.Kill(-pr.cmd.Process.Pid, syscall.SIGKILL)
}

func (pr *process) status() string {
	pr.mu.Lock()
	done, code, runErr, ended := pr.done, pr.exitCode, pr.runErr, pr.ended
	pr.mu.Unlock()

	cmd := oneLineCommand(pr.command)
	if !done {
		return fmt.Sprintf("%s · running · %s · %s",
			pr.id, time.Since(pr.started).Round(time.Second), cmd)
	}
	took := ended.Sub(pr.started).Round(time.Millisecond)
	if code < 0 && runErr != nil {
		return fmt.Sprintf("%s · failed to run: %v · %s · %s", pr.id, runErr, took, cmd)
	}
	return fmt.Sprintf("%s · exit code: %d · %s · %s", pr.id, code, took, cmd)
}

// oneLineCommand keeps a multi-line script recognisable in a status line
// without letting it take over the output.
func oneLineCommand(s string) string {
	s = strings.TrimSpace(strings.ReplaceAll(s, "\n", " "))
	if len(s) > 80 {
		return s[:77] + "..."
	}
	return s
}

// ring is the output buffer: bounded, and with a cursor marking what has
// already been handed to the model. Offsets are absolute — counted from the
// first byte ever written — so dropping from the front does not invalidate
// the cursor, it just moves past it.
type ring struct {
	mu      sync.Mutex
	buf     []byte
	dropped int64
	read    int64
}

func (r *ring) Write(p []byte) (int, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.buf = append(r.buf, p...)
	if over := len(r.buf) - maxBackgroundOutput; over > 0 {
		r.buf = append(r.buf[:0], r.buf[over:]...)
		r.dropped += int64(over)
	}
	return len(p), nil
}

// next returns everything written since the last call, and how many unread
// bytes were dropped to stay inside the cap.
func (r *ring) next() (string, int64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	end := r.dropped + int64(len(r.buf))
	var lost int64
	if r.read < r.dropped {
		lost = r.dropped - r.read
		r.read = r.dropped
	}
	s := string(r.buf[r.read-r.dropped:])
	r.read = end
	return s, lost
}

// BashProcesses is the option that gives the bash tool somewhere to put a
// background process. Without it the tool has no way to own one, and asking
// for background is refused rather than quietly run in the foreground.
func WithProcesses(p *Processes) Option {
	return func(o *options) { o.procs = p }
}

// ProcessOutputArgs is the process_output input.
type ProcessOutputArgs struct {
	ID string `json:"id,omitempty" jsonschema_description:"The handle bash returned, such as bg_1. Omit it to list every background process."`
}

// NewProcessOutput builds the process_output tool. It returns nil when there
// is no manager, so a harness without background support does not advertise a
// tool that cannot work.
func NewProcessOutput(p *Processes) Tool {
	if p == nil {
		return nil
	}
	return New("process_output", strings.TrimSpace(`
Read what a background process has printed since you last asked, and whether it
is still running. Poll this after starting a server to know when it is ready.
Each read returns only new output, so calling it repeatedly is cheap. Omit id
to list every background process.`),
		func(_ context.Context, in ProcessOutputArgs) (string, error) {
			if strings.TrimSpace(in.ID) == "" {
				return p.List(), nil
			}
			return p.Report(strings.TrimSpace(in.ID))
		})
}

// ProcessKillArgs is the process_kill input.
type ProcessKillArgs struct {
	ID string `json:"id" jsonschema_description:"The handle bash returned, such as bg_1"`
}

// NewProcessKill builds the process_kill tool.
func NewProcessKill(p *Processes) Tool {
	if p == nil {
		return nil
	}
	return New("process_kill", strings.TrimSpace(`
Stop a background process and everything it started. Use this rather than pkill
or kill: it knows the process group, and a pattern you pass to pkill can match
the shell running your own command.`),
		func(_ context.Context, in ProcessKillArgs) (string, error) {
			if strings.TrimSpace(in.ID) == "" {
				return "", fmt.Errorf("id is required; call process_output with no id to see what is running")
			}
			return p.Kill(strings.TrimSpace(in.ID))
		})
}

// stageScript writes a command to a temp file and returns the path. The caller
// removes it once the command has finished.
func stageScript(command string) (string, error) {
	f, err := os.CreateTemp("", "codey-bash-*.sh")
	if err != nil {
		return "", fmt.Errorf("cannot stage the command: %w", err)
	}
	if _, err := f.WriteString(command); err != nil {
		f.Close()
		os.Remove(f.Name())
		return "", fmt.Errorf("cannot stage the command: %w", err)
	}
	f.Close()
	return f.Name(), nil
}