// Package scripted is a deterministic Runtime for testing the harness's
// containment of misbehaving LLMs. Each agent gets a Script — a list of
// per-turn LLMResponses the runtime returns in order. The harness exercises
// these against its tool/zone/cost guards.
package scripted

import (
	"context"
	"strings"
	"sync"

	"github.com/flothus/tmux-xterm-research/server-go/internal/harness/runtime"
)

// Runtime is a scripted Runtime suitable for tests.
type Runtime struct {
	// DefaultAutoComplete, when true, makes CallLLM synthesize a "did the
	// work and reported back" turn for any agent with no scripted response.
	// Used by the CLI's --runtime scripted mode so live runs progress
	// without per-agent script setup; tests should leave it false for
	// deterministic per-script behavior.
	DefaultAutoComplete bool

	// RoleBehaviors maps role.id → behavior fn. Overrides the generic
	// auto-complete with role-specific actions: debuggers write reports,
	// plan-writers write plans, workers write code. Production-grade
	// scripted runs produce real artifacts the user can inspect.
	RoleBehaviors map[string]RoleBehavior

	// AgentResolver maps role.id → actual spawned agent id (for the current
	// run). Wired by runlive after SpawnAll so behaviors can emit delegate
	// envelopes with real recipient ids instead of "<role>-AGENTID"
	// placeholders.
	AgentResolver func(role string) string

	mu      sync.Mutex
	scripts map[string][]*runtime.LLMResponse // agentID → per-turn responses
	turn    map[string]int
	status  map[string]string
	roles   map[string]string // agentID → role.id
}

// RoleBehavior takes the incoming envelope + agent metadata and returns
// the LLM response to emit. Provides a hook for tests/the runlive layer to
// install per-role logic without hardcoding behaviors here.
type RoleBehavior func(agentID, role string, in *runtime.LLMRequest) *runtime.LLMResponse

// New returns an empty scripted runtime.
func New() *Runtime {
	return &Runtime{
		scripts:       map[string][]*runtime.LLMResponse{},
		turn:          map[string]int{},
		status:        map[string]string{},
		roles:         map[string]string{},
		RoleBehaviors: map[string]RoleBehavior{},
	}
}

// SetScript installs the per-turn response list for an agent.
func (r *Runtime) SetScript(agentID string, turns []*runtime.LLMResponse) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.scripts[agentID] = turns
	r.turn[agentID] = 0
}

// Spawn records the agent and resets its turn counter.
func (r *Runtime) Spawn(ctx context.Context, spec runtime.SpawnSpec) (string, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.scripts[spec.AgentID]; !ok {
		r.scripts[spec.AgentID] = nil
	}
	r.turn[spec.AgentID] = 0
	r.status[spec.AgentID] = "running"
	r.roles[spec.AgentID] = spec.Role
	return spec.AgentID, nil
}

// CallLLM returns the next scripted response for the agent, OR — if no
// script is set for this agent and DefaultAutoComplete is true — synthesizes
// a "do the work and report back" turn so live runs without per-agent
// scripts (e.g. the CLI with --runtime scripted) still progress.
func (r *Runtime) CallLLM(ctx context.Context, agentID string, req runtime.LLMRequest) (*runtime.LLMResponse, error) {
	r.mu.Lock()
	script, ok := r.scripts[agentID]
	role := r.roles[agentID]
	r.mu.Unlock()
	if !ok {
		if r.DefaultAutoComplete {
			return r.synthesizeForRole(agentID, role, req), nil
		}
		return nil, runtime.ErrUnknownAgent
	}
	r.mu.Lock()
	t := r.turn[agentID]
	r.mu.Unlock()
	if t >= len(script) {
		if r.DefaultAutoComplete {
			return r.synthesizeForRole(agentID, role, req), nil
		}
		return &runtime.LLMResponse{}, nil
	}
	r.mu.Lock()
	resp := script[t]
	r.turn[agentID] = t + 1
	r.mu.Unlock()
	return resp, nil
}

// synthesizeForRole calls the registered RoleBehavior for this role, or
// falls back to the generic `synthesize` if none is registered. Wires the
// production-grade scripted-auto path where each role does its job. After
// the behavior returns, we walk its tool calls and substitute
// "<role>-AGENTID" placeholders with the actual spawned agent id via
// AgentResolver. That way behaviors can be role-symbolic but the runtime
// emits real recipient ids.
func (r *Runtime) synthesizeForRole(agentID, role string, req runtime.LLMRequest) *runtime.LLMResponse {
	r.mu.Lock()
	beh, ok := r.RoleBehaviors[role]
	resolver := r.AgentResolver
	r.mu.Unlock()
	var resp *runtime.LLMResponse
	if ok && beh != nil {
		resp = beh(agentID, role, &req)
	}
	if resp == nil {
		resp = r.synthesize(agentID, req)
	}
	if resolver != nil {
		for i, tc := range resp.ToolCalls {
			if tc.Name != "send_message" {
				continue
			}
			args := tc.Args
			to, _ := args["to"].(string)
			if !strings.HasSuffix(to, "-AGENTID") {
				continue
			}
			roleHint := strings.TrimSuffix(to, "-AGENTID")
			if real := resolver(roleHint); real != "" {
				args["to"] = real
				resp.ToolCalls[i].Args = args
			}
		}
	}
	return resp
}

// synthesize produces a "did the work" response for the given incoming
// message. For delegates, replies with a report to the sender. For other
// types, returns a no-op response. Tokens are nominal so cost ceilings can
// still fire.
func (r *Runtime) synthesize(agentID string, req runtime.LLMRequest) *runtime.LLMResponse {
	if req.IncomingMessage == nil {
		return &runtime.LLMResponse{
			Text:   "[scripted-auto] no incoming message",
			Tokens: runtime.TokenUsage{Prompt: 50, Completion: 25},
		}
	}
	in := req.IncomingMessage
	resp := &runtime.LLMResponse{
		Text:   "[scripted-auto] handled " + string(in.Type),
		Tokens: runtime.TokenUsage{Prompt: 100, Completion: 50},
	}
	switch in.Type {
	case "delegate":
		// Reply with a report so the parent task can complete.
		resp.ToolCalls = []runtime.ToolCall{
			{Name: "send_message", Args: map[string]any{
				"to":      in.From,
				"type":    "report",
				"task_id": in.TaskID,
				"ttl_ms":  60000,
				"id":      "msg_auto_report_" + agentID + "_" + in.ID,
				"run_id":  in.RunID,
				"in_reply_to": in.ID,
				"payload": map[string]any{
					"intent":  "auto-completed by scripted runtime",
					"expects": "none",
				},
			}},
		}
	case "steering", "interrupt":
		// Ack so the orchestrator doesn't escalate.
		resp.ToolCalls = []runtime.ToolCall{
			{Name: "send_message", Args: map[string]any{
				"to":      in.From,
				"type":    "ack",
				"task_id": in.TaskID,
				"ttl_ms":  10000,
				"id":      "msg_auto_ack_" + agentID + "_" + in.ID,
				"run_id":  in.RunID,
				"in_reply_to": in.ID,
				"payload": map[string]any{"expects": "none"},
			}},
		}
	}
	return resp
}

// Terminate marks the agent terminated.
func (r *Runtime) Terminate(ctx context.Context, agentID, reason string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.status[agentID] = "terminated"
	return nil
}

// Health returns the agent's runtime status.
func (r *Runtime) Health(ctx context.Context, agentID string) (string, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if s, ok := r.status[agentID]; ok {
		return s, nil
	}
	return "", runtime.ErrUnknownAgent
}