package toolsrv

import (
	"context"
	"fmt"
	"sync"
	"time"
)

// ProcessKeeper maintains a running ollie-remote process and respawns it
// on failure. Multiple agents can Dial() the same keeper to get independent
// connections. If the process dies, the next Dial() triggers a respawn.
type ProcessKeeper struct {
	mu    sync.Mutex
	ctx   context.Context
	proc  *Process
	spawn func(context.Context) (*Process, error)
}

// NewProcessKeeper creates a keeper with an initial process and a spawn factory.
func NewProcessKeeper(ctx context.Context, proc *Process, spawn func(context.Context) (*Process, error)) *ProcessKeeper {
	return &ProcessKeeper{
		ctx:   ctx,
		proc:  proc,
		spawn: spawn,
	}
}

// Dial returns a new Conn to the managed process. If the process is dead
// (socket unreachable), it respawns before dialing.
func (pk *ProcessKeeper) Dial() (*Conn, error) {
	pk.mu.Lock()
	sock := pk.proc.Socket
	pk.mu.Unlock()

	// Try dialing the existing socket.
	conn, err := Dial(pk.ctx, sock)
	if err == nil {
		return conn, nil
	}

	// Dial failed — process likely dead. Respawn.
	return pk.respawnAndDial()
}

// Socket returns the current process socket path.
func (pk *ProcessKeeper) Socket() string {
	pk.mu.Lock()
	defer pk.mu.Unlock()
	return pk.proc.Socket
}

// Process returns the current underlying Process.
func (pk *ProcessKeeper) Process() *Process {
	pk.mu.Lock()
	defer pk.mu.Unlock()
	return pk.proc
}

// Close shuts down the managed process.
func (pk *ProcessKeeper) Close() {
	pk.mu.Lock()
	defer pk.mu.Unlock()
	if pk.proc != nil {
		pk.proc.Close()
	}
}

// SetContext updates the keeper's context. Used on session resume to
// provide a fresh context after the previous one was cancelled on pause.
func (pk *ProcessKeeper) SetContext(ctx context.Context) {
	pk.mu.Lock()
	defer pk.mu.Unlock()
	pk.ctx = ctx
}

func (pk *ProcessKeeper) respawnAndDial() (*Conn, error) {
	pk.mu.Lock()
	defer pk.mu.Unlock()

	// Another goroutine may have respawned already — try current socket first.
	if conn, err := Dial(pk.ctx, pk.proc.Socket); err == nil {
		return conn, nil
	}

	// Clean up old process.
	pk.proc.Close()

	var lastErr error
	for attempt := 0; attempt < 3; attempt++ {
		if pk.ctx.Err() != nil {
			return nil, pk.ctx.Err()
		}
		proc, err := pk.spawn(pk.ctx)
		if err != nil {
			lastErr = err
			time.Sleep(time.Duration(attempt+1) * 2 * time.Second)
			continue
		}
		pk.proc = proc

		conn, err := Dial(pk.ctx, proc.Socket)
		if err != nil {
			lastErr = err
			proc.Close()
			time.Sleep(time.Duration(attempt+1) * 2 * time.Second)
			continue
		}
		return conn, nil
	}
	return nil, fmt.Errorf("respawn failed after 3 attempts: %w", lastErr)
}