const std = @import("std");
const panto = @import("panto");
const lua_bridge = @import("lua_bridge.zig");
const lua_runtime = @import("lua_runtime.zig");
const extension_loader = @import("extension_loader.zig");
const panto_home = @import("panto_home.zig");
const luarocks_runtime = @import("luarocks_runtime.zig");
const self_exe = @import("self_exe.zig");
const subcommand = @import("subcommand.zig");
const session_paths = @import("session_paths.zig");
const models_toml = @import("models_toml.zig");

// Shorthand alias for the Lua C API. The bridge module owns the actual
// `@cImport`; we re-use it here so the smoke check uses identical types.
const lua = lua_bridge.c;

test {
    // Test contract: deliberate error-path tests should not produce visible
    // log output. Code logs at `.err` for genuine production failures and
    // `.warn` for expected-failure paths exercised by tests; the test
    // runner's logger gates on `std.testing.log_level`, which defaults to
    // `.warn`. Raising it to `.err` silences expected warnings without
    // changing production behavior. Anything that *should* be visible in a
    // passing test must use `std.debug.print` or assert via the testing API.
    std.testing.log_level = .err;

    std.testing.refAllDecls(@This());
    _ = lua_bridge;
    _ = lua_runtime;
    _ = extension_loader;
    _ = panto_home;
    _ = luarocks_runtime;
    _ = self_exe;
    _ = subcommand;
    _ = models_toml;
}

const Receiver = panto.provider.Receiver;
const ReceiverVTable = panto.provider.ReceiverVTable;
const ContentBlockType = panto.provider.ContentBlockType;
const MessageRole = panto.conversation.MessageRole;

/// Receiver that prints streaming deltas to stdout. Thinking blocks are
/// dimmed with ANSI escape codes; text blocks render plain.
///
/// Also captures one `?Usage` per assistant response during a turn. In
/// a tool-using turn the agent loop drives multiple streamStep calls;
/// each one ends with `onMessageComplete(msg, usage)`. We append the
/// usage — including `null` when the wire didn't deliver any — to
/// `per_message_usage` in order so `persistTurn` can pair each captured
/// usage with its corresponding assistant message.
const CLIReceiver = struct {
    stdout: *std.Io.Writer,
    file: *std.Io.File.Writer,
    allocator: std.mem.Allocator,

    /// One slot per assistant message completed during the current
    /// turn, in completion order. `null` means the provider did not
    /// report usage on the wire for that message (typical of
    /// OpenAI-compatible proxies that ignore `stream_options.include_usage`).
    per_message_usage: std.ArrayList(?panto.session_manager.Usage) = .empty,

    pub fn receiver(self: *CLIReceiver) Receiver {
        return .{ .ptr = self, .vtable = &vtable };
    }

    /// Reset usage state at the start of each turn.
    pub fn beginTurn(self: *CLIReceiver) void {
        self.per_message_usage.clearRetainingCapacity();
    }

    pub fn deinit(self: *CLIReceiver) void {
        self.per_message_usage.deinit(self.allocator);
    }

    const vtable: ReceiverVTable = .{
        .onMessageStart = onMessageStart,
        .onBlockStart = onBlockStart,
        .onToolDetails = onToolDetails,
        .onContentDelta = onContentDelta,
        .onBlockComplete = onBlockComplete,
        .onMessageComplete = onMessageComplete,
        .onError = onError,
    };

    /// The print-based CLI defers tool-name rendering to onBlockComplete
    /// to avoid cursor gymnastics (we can't go back and edit the prefix).
    /// Receivers backed by a TUI capable of in-place updates would render
    /// the name as soon as it's known here.
    fn onToolDetails(
        ptr: *anyopaque,
        index: usize,
        id: []const u8,
        name: []const u8,
    ) anyerror!void {
        _ = ptr;
        _ = index;
        _ = id;
        _ = name;
    }

    fn onMessageStart(ptr: *anyopaque, role: MessageRole) anyerror!void {
        _ = ptr;
        _ = role;
    }

    fn onBlockStart(
        ptr: *anyopaque,
        block_type: ContentBlockType,
        index: usize,
    ) anyerror!void {
        _ = index;
        const self: *CLIReceiver = @ptrCast(@alignCast(ptr));
        switch (block_type) {
            .Thinking => try self.stdout.writeAll("\x1b[2m[thinking] "),
            // Tool name is not known reliably at start time (OpenAI may
            // stream id/name across fragments). Open with a bare prefix;
            // the tool name lands at onBlockComplete from the assembled
            // ContentBlock.
            .ToolUse => try self.stdout.writeAll("\n\x1b[36mtool: \x1b[0m"),
            else => {},
        }
        try self.file.flush();
    }

    fn onContentDelta(ptr: *anyopaque, index: usize, delta: []const u8) anyerror!void {
        _ = index;
        const self: *CLIReceiver = @ptrCast(@alignCast(ptr));
        try self.stdout.writeAll(delta);
        try self.file.flush();
    }

    fn onBlockComplete(
        ptr: *anyopaque,
        index: usize,
        block: panto.conversation.ContentBlock,
    ) anyerror!void {
        _ = index;
        const self: *CLIReceiver = @ptrCast(@alignCast(ptr));
        switch (block) {
            .Thinking => try self.stdout.writeAll("\x1b[0m\n"),
            // Append the tool name now that we know it for certain.
            .ToolUse => |tu| try self.stdout.print("\x1b[36m : ({s})\x1b[0m\n", .{tu.name}),
            else => {},
        }
        try self.file.flush();
    }

    fn onMessageComplete(
        ptr: *anyopaque,
        message: panto.conversation.Message,
        usage: ?panto.session_manager.Usage,
    ) anyerror!void {
        const self: *CLIReceiver = @ptrCast(@alignCast(ptr));
        // Only assistant messages come through streaming. The receiver
        // contract says onMessageComplete fires exactly once per
        // streamStep, with role=.assistant; record the usage slot in
        // turn order regardless of whether the wire actually had usage.
        if (message.role == .assistant) {
            try self.per_message_usage.append(self.allocator, usage);
        }
        try self.stdout.writeAll("\n");
        try self.file.flush();
    }

    /// Reset any in-progress display state after a failed turn. Errors here
    /// must be swallowed — we're already in the error path.
    fn onError(ptr: *anyopaque, err: anyerror) void {
        _ = &err;
        const self: *CLIReceiver = @ptrCast(@alignCast(ptr));
        // If we were rendering a Thinking block, clear the dim style so
        // subsequent output (including the "[error: ...]" line) is readable.
        self.stdout.writeAll("\x1b[0m") catch {};
        self.file.flush() catch {};
    }
};

/// Spin up a Lua interpreter, run a no-op, tear it down. Catches
/// link/compile errors on the Lua dependency at the earliest moment.
/// Logs only in debug builds; release builds run silently.
fn luaSmokeCheck() void {
    const L = lua.luaL_newstate() orelse {
        std.log.err("lua: luaL_newstate returned null", .{});
        return;
    };
    defer lua.lua_close(L);
    lua.luaL_openlibs(L);

    // Push _VERSION to confirm libs loaded.
    _ = lua.lua_getglobal(L, "_VERSION");
    const ver = lua.lua_tolstring(L, -1, null);
    if (ver != null) {
        std.log.debug("lua: {s} linked OK", .{ver});
    }
    lua.lua_settop(L, 0);
}

fn loadConfig(environ_map: *const std.process.Environ.Map) !panto.config.Config {
    const style_str = environ_map.get("PANTO_API_STYLE") orelse "openai_chat";
    const style = std.meta.stringToEnum(panto.config.APIStyle, style_str) orelse {
        std.debug.print(
            "error: PANTO_API_STYLE must be one of: openai_chat, anthropic_messages (got: {s})\n",
            .{style_str},
        );
        return error.InvalidApiStyle;
    };

    switch (style) {
        .openai_chat => {
            const api_key = environ_map.get("OPENAI_API_KEY") orelse {
                std.debug.print("error: OPENAI_API_KEY is required\n", .{});
                return error.MissingApiKey;
            };
            const base_url = environ_map.get("OPENAI_BASE_URL") orelse "https://api.openai.com/v1";
            const model = environ_map.get("OPENAI_MODEL") orelse {
                std.debug.print("error: OPENAI_MODEL is required\n", .{});
                return error.MissingModel;
            };
            const reasoning: panto.config.ReasoningEffort =
                if (environ_map.get("OPENAI_REASONING")) |val|
                    std.meta.stringToEnum(panto.config.ReasoningEffort, val) orelse {
                        std.debug.print(
                            "error: OPENAI_REASONING must be one of: default, off, minimal, low, medium, high (got: {s})\n",
                            .{val},
                        );
                        return error.InvalidReasoning;
                    }
                else
                    .default;
            return .{ .openai_chat = .{
                .api_key = api_key,
                .base_url = base_url,
                .model = model,
                .reasoning = reasoning,
            } };
        },
        .anthropic_messages => {
            const api_key = environ_map.get("ANTHROPIC_API_KEY") orelse {
                std.debug.print("error: ANTHROPIC_API_KEY is required\n", .{});
                return error.MissingApiKey;
            };
            const base_url = environ_map.get("ANTHROPIC_BASE_URL") orelse "https://api.anthropic.com";
            const model = environ_map.get("ANTHROPIC_MODEL") orelse {
                std.debug.print("error: ANTHROPIC_MODEL is required\n", .{});
                return error.MissingModel;
            };
            const api_version = environ_map.get("ANTHROPIC_API_VERSION") orelse "2023-06-01";
            const max_tokens: u32 = if (environ_map.get("ANTHROPIC_MAX_TOKENS")) |val|
                std.fmt.parseInt(u32, val, 10) catch {
                    std.debug.print(
                        "error: ANTHROPIC_MAX_TOKENS must be a positive integer (got: {s})\n",
                        .{val},
                    );
                    return error.InvalidMaxTokens;
                }
            else
                4096;
            return .{ .anthropic_messages = .{
                .api_key = api_key,
                .base_url = base_url,
                .model = model,
                .api_version = api_version,
                .max_tokens = max_tokens,
            } };
        },
    }
}

pub fn main(init: std.process.Init) !void {
    const alloc = init.gpa;
    const io = init.io;

    // Smoke test: prove Lua is linked. Slice 1 — no extension runtime
    // wired up yet, this just confirms the static lib comes through.
    luaSmokeCheck();

    // Resolve the absolute path of the running panto binary. Needed
    // both by `panto lua` (we re-exec ourselves through a wrapper
    // luarocks invokes) and by the agent's bootstrap.
    const panto_path = try self_exe.selfExePathAlloc(alloc);
    defer alloc.free(panto_path);

    // Subcommand dispatch: `panto lua` and `panto bootstrap` short
    // out of the agent loop, but still run the same luarocks bootstrap
    // pipeline so first-run setup happens consistently.
    switch (try subcommand.dispatch(
        alloc,
        io,
        init.environ_map,
        init.minimal.args,
        panto_path,
    )) {
        .done => return,
        .agent => {},
    }

    const config = try loadConfig(init.environ_map);

    // Parse the agent-mode flags. Currently only `--resume [<id>]`.
    const cli_flags = try parseAgentFlags(alloc, init.minimal.args);
    defer cli_flags.deinit(alloc);

    var stdout_buffer: [4096]u8 = undefined;
    var stdout_file = std.Io.File.stdout().writer(io, &stdout_buffer);
    const stdout = &stdout_file.interface;

    var stdin_buffer: [4096]u8 = undefined;
    var stdin_file = std.Io.File.stdin().reader(io, &stdin_buffer);
    const stdin = &stdin_file.interface;

    // Resolve where this project's sessions live.
    var cwd_buf: [std.fs.max_path_bytes]u8 = undefined;
    const cwd_n = try std.process.currentPath(io, &cwd_buf);
    const cwd = cwd_buf[0..cwd_n];
    const session_dir = try session_paths.sessionDirForCwd(alloc, init.environ_map, cwd);
    defer alloc.free(session_dir);

    // Load the user's models.toml — a missing file is fine (empty
    // registry). Cost lookups against an empty registry return null,
    // and the display layer will format that as "unknown."
    const models_toml_path = try models_toml.configPath(alloc, init.environ_map);
    defer alloc.free(models_toml_path);
    var pricing_registry = try models_toml.loadFromPath(alloc, io, models_toml_path);
    defer pricing_registry.deinit();
    std.log.debug("models.toml: {d} entries from {s}", .{ pricing_registry.count(), models_toml_path });

    const banner_model_initial: []const u8 = switch (config) {
        inline else => |c| c.model,
    };
    const banner_provider_initial: []const u8 = @tagName(config);

    // Create or resume the session. Resume failures (missing/ambiguous id)
    // are user errors — print a tidy message and exit 1 rather than
    // printing a Zig stack trace.
    var session_mgr = openSession(
        alloc,
        io,
        session_dir,
        cwd,
        cli_flags,
        stdout,
        &stdout_file,
    ) catch |err| switch (err) {
        error.SessionNotFound, error.AmbiguousSessionId => std.process.exit(1),
        else => return err,
    };
    defer session_mgr.deinit();

    var conv = panto.conversation.Conversation.init(alloc);
    defer conv.deinit();

    if (session_mgr.getEntries().len > 0) {
        // Resumed an existing session — rebuild the conversation from the
        // log. The system prompt is part of the log.
        conv.deinit();
        conv = try session_mgr.rebuildConversation();
        try stdout.print(
            "resumed session {s} ({d} entries)\n",
            .{ session_mgr.getSessionId()[0..@min(8, session_mgr.getSessionId().len)], session_mgr.getEntries().len },
        );
    } else {
        // Fresh session — install the default system prompt and record it.
        const system_text = "You are a helpful assistant.";
        try conv.addSystemMessage(system_text);
        try appendSystemToSession(alloc, &session_mgr, system_text);
    }

    const prov = try panto.provider.Provider.init(alloc, io, config);
    var agent = panto.agent.Agent.init(alloc, io, prov);
    defer agent.deinit();

    // Spin up the long-lived Lua runtime. All Lua extensions load into
    // one `lua_State`; module-global state survives across calls. The
    // runtime registers with the agent as a single `ToolSource` named
    // `panto-lua`.
    var rt = try lua_runtime.LuaRuntime.create(alloc);
    defer rt.deinit();

    // Bootstrap luarocks against the Lua runtime's lua_State — same
    // pipeline as `panto lua` and `panto bootstrap`. After this,
    // `require("luarocks.*")` works and any pinned batteries from the
    // manifest are installed under $PANTO_HOME.
    const luarocks_rt = try luarocks_runtime.bootstrap(
        alloc,
        io,
        init.environ_map,
        rt.L,
        panto_path,
    );
    defer luarocks_rt.deinit();

    // luv is installed (or already present) at this point; wire the
    // libuv-driven coroutine scheduler before any extensions get a
    // chance to register tools that might want to yield.
    try rt.installScheduler();

    // Discover Lua extensions across three layers — system
    // ($PANTO_HOME/agent), user ($XDG_CONFIG_HOME/panto or
    // $HOME/.config/panto), and project (./.panto). Project shadows
    // user shadows system; tool-name collisions across surviving
    // entries abort startup.
    const n_ext_tools = extension_loader.discoverAndLoad(
        alloc,
        io,
        init.environ_map,
        luarocks_rt.layout.agent_dir,
        rt,
    ) catch |err| {
        std.log.err("extension discovery failed: {t}", .{err});
        return err;
    };
    std.log.debug("extensions: {d} tool(s) registered", .{n_ext_tools});

    if (n_ext_tools > 0) {
        try agent.registerToolSource(rt.toolSource());
    }

    const banner_base: []const u8 = switch (config) {
        inline else => |c| c.base_url,
    };
    try stdout.print(
        "panto — {s}: {s} @ {s}\n",
        .{ banner_provider_initial, banner_model_initial, banner_base },
    );
    try stdout.print("> ", .{});
    try stdout_file.flush();

    var cli_recv = CLIReceiver{
        .stdout = stdout,
        .file = &stdout_file,
        .allocator = alloc,
    };
    defer cli_recv.deinit();
    var recv = cli_recv.receiver();

    while (true) {
        const maybe_line = stdin.takeDelimiter('\n') catch |err| {
            std.debug.print("read error: {}\n", .{err});
            return;
        };
        const line = maybe_line orelse {
            try stdout.writeAll("\n");
            try stdout_file.flush();
            break;
        };

        if (line.len == 0) {
            try stdout.writeAll("\n> ");
            try stdout_file.flush();
            continue;
        }

        try conv.addUserMessage(line);
        try appendUserPromptToSession(
            alloc,
            &session_mgr,
            line,
            banner_provider_initial,
            banner_model_initial,
        );

        const entries_before_step = conv.messages.items.len;

        cli_recv.beginTurn();
        agent.runStep(&conv, &recv) catch |err| {
            try stdout.print("\n[error: {s}]\n", .{@errorName(err)});
        };

        // Persist whatever new entries the agent produced this turn. The
        // agent loop may have appended:
        //   - assistant message(s) (one per provider response)
        //   - user messages containing ToolResult blocks (one per tool round)
        // Each assistant message gets paired with the Usage that the
        // receiver captured at its onMessageComplete time (or null if
        // the provider didn't emit usage that round).
        try persistTurn(
            alloc,
            &session_mgr,
            &conv,
            entries_before_step,
            banner_provider_initial,
            banner_model_initial,
            cli_recv.per_message_usage.items,
        );

        try stdout.writeAll("\n> ");
        try stdout_file.flush();
    }
}

// -----------------------------------------------------------------------------
// CLI flag parsing
// -----------------------------------------------------------------------------

const AgentFlags = struct {
    /// `--resume` without an id: resume the most recent session.
    /// `--resume <id>`: resume the session whose id has this prefix.
    /// Not present: start a new session.
    resume_kind: ResumeKind = .none,
    resume_id: ?[]const u8 = null, // owned

    pub fn deinit(self: AgentFlags, alloc: std.mem.Allocator) void {
        if (self.resume_id) |id| alloc.free(id);
    }
};

const ResumeKind = enum { none, most_recent, by_id };

fn parseAgentFlags(alloc: std.mem.Allocator, args: std.process.Args) !AgentFlags {
    var flags: AgentFlags = .{};
    errdefer flags.deinit(alloc);

    var it = args.iterate();
    defer it.deinit();
    _ = it.next(); // argv[0]

    while (it.next()) |a| {
        if (std.mem.eql(u8, a, "--resume")) {
            // Peek at the next arg. If it exists and doesn't start with `-`,
            // treat it as a session id (or prefix).
            const next = it.next();
            if (next) |id| {
                if (id.len > 0 and id[0] != '-') {
                    flags.resume_kind = .by_id;
                    flags.resume_id = try alloc.dupe(u8, id);
                    continue;
                } else {
                    // Not an id; rewind by treating it as a separate flag.
                    // The Args API doesn't support rewind, so handle inline.
                    flags.resume_kind = .most_recent;
                    if (std.mem.eql(u8, id, "--resume")) {
                        // back-to-back --resume; second resets, fine.
                        continue;
                    }
                    // Otherwise, fall through and process this token as a flag.
                    // (Currently we don't have other flags; ignore unknowns.)
                    std.log.warn("panto: ignoring unknown argument '{s}'", .{id});
                    continue;
                }
            }
            flags.resume_kind = .most_recent;
            continue;
        }
        // Future agent-mode flags would land here. Unknown args are tolerated
        // (the user might be passing something we don't recognize yet).
    }
    return flags;
}

// -----------------------------------------------------------------------------
// Session bootstrap
// -----------------------------------------------------------------------------

fn openSession(
    alloc: std.mem.Allocator,
    io: std.Io,
    session_dir: []const u8,
    cwd: []const u8,
    flags: AgentFlags,
    stdout: *std.Io.Writer,
    stdout_file: *std.Io.File.Writer,
) !panto.session_manager.SessionManager {
    switch (flags.resume_kind) {
        .none => return try panto.session_manager.SessionManager.init(
            alloc,
            io,
            session_dir,
            cwd,
        ),
        .most_recent => {
            const path_opt = panto.session_manager.findMostRecentSession(alloc, io, session_dir) catch null;
            if (path_opt) |path| {
                defer alloc.free(path);
                return try panto.session_manager.SessionManager.open(alloc, io, path);
            }
            try stdout.print("no sessions to resume; starting fresh.\n", .{});
            try stdout_file.flush();
            return try panto.session_manager.SessionManager.init(
                alloc,
                io,
                session_dir,
                cwd,
            );
        },
        .by_id => {
            const id = flags.resume_id.?;
            const path = panto.session_manager.resolveSessionId(alloc, io, session_dir, id) catch |err| switch (err) {
                error.SessionNotFound => {
                    try stdout.print("error: no session matching '{s}' in {s}\n", .{ id, session_dir });
                    try stdout_file.flush();
                    return err;
                },
                error.AmbiguousSessionId => {
                    try stdout.print("error: session id '{s}' is ambiguous\n", .{id});
                    try stdout_file.flush();
                    return err;
                },
                else => return err,
            };
            defer alloc.free(path);
            return try panto.session_manager.SessionManager.open(alloc, io, path);
        },
    }
}

// -----------------------------------------------------------------------------
// Session append helpers — bridge in-memory ContentBlocks to on-disk entries.
// -----------------------------------------------------------------------------

fn appendSystemToSession(
    alloc: std.mem.Allocator,
    mgr: *panto.session_manager.SessionManager,
    text: []const u8,
) !void {
    const blocks = try alloc.alloc(panto.session_manager.DiskContentBlock, 1);
    blocks[0] = .{ .text = .{ .text = try alloc.dupe(u8, text) } };
    _ = try mgr.appendMessage(
        .{ .role = .system, .content = blocks },
        null,
        null,
    );
}

fn appendUserPromptToSession(
    alloc: std.mem.Allocator,
    mgr: *panto.session_manager.SessionManager,
    text: []const u8,
    provider: []const u8,
    model: []const u8,
) !void {
    const blocks = try alloc.alloc(panto.session_manager.DiskContentBlock, 1);
    blocks[0] = .{ .text = .{ .text = try alloc.dupe(u8, text) } };
    _ = try mgr.appendMessage(
        .{ .role = .user, .content = blocks },
        provider,
        model,
    );
}

/// After the agent loop has driven a turn to completion, persist every
/// new in-memory message at index `>= start_index` to the session log.
///
/// Each in-memory message is mapped to disk:
///   - assistant → assistant entry with provider/model/stop_reason metadata
///     and Usage (from `per_message_usage`, one slot per assistant message in
///     completion order).
///   - user (with ToolResult blocks) → user entry stamped with provider/model.
///
/// `stop_reason` is recorded as `"stop"` for now; surfacing the real wire
/// value requires plumbing it through the Receiver vtable (future work,
/// separate from token plumbing).
fn persistTurn(
    alloc: std.mem.Allocator,
    mgr: *panto.session_manager.SessionManager,
    conv: *const panto.conversation.Conversation,
    start_index: usize,
    provider: []const u8,
    model: []const u8,
    per_message_usage: []const ?panto.session_manager.Usage,
) !void {
    var i = start_index;
    var assistant_seen: usize = 0;
    while (i < conv.messages.items.len) : (i += 1) {
        const msg = conv.messages.items[i];
        const blocks = try alloc.alloc(panto.session_manager.DiskContentBlock, msg.content.items.len);
        var allocated: usize = 0;
        errdefer {
            for (blocks[0..allocated]) |b| b.deinit(alloc);
            alloc.free(blocks);
        }
        for (msg.content.items) |block| {
            blocks[allocated] = try panto.session.contentBlockToDisk(alloc, block);
            allocated += 1;
        }
        switch (msg.role) {
            .system => {
                // Mid-turn system messages aren't a thing in pantograph today.
                // Treat as harmless and persist verbatim, sans stamps.
                _ = try mgr.appendMessage(
                    .{ .role = .system, .content = blocks },
                    null,
                    null,
                );
            },
            .user => {
                _ = try mgr.appendMessage(
                    .{ .role = .user, .content = blocks },
                    provider,
                    model,
                );
            },
            .assistant => {
                // Pair this assistant message with its usage, if the
                // receiver captured one for this position. (A turn
                // ending in a stream error can have fewer usage entries
                // than assistant messages; treat as null and move on.)
                const usage: ?panto.session_manager.Usage = if (assistant_seen < per_message_usage.len)
                    per_message_usage[assistant_seen]
                else
                    null;
                assistant_seen += 1;
                _ = try mgr.appendMessage(
                    .{
                        .role = .assistant,
                        .content = blocks,
                        .provider = try alloc.dupe(u8, provider),
                        .model = try alloc.dupe(u8, model),
                        .stop_reason = try alloc.dupe(u8, "stop"),
                        .usage = usage,
                    },
                    null,
                    null,
                );
            },
        }
    }
}