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 Self = @This();

const default_system_prompt =
    \\You are a web browsing assistant powered by the Lightpanda browser.
    \\You can navigate to websites, read their content, interact with forms,
    \\click links, and extract information.
    \\
    \\When helping the user, navigate to relevant pages and extract information.
    \\Use the semanticTree or interactiveElements tools to understand page structure
    \\before clicking or filling forms. Be concise in your responses.
    \\
    \\IMPORTANT RULES:
    \\- NEVER use backendNodeId with click, fill, hover, selectOption, or setChecked.
    \\  Always use a CSS selector. Use findElement to resolve a description into a
    \\  CSS selector if needed.
    \\  Example: click with selector "#login-btn", NOT with backendNodeId 42.
    \\- Use specific CSS selectors that uniquely identify elements. Include
    \\  distinguishing attributes like value, name, or position to avoid ambiguity.
    \\  Example: input[type="submit"][value="login"], NOT just input[type="submit"].
    \\- When filling credentials, pass environment variable references like
    \\  $LP_USERNAME and $LP_PASSWORD directly as the value — they will be
    \\  resolved automatically. Do NOT use getEnv to resolve them first.
;

const self_heal_prompt_prefix =
    \\A Pandascript command failed during replay. The original intent was:
    \\
;

const self_heal_prompt_suffix =
    \\
    \\The command that failed was:
    \\
;

const self_heal_prompt_page_state =
    \\
    \\Please analyze the current page state and execute the equivalent action.
    \\Use the available tools to accomplish the original intent.
;

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 login form.
;

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.
;

allocator: std.mem.Allocator,
ai_client: ?zenai.provider.Client,
tool_executor: *ToolExecutor,
terminal: Terminal,
cmd_executor: CommandExecutor,
recorder: Recorder,
messages: std.ArrayListUnmanaged(zenai.provider.Message),
message_arena: std.heap.ArenaAllocator,
tools: []const zenai.provider.Tool,
model: []const u8,
system_prompt: []const u8,
script_file: ?[]const u8,
self_heal: bool,

pub fn init(allocator: std.mem.Allocator, app: *App, opts: Config.Agent) !*Self {
    const is_script_mode = opts.script_file != null and !opts.save;

    // API key is only required for REPL mode and self-healing
    const api_key: ?[:0]const u8 = getEnvApiKey(opts.provider) orelse if (!is_script_mode) {
        log.fatal(.app, "missing API key", .{
            .hint = "Set ANTHROPIC_API_KEY, OPENAI_API_KEY, or GOOGLE_API_KEY",
        });
        return error.MissingApiKey;
    } else null;

    const tool_executor = try ToolExecutor.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;

    const tools = tool_executor.getTools() catch {
        log.fatal(.app, "failed to initialize tools", .{});
        return error.ToolInitFailed;
    };

    // Persist REPL history in a cwd-relative `.lp-history`. Skipped in pure
    // replay mode (no REPL is opened).
    const history_path: ?[:0]const u8 = if (is_script_mode) null else ".lp-history";

    self.* = .{
        .allocator = allocator,
        .ai_client = ai_client,
        .tool_executor = tool_executor,
        .terminal = Terminal.init(history_path),
        .cmd_executor = undefined,
        .recorder = Recorder.init(if (opts.save) opts.script_file else null),
        .messages = .empty,
        .message_arena = std.heap.ArenaAllocator.init(allocator),
        .tools = tools,
        .model = opts.model orelse defaultModel(opts.provider),
        .system_prompt = opts.system_prompt orelse default_system_prompt,
        .script_file = if (!opts.save) opts.script_file else null,
        .self_heal = opts.self_heal,
    };

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

    return self;
}

pub fn deinit(self: *Self) void {
    self.recorder.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, false if a script command failed.
pub fn run(self: *Self) bool {
    if (self.script_file) |path| {
        return self.runScript(path);
    }
    self.runRepl();
    return true;
}

fn runRepl(self: *Self) void {
    self.terminal.printInfo("Lightpanda Agent (type 'quit' to exit)");
    log.debug(.app, "tools loaded", .{ .count = self.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("Ready.");
    }

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

        if (line.len == 0) continue;

        const cmd = Command.parse(line);
        switch (cmd) {
            .exit => break,
            .comment => continue,
            .login => {
                self.processUserMessage(login_prompt, line) catch |err| {
                    self.printAllocError("LOGIN failed: {s}", .{@errorName(err)});
                };
            },
            .accept_cookies => {
                self.processUserMessage(accept_cookies_prompt, line) catch |err| {
                    self.printAllocError("ACCEPT_COOKIES failed: {s}", .{@errorName(err)});
                };
            },
            .natural_language => {
                if (std.mem.eql(u8, line, "quit")) break;

                self.processUserMessage(line, line) catch |err| {
                    self.printAllocError("Request failed: {s}", .{@errorName(err)});
                };
            },
            else => {
                self.cmd_executor.execute(cmd);
                self.recorder.record(cmd);
            },
        }
    }

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

fn printAllocError(self: *Self, comptime fmt: []const u8, args: anytype) void {
    self.terminal.printErrorFmt(fmt, args);
}

fn runScript(self: *Self, path: []const u8) bool {
    const file = std.fs.cwd().openFile(path, .{}) catch |err| {
        self.printAllocError("Failed to open script '{s}': {s}", .{ path, @errorName(err) });
        return false;
    };
    defer file.close();

    const content = file.readToEndAlloc(self.allocator, 10 * 1024 * 1024) catch |err| {
        self.printAllocError("Failed to read script: {s}", .{@errorName(err)});
        return false;
    };
    defer self.allocator.free(content);

    self.terminal.printInfoFmt("Running script: {s}", .{path});

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

    var iter = Command.ScriptIterator.init(content, script_arena.allocator());
    var last_intent: ?[]const u8 = null;

    while (iter.next()) |entry| {
        switch (entry.command) {
            .exit => {
                self.terminal.printInfo("EXIT — stopping script.");
                return true;
            },
            .comment => {
                if (std.mem.startsWith(u8, entry.raw_line, "# INTENT:")) {
                    last_intent = std.mem.trim(u8, entry.raw_line["# INTENT:".len..], &std.ascii.whitespace);
                }
                continue;
            },
            .natural_language => {
                self.printAllocError("line {d}: unrecognized command: {s}", .{ entry.line_num, entry.raw_line });
                return false;
            },
            .login, .accept_cookies => {
                if (self.ai_client == null) {
                    self.printAllocError("line {d}: {s} requires an API key for LLM resolution", .{
                        entry.line_num,
                        entry.raw_line,
                    });
                    return false;
                }
                const prompt = if (entry.command == .login) login_prompt else accept_cookies_prompt;
                self.processUserMessage(prompt, "") catch |err| {
                    self.printAllocError("line {d}: {s} failed: {s}", .{
                        entry.line_num,
                        entry.raw_line,
                        @errorName(err),
                    });
                    return false;
                };
            },
            else => {
                self.terminal.printInfoFmt("[{d}] {s}", .{ entry.line_num, entry.raw_line });

                var cmd_arena = std.heap.ArenaAllocator.init(self.allocator);
                defer cmd_arena.deinit();

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

                if (result.failed) {
                    if (self.self_heal and self.ai_client != null) {
                        self.terminal.printInfo("Command failed, attempting self-healing...");
                        if (self.attemptSelfHeal(last_intent, entry.raw_line)) {
                            continue;
                        }
                    }
                    self.printAllocError("line {d}: command failed: {s}", .{
                        entry.line_num,
                        entry.raw_line,
                    });
                    return false;
                }
            },
        }
    }

    self.terminal.printInfo("Script completed.");
    return true;
}

const self_heal_max_attempts = 3;

fn attemptSelfHeal(self: *Self, intent: ?[]const u8, failed_command: []const u8) bool {
    var heal_arena = std.heap.ArenaAllocator.init(self.allocator);
    defer heal_arena.deinit();
    const ha = heal_arena.allocator();

    const prompt = std.fmt.allocPrint(ha, "{s}{s}{s}{s}{s}", .{
        self_heal_prompt_prefix,
        intent orelse "(no recorded intent)",
        self_heal_prompt_suffix,
        failed_command,
        self_heal_prompt_page_state,
    }) catch return false;

    var attempt: u8 = 0;
    while (attempt < self_heal_max_attempts) : (attempt += 1) {
        self.processUserMessage(prompt, "") catch |err| {
            self.printAllocError("self-heal attempt {d}/{d} failed: {s}", .{
                attempt + 1,
                self_heal_max_attempts,
                @errorName(err),
            });
            continue;
        };
        return true;
    }
    return false;
}

fn processUserMessage(self: *Self, user_input: []const u8, record_comment: []const u8) !void {
    const ma = self.message_arena.allocator();

    if (self.messages.items.len == 0) {
        try self.messages.append(self.allocator, .{
            .role = .system,
            .content = self.system_prompt,
        });
    }

    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;

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

    var recorded_any = false;
    for (result.tool_calls_made) |tc| {
        if (!std.mem.startsWith(u8, tc.result, "Error:")) {
            if (toolCallToCommand(ma, tc.name, tc.arguments)) |cmd| {
                if (!recorded_any) {
                    if (record_comment.len > 0) self.recorder.recordComment(record_comment);
                    recorded_any = true;
                }
                self.recorder.record(cmd);
            }
        }
    }

    if (result.text) |text| {
        std.debug.print("\n", .{});
        self.terminal.printAssistant(text);
        std.debug.print("\n", .{});
    } else {
        self.terminal.printInfo("(no response from model)");
    }
}

fn handleToolCall(ctx: *anyopaque, allocator: std.mem.Allocator, tool_name: []const u8, arguments: []const u8) []const u8 {
    const self: *Self = @ptrCast(@alignCast(ctx));
    self.terminal.printToolCall(tool_name, arguments);
    const tool_result = self.tool_executor.call(allocator, tool_name, arguments) catch |err| blk: {
        break :blk std.fmt.allocPrint(allocator, "Error: {s}", .{@errorName(err)}) catch "Error: tool execution failed";
    };
    self.terminal.printToolResult(tool_name, tool_result);
    return tool_result;
}

fn toolCallToCommand(arena: std.mem.Allocator, tool_name: []const u8, arguments: []const u8) ?Command.Command {
    const action = std.meta.stringToEnum(lp.tools.Action, tool_name) orelse return null;
    const parsed = std.json.parseFromSlice(std.json.Value, arena, arguments, .{}) catch return null;
    const obj = switch (parsed.value) {
        .object => |o| o,
        else => return null,
    };

    const getString = struct {
        fn f(o: std.json.ObjectMap, key: []const u8) ?[]const u8 {
            return switch (o.get(key) orelse return null) {
                .string => |s| s,
                else => null,
            };
        }
    }.f;

    return switch (action) {
        .goto => .{ .goto = getString(obj, "url") orelse return null },
        .click => .{ .click = getString(obj, "selector") orelse return null },
        .hover => .{ .hover = getString(obj, "selector") orelse return null },
        .eval => .{ .eval_js = getString(obj, "script") orelse return null },
        .waitForSelector => .{ .wait = getString(obj, "selector") orelse return null },
        .fill => .{ .type_cmd = .{
            .selector = getString(obj, "selector") orelse return null,
            .value = getString(obj, "value") orelse return null,
        } },
        .selectOption => .{ .select = .{
            .selector = getString(obj, "selector") orelse return null,
            .value = getString(obj, "value") orelse return null,
        } },
        .setChecked => .{ .check = .{
            .selector = getString(obj, "selector") orelse return null,
            .checked = switch (obj.get("checked") orelse return null) {
                .bool => |b| b,
                else => return null,
            },
        } },
        .scroll => blk: {
            if (obj.get("backendNodeId") != null) break :blk null;
            const x: i32 = switch (obj.get("x") orelse std.json.Value{ .integer = 0 }) {
                .integer => |i| @intCast(i),
                else => 0,
            };
            const y: i32 = switch (obj.get("y") orelse std.json.Value{ .integer = 0 }) {
                .integer => |i| @intCast(i),
                else => 0,
            };
            break :blk .{ .scroll = .{ .x = x, .y = y } };
        },
        else => null,
    };
}

fn getEnvApiKey(provider_type: Config.AiProvider) ?[:0]const u8 {
    return switch (provider_type) {
        .anthropic => std.posix.getenv("ANTHROPIC_API_KEY"),
        .openai => std.posix.getenv("OPENAI_API_KEY"),
        .gemini => std.posix.getenv("GOOGLE_API_KEY") orelse std.posix.getenv("GEMINI_API_KEY"),
        .ollama => "ollama",
    };
}

fn defaultModel(provider_type: Config.AiProvider) []const u8 {
    return switch (provider_type) {
        .anthropic => "claude-haiku-4-5-20251001",
        .openai => "gpt-4o-mini",
        .gemini => "gemini-2.5-flash",
        .ollama => "gemma3",
    };
}