const std = @import("std");
const zenai = @import("zenai");
const lp = @import("lightpanda");

const log = lp.log;
const Config = lp.Config;
const App = @import("../App.zig");
const ToolExecutor = @import("ToolExecutor.zig");
const Terminal = @import("Terminal.zig");
const Command = @import("Command.zig");
const CommandExecutor = @import("CommandExecutor.zig");
const Recorder = @import("Recorder.zig");
const Verifier = @import("Verifier.zig");
const SlashCommand = @import("SlashCommand.zig");
const script = lp.script;

const Self = @This();

const default_system_prompt = script.mcp_driver_guidance ++
    \\
    \\Agent-specific behavior:
    \\- Call a tool for every browser action. NEVER claim you performed an
    \\  action, visited a page, or saw content without actually calling the
    \\  corresponding tool. If a task needs a capability Lightpanda lacks
    \\  (images, PDFs, audio), say so honestly rather than improvising.
    \\- Be decisive and concise. Prefer few, well-chosen tool calls over many
    \\  probes. If extraction repeatedly fails or the site errors, commit to a
    \\  best-effort answer rather than thrashing.
    \\- An honest "the site blocked access" beats a fabricated answer every time.
    \\- If the user asks for account-scoped information (their karma, profile,
    \\  history, inbox, dashboard, settings, etc.) and the page shows you are
    \\  not signed in, attempt to log in proactively before reporting that the
    \\  data is unavailable. Find the login link or form on the current page
    \\  (interactiveElements or findElement), dismiss any cookie banner first,
    \\  then fill the username field with $LP_USERNAME and the password field
    \\  with $LP_PASSWORD and submit. Only fall back to "I couldn't access X"
    \\  if the form is missing or the credentials are rejected — and say which.
;

const self_heal_prompt_prefix =
    \\A PandaScript command failed during replay. The command that failed was:
    \\
;

const self_heal_prompt_page_state =
    \\
    \\The current page URL is:
    \\
;

const self_heal_prompt_instructions =
    \\
    \\IMPORTANT:
    \\- Do NOT navigate away from the current page. The page is already loaded and
    \\  contains the element you need — the selector just needs to be fixed.
    \\- Use the tree or interactiveElements tools WITHOUT a url parameter to inspect
    \\  the current page, find the correct selector, and execute the equivalent action.
    \\- If the action is blocked by a popup, cookie banner, or surprise modal,
    \\  handle it first (e.g., click "Accept") before executing the fixed command.
    \\- ONLY fix the failed command and handle immediate blockers. STOP immediately
    \\  once the intent of the original command is achieved.
    \\  The script will continue executing the remaining commands after the heal.
;

const login_prompt =
    \\Find the login form on the current page. Fill in the credentials using
    \\environment variables (look for $LP_EMAIL or $LP_USERNAME for the username
    \\field, and $LP_PASSWORD for the password field). Handle any cookie banners
    \\or popups first, then submit the form by clicking its submit button or
    \\pressing Enter in a filled field — there is no dedicated submit tool.
;

const accept_cookies_prompt =
    \\Find and dismiss the cookie consent banner on the current page.
    \\Look for "Accept", "Accept All", "I agree", or similar buttons and click them.
;

const synthesis_prompt =
    \\You have used your tool budget or cannot finish the exploration.
    \\Give your best final answer NOW based ONLY on what you actually observed
    \\via tool calls in this conversation. Do NOT fall back to prior knowledge —
    \\if your snapshots show only cookie banners, 403/access-denied pages,
    \\blocked search results, or empty bodies, say that explicitly
    \\(e.g. "the page was blocked by a cookie wall and I could not extract X").
    \\Do not invent details that are not visible in the tool outputs above.
    \\Do not call any more tools.
    \\Respond with ONLY the answer — one word, one number, one short phrase,
    \\or a brief honest explanation of why the page could not be read.
    \\No prefix, no markdown.
;

allocator: std.mem.Allocator,
ai_client: ?zenai.provider.Client,
tool_executor: *ToolExecutor,
terminal: Terminal,
cmd_executor: CommandExecutor,
verifier: Verifier,
recorder: Recorder,
messages: std.ArrayList(zenai.provider.Message),
message_arena: std.heap.ArenaAllocator,
model: []const u8,
system_prompt: []const u8,
script_file: ?[]const u8,
self_heal: bool,
interactive: bool,
one_shot_task: ?[]const u8,
one_shot_attachments: ?[]const []const u8,
slash_schemas: []const SlashCommand.SchemaInfo,

pub fn init(allocator: std.mem.Allocator, app: *App, opts: Config.Agent) !*Self {
    const is_one_shot = opts.task != null;
    const will_repl = !is_one_shot and (opts.interactive or opts.script_file == null);
    const needs_llm = will_repl or is_one_shot;

    if (opts.provider == null) {
        const required_by: ?[]const u8 = if (opts.self_heal)
            "--self-heal requires --provider; drop one or add the other"
        else if (is_one_shot)
            "--task requires --provider"
        else
            null;
        if (required_by) |hint| {
            log.fatal(.app, "missing --provider", .{ .hint = hint });
            return error.MissingProvider;
        }
    }

    const api_key = try resolveApiKey(opts.provider, needs_llm);

    const tool_executor: *ToolExecutor = try .init(allocator, app);
    errdefer tool_executor.deinit();

    const self = try allocator.create(Self);
    errdefer allocator.destroy(self);

    const ai_client: ?zenai.provider.Client = if (api_key) |key| switch (opts.provider.?) {
        inline else => |tag| blk: {
            const ProviderClient = zenai.provider.Client;
            const ClientPtr = @FieldType(ProviderClient, @tagName(tag));
            const Client = @typeInfo(ClientPtr).pointer.child;
            const client = try allocator.create(Client);
            const url: ?[]const u8 = opts.base_url orelse if (tag == .ollama) "http://localhost:11434/v1" else null;
            client.* = Client.init(allocator, key, if (url) |u| .{ .base_url = u } else .{});
            break :blk @unionInit(ProviderClient, @tagName(tag), client);
        },
    } else null;
    errdefer if (ai_client) |c| switch (c) {
        inline else => |client| {
            client.deinit();
            allocator.destroy(client);
        },
    };

    const history_path: ?[:0]const u8 = if (will_repl) ".lp-history" else null;

    // `-i <file>` means "replay then grow this file"; a script path alone is
    // pure replay and must not be mutated.
    const recorder_path: ?[]const u8 = if (opts.interactive) opts.script_file else null;

    // Reuse the executor's schema arena: the parsed schemas in `tools` already
    // live there, and the cache must outlive only as long as those do.
    const slash_schemas: []const SlashCommand.SchemaInfo = if (will_repl)
        SlashCommand.buildSchemas(tool_executor.schemaAllocator(), tool_executor.tools) catch {
            log.fatal(.app, "failed to build slash schemas", .{});
            return error.SlashSchemaInitFailed;
        }
    else
        &.{};

    self.* = .{
        .allocator = allocator,
        .ai_client = ai_client,
        .tool_executor = tool_executor,
        .terminal = .init(allocator, history_path, opts.verbosity, will_repl),
        .cmd_executor = undefined,
        .verifier = .{ .session = tool_executor.session, .node_registry = &tool_executor.node_registry },
        .recorder = .init(allocator, recorder_path),
        .messages = .empty,
        .message_arena = .init(allocator),
        .model = if (opts.provider) |p| (opts.model orelse zenai.provider.defaultModel(p)) else "",
        .system_prompt = opts.system_prompt orelse default_system_prompt,
        .script_file = opts.script_file,
        .self_heal = opts.self_heal,
        .interactive = opts.interactive,
        .one_shot_task = opts.task,
        .one_shot_attachments = if (opts.task_attachments.items.len == 0) null else opts.task_attachments.items,
        .slash_schemas = slash_schemas,
    };

    self.cmd_executor = CommandExecutor.init(allocator, tool_executor, &self.terminal);

    Terminal.setSlashSchemas(slash_schemas);

    return self;
}

pub fn deinit(self: *Self) void {
    self.recorder.deinit();
    self.terminal.deinit();
    self.message_arena.deinit();
    self.messages.deinit(self.allocator);
    self.tool_executor.deinit();
    if (self.ai_client) |ai_client| {
        switch (ai_client) {
            inline else => |c| {
                c.deinit();
                self.allocator.destroy(c);
            },
        }
    }
    self.allocator.destroy(self);
}

/// Returns true on success.
pub fn run(self: *Self) bool {
    if (self.one_shot_task) |task| return self.runTurn(task, null, self.one_shot_attachments, "Request");
    if (self.script_file) |path| {
        const script_ok = self.runScript(path);
        if (!self.interactive) return script_ok;
    }
    self.runRepl();
    return true;
}

/// Final answer goes to stdout; errors go to stderr, so a caller can
/// pipe stdout to capture a clean answer.
fn runTurn(self: *Self, prompt: []const u8, record_comment: ?[]const u8, attachments: ?[]const []const u8, label: []const u8) bool {
    const text = self.processUserMessage(prompt, record_comment, attachments) catch |err| switch (err) {
        // buildUserMessageParts has already logged the detail.
        error.UnsupportedAttachment, error.AttachmentReadFailed => return false,
        else => {
            self.terminal.printErrorFmt("{s} failed: {s}", .{ label, @errorName(err) });
            return false;
        },
    };
    if (text) |t| self.terminal.printAssistant(t) else self.terminal.printInfo("(no response from model)");
    return true;
}

fn runRepl(self: *Self) void {
    self.terminal.printInfo("Lightpanda Agent (type '/quit' to exit)");
    self.terminal.printInfo("Tab completes/cycles through commands; the dim grey ghost shows the first match.");
    log.debug(.app, "tools loaded", .{ .count = self.tool_executor.tools.len });
    if (self.ai_client) |ai_client| {
        self.terminal.printInfoFmt("Provider: {s}, Model: {s}", .{ @tagName(std.meta.activeTag(ai_client)), self.model });
    } else {
        self.terminal.printInfo("Dumb REPL (no --provider) — PandaScript only. Pass --provider for natural-language, LOGIN, and ACCEPT_COOKIES.");
    }

    repl: while (true) {
        const line = self.terminal.readLine("> ") orelse break;
        defer self.terminal.freeLine(line);

        if (line.len == 0) continue;

        if (line[0] == '/') {
            if (self.handleSlash(line[1..])) break :repl;
            continue :repl;
        }

        const cmd = Command.parse(line);

        if (cmd.needsLlm() and self.ai_client == null) {
            self.terminal.printError("This command requires --provider. PandaScript commands (GOTO, CLICK, EXTRACT, ...) work without one.");
            continue;
        }

        switch (cmd) {
            .comment => continue :repl,
            .login => _ = self.runTurn(login_prompt, line, null, "LOGIN"),
            .accept_cookies => _ = self.runTurn(accept_cookies_prompt, line, null, "ACCEPT_COOKIES"),
            .natural_language => _ = self.runTurn(line, line, null, "Request"),
            else => {
                self.cmd_executor.execute(cmd);
                self.recorder.record(cmd);
            },
        }
    }

    self.terminal.printInfo("Goodbye!");
}

/// Handle a REPL line that started with `/`. Returns `true` if the user asked
/// to quit (`/quit`), `false` otherwise. All errors are printed and
/// swallowed — the REPL must not die from a malformed slash command.
fn handleSlash(self: *Self, body: []const u8) bool {
    const split = SlashCommand.splitNameRest(body) orelse {
        self.terminal.printError("Empty slash command. Try /help.");
        return false;
    };
    const name = split.name;
    const rest = split.rest;

    if (std.mem.eql(u8, name, "quit")) {
        return true;
    }
    if (std.mem.eql(u8, name, "help")) {
        self.printSlashHelp(rest);
        return false;
    }

    const schema = SlashCommand.findSchema(self.slash_schemas, name) orelse {
        self.printSlashParseError(error.UnknownTool, name);
        return false;
    };

    var arena: std.heap.ArenaAllocator = .init(self.allocator);
    defer arena.deinit();
    const aa = arena.allocator();

    const args_json = SlashCommand.parseArgs(aa, schema, rest) catch |err| {
        self.printSlashParseError(err, name);
        return false;
    };

    if (std.mem.eql(u8, schema.tool_name, @tagName(lp.tools.Action.eval))) {
        // callEval surfaces the is_error flag separately from the text;
        // tool_executor.call discards it.
        const eval_script = extractEvalScript(aa, args_json) catch {
            self.terminal.printError("eval requires a `script` argument.");
            return false;
        };
        const result = self.tool_executor.callEval(aa, eval_script);
        if (result.is_error) {
            self.terminal.printErrorFmt("eval: {s}", .{result.text});
        } else {
            self.terminal.printToolResult(schema.tool_name, result.text);
        }
        return false;
    }

    const result = self.tool_executor.call(aa, schema.tool_name, args_json) catch |err| {
        self.terminal.printErrorFmt("{s}: {s}", .{ schema.tool_name, @errorName(err) });
        return false;
    };
    self.terminal.printToolResult(schema.tool_name, result);
    return false;
}

fn printSlashHelp(self: *Self, target: []const u8) void {
    if (target.len == 0) {
        self.terminal.printInfo("Slash commands (no LLM, REPL only):");
        for (self.slash_schemas) |s| {
            const summary = firstSentence(s.description);
            self.terminal.printInfoFmt("  /{s} — {s}", .{ s.tool_name, summary });
        }
        self.terminal.printInfo("Meta: /help [name], /quit");
        return;
    }
    const lookup = if (target[0] == '/') target[1..] else target;
    const schema = SlashCommand.findSchema(self.slash_schemas, lookup) orelse {
        self.terminal.printErrorFmt("unknown tool: {s}", .{lookup});
        return;
    };
    self.terminal.printInfoFmt("/{s} — {s}", .{ schema.tool_name, schema.description });
    self.terminal.printInfoFmt("schema: {s}", .{schema.input_schema_raw});
}

fn printSlashParseError(self: *Self, err: SlashCommand.ParseError, name: []const u8) void {
    switch (err) {
        error.UnknownTool => self.terminal.printErrorFmt("unknown tool '{s}'. Try /help.", .{name}),
        error.MissingName => self.terminal.printError("missing tool name. Try /help."),
        error.MissingRequired => self.terminal.printErrorFmt("{s}: missing required argument. Try /help {s}.", .{ name, name }),
        error.MalformedKv => self.terminal.printErrorFmt("{s}: malformed key=value. Use key=value or {{json}}.", .{name}),
        error.PositionalNotAllowed => self.terminal.printErrorFmt("{s}: positional only works for tools with one required field. Use key=value.", .{name}),
        error.UnterminatedQuote => self.terminal.printErrorFmt("{s}: unterminated quote.", .{name}),
        error.OutOfMemory => self.terminal.printError("out of memory"),
    }
}

fn firstSentence(text: []const u8) []const u8 {
    // Sentence boundary = period followed by whitespace (or end of string).
    // Plain "." is too aggressive — descriptions reference "console.log",
    // "JSON-LD, OpenGraph, etc.", and similar abbreviations.
    var i: usize = 0;
    while (std.mem.indexOfScalarPos(u8, text, i, '.')) |idx| : (i = idx + 1) {
        if (idx + 1 == text.len or std.ascii.isWhitespace(text[idx + 1])) return text[0..idx];
    }
    return text;
}

fn extractEvalScript(arena: std.mem.Allocator, args_json: []const u8) ![]const u8 {
    if (args_json.len == 0) return error.MissingScript;
    const parsed = std.json.parseFromSliceLeaky(struct { script: []const u8 }, arena, args_json, .{ .ignore_unknown_fields = true }) catch return error.MissingScript;
    return parsed.script;
}

const Replacement = script.Replacement;

fn runScript(self: *Self, path: []const u8) bool {
    self.terminal.printInfoFmt("Running script: {s}", .{path});

    var script_arena: std.heap.ArenaAllocator = .init(self.allocator);
    defer script_arena.deinit();
    const sa = script_arena.allocator();

    const content = std.fs.cwd().readFileAlloc(sa, path, 10 * 1024 * 1024) catch |err| {
        self.terminal.printErrorFmt("Failed to read script '{s}': {s}", .{ path, @errorName(err) });
        return false;
    };

    var iter: Command.ScriptIterator = .init(sa, content);
    var last_comment: ?[]const u8 = null;
    var replacements: std.ArrayList(Replacement) = .empty;

    while (iter.next()) |entry| {
        switch (entry.command) {
            .comment => {
                // Recorded scripts prefix LLM-generated commands with the
                // natural-language prompt that produced them; keep the
                // last one around so self-heal can use it as context.
                if (entry.raw_line.len > 2 and entry.raw_line[0] == '#') {
                    last_comment = std.mem.trim(u8, entry.raw_line[1..], &std.ascii.whitespace);
                }
                continue;
            },
            .natural_language => {
                self.terminal.printErrorFmt("line {d}: unrecognized command: {s}", .{ entry.line_num, entry.raw_line });
                self.flushReplacements(path, content, replacements.items);
                return false;
            },
            .login, .accept_cookies => {
                if (self.ai_client == null) {
                    self.terminal.printErrorFmt("line {d}: {s} requires --provider", .{
                        entry.line_num,
                        entry.raw_line,
                    });
                    self.flushReplacements(path, content, replacements.items);
                    return false;
                }
                const prompt = if (entry.command == .login) login_prompt else accept_cookies_prompt;
                const text = self.processUserMessage(prompt, null, null) catch |err| {
                    self.terminal.printErrorFmt("line {d}: {s} failed: {s}", .{
                        entry.line_num,
                        entry.raw_line,
                        @errorName(err),
                    });
                    self.flushReplacements(path, content, replacements.items);
                    return false;
                };
                if (text) |t| self.terminal.printAssistant(t);
            },
            else => {
                self.terminal.printInfoFmt("[{d}] {s}", .{ entry.line_num, entry.raw_line });
                switch (self.runActionEntry(sa, entry, last_comment)) {
                    .ok => {},
                    .healed => |r| replacements.append(sa, r) catch |err| {
                        self.terminal.printErrorFmt(
                            "line {d}: out of memory recording heal: {s} (script left unchanged)",
                            .{ entry.line_num, @errorName(err) },
                        );
                        return false;
                    },
                    .fail => {
                        self.flushReplacements(path, content, replacements.items);
                        return false;
                    },
                }
            },
        }
    }

    self.flushReplacements(path, content, replacements.items);
    self.terminal.printInfo("Script completed.");
    return true;
}

const ActionOutcome = union(enum) {
    ok,
    healed: Replacement,
    /// The per-line error has already been printed; caller must not re-report.
    fail,
};

/// Execute one action-style script entry, including post-execution
/// verification, transient-failure retry, and LLM self-heal escalation.
fn runActionEntry(self: *Self, sa: std.mem.Allocator, entry: Command.ScriptIterator.Entry, last_comment: ?[]const u8) ActionOutcome {
    var cmd_arena: std.heap.ArenaAllocator = .init(self.allocator);
    defer cmd_arena.deinit();
    const ca = cmd_arena.allocator();

    const result = self.cmd_executor.executeWithResult(ca, entry.command);
    self.cmd_executor.printResult(entry.command, result);

    const verification = if (!result.failed and self.self_heal)
        self.verifier.verify(ca, entry.command)
    else
        Verifier.VerifyResult{ .result = .passed };

    if (!result.failed and verification.result != .failed) return .ok;

    if (self.self_heal and self.ai_client != null) {
        // Verification-only failures often resolve with a brief wait
        // (animations, lazy-load); skip the LLM round-trip when they do.
        if (!result.failed and isRetryable(entry.command) and self.retryCommand(ca, entry.command)) {
            return .ok;
        }

        const msg = if (result.failed)
            "Command failed, attempting self-healing..."
        else
            "Command succeeded but verification failed, attempting self-healing...";
        self.terminal.printInfo(msg);

        if (self.attemptSelfHeal(sa, entry.raw_line, verification.reason, last_comment)) |healed_cmds| {
            const replacement = script.formatHealReplacement(sa, entry.raw_span, entry.raw_line, healed_cmds) catch |err| {
                self.terminal.printErrorFmt(
                    "line {d}: failed to record heal: {s} (script left unchanged)",
                    .{ entry.line_num, @errorName(err) },
                );
                return .fail;
            };
            return .{ .healed = replacement };
        }
    }
    self.terminal.printErrorFmt("line {d}: command failed: {s}", .{
        entry.line_num,
        entry.raw_line,
    });
    return .fail;
}

/// Re-run a verification-failed command with bounded backoff. Returns true
/// once both execution and verification pass, false after 3 attempts.
fn retryCommand(self: *Self, ca: std.mem.Allocator, cmd: Command.Command) bool {
    for (0..3) |i| {
        std.Thread.sleep((500 + i * 250) * std.time.ns_per_ms);
        self.terminal.printInfo("Retrying command...");
        const retry_result = self.cmd_executor.executeWithResult(ca, cmd);
        if (retry_result.failed) continue;
        if (self.verifier.verify(ca, cmd).result == .failed) continue;
        self.cmd_executor.printResult(cmd, retry_result);
        return true;
    }
    return false;
}

fn flushReplacements(self: *Self, path: []const u8, content: []const u8, replacements: []const Replacement) void {
    if (replacements.len == 0) return;
    script.writeAtomic(self.allocator, std.fs.cwd(), path, content, replacements) catch |err| {
        self.terminal.printErrorFmt(
            "Failed to update script {s}: {s} (script left unchanged)",
            .{ path, @errorName(err) },
        );
        return;
    };
    self.terminal.printInfoFmt(
        "Script updated with {d} healed command(s); backup at {s}.bak",
        .{ replacements.len, path },
    );
}

fn isRetryable(cmd: Command.Command) bool {
    return switch (cmd) {
        .type_cmd, .check, .select => true,
        else => false,
    };
}

const self_heal_max_attempts = 3;

fn ensureSystemPrompt(self: *Self) !void {
    if (self.messages.items.len == 0) {
        try self.messages.append(self.allocator, .{
            .role = .system,
            .content = self.system_prompt,
        });
    }
}

// Drop older turns once `prune_high` is hit; survivors are deep-copied so
// the old arena (which still pins dropped strings) can be released.
const prune_high = 30;
const prune_keep = 20;

fn pruneMessages(self: *Self) void {
    const msgs = self.messages.items;
    if (msgs.len <= prune_high) return;

    const tail_start = zenai.provider.safeTruncationStart(msgs, msgs.len - prune_keep) orelse return;

    // Dupe the kept tail into a scratch slice in the new arena first. Only
    // mutate self.messages once every dupe has succeeded — otherwise a
    // partial failure would leave self.messages.items[1..] pointing into
    // the freed `new_arena`.
    var new_arena: std.heap.ArenaAllocator = .init(self.allocator);
    const duped = zenai.provider.dupeMessages(new_arena.allocator(), msgs[tail_start..]) catch {
        new_arena.deinit();
        return;
    };

    // System prompt at index 0 lives outside the arena and is preserved.
    @memcpy(self.messages.items[1..][0..duped.len], duped);
    self.messages.shrinkRetainingCapacity(1 + duped.len);
    self.message_arena.deinit();
    self.message_arena = new_arena;
}

/// Self-heal must only patch the current page; navigation and arbitrary
/// scripting are blocked even if the model emits them via `goto` / `eval`.
/// docs/agent.md guarantees "no navigation away from the current page".
fn isHealAllowed(cmd: Command.Command) bool {
    return switch (cmd) {
        .goto, .eval_js => false,
        else => true,
    };
}

/// Runs a single LLM turn, captures the commands it called without recording
/// them — so the caller can splice healed commands into the script directly.
fn runHealTurn(self: *Self, arena: std.mem.Allocator, prompt: []const u8) ![]Command.Command {
    const ma = self.message_arena.allocator();

    try self.ensureSystemPrompt();

    try self.messages.append(self.allocator, .{
        .role = .user,
        .content = try ma.dupe(u8, prompt),
    });

    const provider_client = self.ai_client orelse return error.NoAiClient;

    self.terminal.spinner.start();
    var result = provider_client.runTools(
        self.model,
        &self.messages,
        self.allocator,
        ma,
        .{ .context = @ptrCast(self), .callFn = &handleToolCall },
        .{
            .tools = self.tool_executor.tools,
            .max_tool_calls = 4,
            .max_tokens = 4096,
            .tool_choice = .auto,
        },
    ) catch |err| {
        self.terminal.spinner.cancel();
        log.err(.app, "AI API error", .{ .err = err });
        return error.ApiError;
    };
    self.terminal.spinner.stop();
    defer result.deinit();

    var cmds: std.ArrayList(Command.Command) = .empty;
    for (result.tool_calls_made) |tc| {
        if (tc.is_error) continue;
        const cmd = Command.fromToolCall(ma, tc.name, tc.arguments) orelse continue;
        if (!isHealAllowed(cmd)) {
            self.terminal.printInfoFmt(
                "self-heal: ignoring {s} (navigation and eval are not allowed during heal)",
                .{tc.name},
            );
            continue;
        }
        try cmds.append(arena, cmd);
    }

    if (result.text) |text| {
        self.terminal.printAssistant(text);
    }

    return cmds.toOwnedSlice(arena);
}

fn attemptSelfHeal(self: *Self, arena: std.mem.Allocator, failed_command: []const u8, verify_context: ?[]const u8, context_comment: ?[]const u8) ?[]Command.Command {
    const ha = self.message_arena.allocator();

    const verify_section = if (verify_context) |ctx|
        std.fmt.allocPrint(ha, "\n\nVerification detected a problem:\n{s}", .{ctx}) catch return null
    else
        "";

    const comment_section = if (context_comment) |c|
        std.fmt.allocPrint(ha, "\n\nThe original user request that generated this command was:\n{s}", .{c}) catch return null
    else
        "";

    const prompt = std.fmt.allocPrint(ha, "{s}{s}{s}{s}{s}{s}{s}", .{
        self_heal_prompt_prefix,
        failed_command,
        self_heal_prompt_page_state,
        self.tool_executor.getCurrentUrl(),
        comment_section,
        verify_section,
        self_heal_prompt_instructions,
    }) catch return null;

    // Save message count so we can roll back between attempts — each failed
    // heal turn would otherwise accumulate in context, confusing the next try.
    const msg_baseline = self.messages.items.len;

    var attempt: u8 = 0;
    while (attempt < self_heal_max_attempts) : (attempt += 1) {
        const cmds = self.runHealTurn(arena, prompt) catch |err| {
            self.terminal.printErrorFmt("self-heal attempt {d}/{d} failed: {s}", .{
                attempt + 1,
                self_heal_max_attempts,
                @errorName(err),
            });
            self.messages.shrinkRetainingCapacity(msg_baseline);
            continue;
        };
        if (cmds.len > 0) {
            self.pruneMessages();
            return cmds;
        }
        self.messages.shrinkRetainingCapacity(msg_baseline);
    }
    return null;
}

/// Returned text lives in `message_arena`, so it's only valid until the
/// next prune. `null` means the model emitted nothing even after the
/// synthesis turn.
fn processUserMessage(self: *Self, user_input: []const u8, record_comment: ?[]const u8, attachments: ?[]const []const u8) !?[]const u8 {
    const ma = self.message_arena.allocator();

    try self.ensureSystemPrompt();

    // Attachments only ride on the very first user turn (just after the
    // system prompt) — wired into the message's rich `parts`.
    const turn_attachments: ?[]const []const u8 =
        if (self.messages.items.len == 1) attachments else null;

    if (turn_attachments) |paths| {
        const parts = try buildUserMessageParts(self, ma, user_input, paths);
        try self.messages.append(self.allocator, .{
            .role = .user,
            .parts = parts,
        });
    } else {
        try self.messages.append(self.allocator, .{
            .role = .user,
            .content = try ma.dupe(u8, user_input),
        });
    }

    const provider_client = self.ai_client orelse return error.NoAiClient;

    self.terminal.spinner.start();
    var result = provider_client.runTools(
        self.model,
        &self.messages,
        self.allocator,
        ma,
        .{ .context = @ptrCast(self), .callFn = &handleToolCall },
        .{
            .tools = self.tool_executor.tools,
            .max_turns = 30,
            // Safety net; max_turns is the primary terminal.
            .max_tool_calls = 200,
            .max_tokens = 4096,
            .tool_choice = .auto,
            // Cap per-turn reasoning so thinking models don't burn
            // minutes per turn. Ignored by non-thinking models.
            .thinking_budget = 2048,
        },
    ) catch |err| {
        self.terminal.spinner.cancel();
        log.err(.app, "AI API error", .{ .err = err });
        return error.ApiError;
    };
    self.terminal.spinner.stop();
    defer result.deinit();

    if (self.recorder.file != null) {
        var recorded_any = false;
        for (result.tool_calls_made) |tc| {
            if (tc.is_error) continue;
            const cmd = Command.fromToolCall(ma, tc.name, tc.arguments) orelse continue;
            if (!recorded_any) {
                if (record_comment) |c| self.recorder.recordComment(c);
                recorded_any = true;
            }
            self.recorder.record(cmd);
        }
    }

    // `result.text` and `synth.text` are owned by their RunToolsResult arenas,
    // which are deinited at the end of this function. Dupe into the agent's
    // `message_arena` so the returned slice outlives those arenas.
    const final_text: ?[]const u8 = blk: {
        if (result.text) |text| break :blk try ma.dupe(u8, text);

        // Tool loop ended without a final text — force one more turn that
        // forbids tools and pretraining fallback. Without this, models
        // confabulate answers when the page was blocked or empty.
        log.info(.app, "synthesizing final answer", .{});
        try self.messages.append(self.allocator, .{
            .role = .user,
            .content = try ma.dupe(u8, synthesis_prompt),
        });

        var synth = provider_client.runTools(
            self.model,
            &self.messages,
            self.allocator,
            ma,
            .{ .context = @ptrCast(self), .callFn = &handleToolCall },
            .{
                .tools = self.tool_executor.tools,
                .max_turns = 1,
                .max_tokens = 4096,
                .tool_choice = .none,
                // Cap thinking on the finalize turn. Fully disabling it (0)
                // leaves reasoning-heavy tasks with no answer at all; letting
                // it run unbounded lets models fill the turn with thoughts
                // and emit nothing as the final text. 512 tokens is enough
                // for the model to pick its answer but not to freewheel.
                .thinking_budget = 512,
            },
        ) catch |err| {
            log.err(.app, "AI synthesis error", .{ .err = err });
            break :blk null;
        };
        defer synth.deinit();

        break :blk if (synth.text) |text| try ma.dupe(u8, text) else null;
    };

    self.pruneMessages();
    return final_text;
}

/// Build a `parts`-based user message when `--task-attachment` was given.
/// Text-ish files are inlined into the text prefix (surrounded by clear
/// markers); binary files (image/audio/pdf) are base64-encoded and sent as
/// provider inline-data parts. Unknown extensions error out so the caller
/// fails loudly instead of silently dropping the attachment.
fn buildUserMessageParts(
    self: *Self,
    ma: std.mem.Allocator,
    user_input: []const u8,
    paths: []const []const u8,
) ![]const zenai.provider.ContentPart {
    var text_prefix: std.ArrayList(u8) = .empty;
    var inline_parts: std.ArrayList(zenai.provider.ContentPart) = .empty;

    for (paths) |path| {
        const mime = zenai.provider.inferInlineMimeType(path) orelse {
            log.err(.app, "unsupported attachment", .{ .path = path });
            self.terminal.printErrorFmt("unsupported attachment type: {s}", .{path});
            return error.UnsupportedAttachment;
        };

        if (std.mem.startsWith(u8, mime, "text/")) {
            const bytes = std.fs.cwd().readFileAlloc(ma, path, 512 * 1024) catch |err| {
                log.err(.app, "read attachment failed", .{ .path = path, .err = err });
                self.terminal.printErrorFmt("could not read attachment: {s}", .{path});
                return error.AttachmentReadFailed;
            };
            try text_prefix.writer(ma).print(
                "[Attached file: {s}]\n{s}\n[End of attachment]\n\n",
                .{ path, bytes },
            );
        } else {
            const raw = std.fs.cwd().readFileAlloc(ma, path, 20 * 1024 * 1024) catch |err| {
                log.err(.app, "read attachment failed", .{ .path = path, .err = err });
                self.terminal.printErrorFmt("could not read attachment: {s}", .{path});
                return error.AttachmentReadFailed;
            };
            const b64_len = std.base64.standard.Encoder.calcSize(raw.len);
            const b64 = try ma.alloc(u8, b64_len);
            _ = std.base64.standard.Encoder.encode(b64, raw);
            try inline_parts.append(ma, .{ .image = .{
                .data = b64,
                .mime_type = try ma.dupe(u8, mime),
            } });
        }
    }

    var parts: std.ArrayList(zenai.provider.ContentPart) = .empty;
    try text_prefix.appendSlice(ma, user_input);
    try parts.append(ma, .{ .text = try text_prefix.toOwnedSlice(ma) });
    for (inline_parts.items) |p| try parts.append(ma, p);
    return parts.toOwnedSlice(ma);
}

// Cap per-call tool output so heavy pages don't balloon the message arena
// (and the next request body) without bound.
const tool_output_max_bytes: usize = 1 * 1024 * 1024;

fn capToolOutput(allocator: std.mem.Allocator, output: []const u8) []const u8 {
    if (output.len <= tool_output_max_bytes) return output;
    const prefix = output[0..tool_output_max_bytes];
    return std.fmt.allocPrint(
        allocator,
        "{s}\n...[truncated, original {d} bytes]",
        .{ prefix, output.len },
    ) catch prefix;
}

fn handleToolCall(ctx: *anyopaque, allocator: std.mem.Allocator, tool_name: []const u8, arguments: []const u8) zenai.provider.Client.ToolHandler.Result {
    const self: *Self = @ptrCast(@alignCast(ctx));
    self.terminal.spinner.setTool(tool_name, arguments);
    defer self.terminal.spinner.setThinking();
    if (self.tool_executor.call(allocator, tool_name, arguments)) |output| {
        const capped = capToolOutput(allocator, output);
        self.terminal.agentToolDone(tool_name, arguments, true);
        // Only `high` keeps the per-call body line — benchmark harness parses it.
        if (self.terminal.verbosity == .high) self.terminal.printToolResult(tool_name, capped);
        return .{ .content = capped };
    } else |err| {
        const msg = std.fmt.allocPrint(allocator, "Error: {s}", .{@errorName(err)}) catch "Error: tool execution failed";
        // Errors go back to the model so it can self-correct. The red
        // bullet (per-line) or red spinner label is the user-facing
        // failure signal; we only print the body at `high` (harness).
        self.terminal.agentToolDone(tool_name, arguments, false);
        if (self.terminal.verbosity == .high) self.terminal.printToolResult(tool_name, msg);
        return .{ .content = msg, .is_error = true };
    }
}

/// An API key is only required when an LLM turn will actually run. Without a
/// provider, no key is needed.
fn resolveApiKey(provider: ?Config.AiProvider, needs_llm: bool) !?[:0]const u8 {
    const p = provider orelse return null;
    if (zenai.provider.envApiKey(p)) |key| return key;
    if (!needs_llm) return null;
    log.fatal(.app, "missing API key", .{
        .hint = "Set ANTHROPIC_API_KEY, OPENAI_API_KEY, or GOOGLE_API_KEY",
    });
    return error.MissingApiKey;
}

// --- Tests ---

test "isHealAllowed: blocks goto and eval_js, allows page-local commands" {
    try std.testing.expect(!isHealAllowed(.{ .goto = "https://x" }));
    try std.testing.expect(!isHealAllowed(.{ .eval_js = "alert(1)" }));

    try std.testing.expect(isHealAllowed(.{ .click = ".btn" }));
    try std.testing.expect(isHealAllowed(.{ .hover = ".menu" }));
    try std.testing.expect(isHealAllowed(.{ .wait = ".loaded" }));
    try std.testing.expect(isHealAllowed(.{ .type_cmd = .{ .selector = "#u", .value = "x" } }));
    try std.testing.expect(isHealAllowed(.{ .select = .{ .selector = "#s", .value = "x" } }));
    try std.testing.expect(isHealAllowed(.{ .check = .{ .selector = "#c", .checked = true } }));
    try std.testing.expect(isHealAllowed(.{ .scroll = .{ .x = 0, .y = 100 } }));
}
