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");
const config_file = @import("config_file.zig");
const glob = @import("glob.zig");
const system_prompt = @import("system_prompt.zig");
const command = @import("command.zig");
const command_compaction = @import("compaction.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;
    _ = config_file;
    _ = glob;
    _ = system_prompt;
    _ = command;
    _ = command_compaction;
}

const ContentBlockType = panto.ContentBlockType;
const MessageRole = panto.MessageRole;
const Event = panto.Event;

/// Display state for the pull-stream CLI renderer. Prints streaming deltas
/// to stdout; thinking blocks are dimmed with ANSI escape codes, text blocks
/// render plain. Persistence is owned by the agent; the renderer is
/// display-only.
const CLIRenderer = struct {
    stdout: *std.Io.Writer,
    file: *std.Io.File.Writer,
    allocator: std.mem.Allocator,

    /// Per-turn reset hook. Currently a no-op (no per-turn render state),
    /// retained as the REPL's turn-boundary signal.
    pub fn beginTurn(self: *CLIRenderer) void {
        _ = self;
    }

    pub fn deinit(self: *CLIRenderer) void {
        _ = self;
    }

    /// Render one pull `Event`. Mirrors the prior push receiver's output
    /// exactly: dim `[thinking]`, a `tool:` prefix opened at block start
    /// with the name appended at block_complete, a trailing newline at each
    /// `message_complete`, and a dim retry status line.
    fn render(self: *CLIRenderer, ev: Event) !void {
        switch (ev) {
            .message_start => {},
            .block_start => |b| {
                switch (b.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 name lands at block_complete from the
                    // assembled ContentBlock.
                    .ToolUse => try self.stdout.writeAll("\n\x1b[36mtool: \x1b[0m"),
                    else => {},
                }
                try self.file.flush();
            },
            // The print-based CLI defers tool-name rendering to
            // block_complete to avoid cursor gymnastics (we can't go back and
            // edit the prefix), so tool_details is a no-op here.
            .tool_details => {},
            .content_delta => |d| {
                try self.stdout.writeAll(d.delta);
                try self.file.flush();
            },
            .block_complete => |b| {
                switch (b.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();
            },
            .message_complete => {
                try self.stdout.writeAll("\n");
                try self.file.flush();
            },
            // Surface provider retry scheduling as a dim status line so the
            // user can see the agent is waiting on the provider, not hung.
            .provider_retry => |info| {
                if (info.compaction) {
                    self.stdout.print(
                        "\x1b[2mcontext overflow: compacting and retrying\x1b[0m\n",
                        .{},
                    ) catch {};
                } else {
                    const secs = @as(f64, @floatFromInt(info.delay_ms)) / 1000.0;
                    self.stdout.print(
                        "\x1b[2mprovider unavailable ({s}): retrying in {d:.1}s (attempt {d}/{d})\x1b[0m\n",
                        .{ @errorName(info.err), secs, info.attempt + 1, info.max_attempts },
                    ) catch {};
                }
                self.file.flush() catch {};
            },
            .tool_dispatch_start, .tool_dispatch_complete, .turn_complete => {},
        }
    }

    /// Reset any in-progress display state after a failed turn. Errors here
    /// must be swallowed — we're already in the error path.
    fn renderError(self: *CLIRenderer, err: anyerror) void {
        _ = &err;
        // 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 {};
    }
};

/// Drive a whole turn: open the pull `Stream` and render every event until
/// it terminates (`turn_complete` → `next()` returns null). A failure
/// surfaces as the error from `next()`; the caller renders it. The stream is
/// always `deinit`ed (persisting the turn tail) on every exit path.
fn driveTurn(agent: *panto.agent.Agent, message: panto.agent.Agent.UserMessage, renderer: *CLIRenderer) !void {
    var stream = try agent.run(message);
    defer stream.deinit();
    while (try stream.next()) |ev| try renderer.render(ev);
}

/// 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);
}

/// Print a friendly message for a config.toml load failure before the
/// process exits non-zero.
fn reportConfigError(err: anyerror) void {
    switch (err) {
        error.InvalidConfigToml => std.debug.print(
            "error: a config.toml failed to parse (see log above).\n",
            .{},
        ),
        error.InvalidProvider => std.debug.print(
            "error: a [providers.*] entry is malformed (need style + base_url).\n",
            .{},
        ),
        error.InvalidModelRef => std.debug.print(
            "error: defaults.model must look like \"provider:model\".\n",
            .{},
        ),
        error.UnknownDefaultProvider => std.debug.print(
            "error: defaults.model names a provider that isn't configured (or whose API key env var isn't set).\n",
            .{},
        ),
        error.PolicyConflict => std.debug.print(
            "error: a tool/extension pattern appears in both allow and deny.\n",
            .{},
        ),
        error.NoHomeDirectory => std.debug.print(
            "error: cannot resolve config paths — set HOME or XDG_CONFIG_HOME/XDG_DATA_HOME.\n",
            .{},
        ),
        else => std.debug.print("error: failed to load config: {s}\n", .{@errorName(err)}),
    }
}

/// Print a friendly message for a model-selection failure.
fn reportModelError(err: anyerror, cfg: *const config_file.Config) void {
    switch (err) {
        error.NoModelSelected => {
            std.debug.print(
                "error: no model selected. Set `defaults.model = \"provider:alias\"` in config.toml.\n",
                .{},
            );
            if (cfg.providers.len == 0) {
                std.debug.print(
                    "  (no providers resolved — check that an API key or its env var is present.)\n",
                    .{},
                );
            }
        },
        error.UnknownProvider => std.debug.print(
            "error: the selected model names an unconfigured provider.\n",
            .{},
        ),
        else => std.debug.print("error: model selection failed: {s}\n", .{@errorName(err)}),
    }
}

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 => {},
    }

    // 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 merged config.toml (base → user → project). Providers,
    // default model, and tool/extension policy all come from here.
    var app_config = config_file.load(alloc, io, init.environ_map, cwd) catch |err| {
        reportConfigError(err);
        std.process.exit(1);
    };
    defer app_config.deinit();

    // Load models.toml — model aliases (wire name + knobs) and pricing.
    // A missing file is fine (empty registries); cost lookups return
    // null, formatted as "unknown" by the display layer.
    const models_toml_path = try models_toml.configPath(alloc, init.environ_map);
    defer alloc.free(models_toml_path);
    var models = try models_toml.loadFromPath(alloc, io, models_toml_path);
    defer models.deinit();
    std.log.debug(
        "models.toml: {d} model def(s), {d} price entr(ies) from {s}",
        .{ models.defs.count(), models.pricing.count(), models_toml_path },
    );
    // `models.pricing` is parsed and ready; the print-based CLI doesn't
    // surface per-turn cost yet (that lands with the TUI frontend).

    // Choose the active model and assemble the provider config.
    const model_ref = app_config.selectModel(&models.defs, null) catch |err| {
        reportModelError(err, &app_config);
        std.process.exit(1);
    };
    const provider_config = config_file.buildProviderConfig(&app_config, &models.defs, model_ref) catch |err| {
        reportModelError(err, &app_config);
        std.process.exit(1);
    };

    // Resolve the optional compaction-model override into a ProviderConfig.
    // On any resolution failure we log and fall back to no override (the
    // active chat model is used for compaction).
    const compaction_provider_config: ?panto.ProviderConfig = blk: {
        const ref = app_config.compaction_model_ref orelse break :blk null;
        break :blk config_file.buildProviderConfig(&app_config, &models.defs, ref) catch |err| {
            std.log.warn("compaction_model resolution failed ({t}); using active model", .{err});
            break :blk null;
        };
    };
    const compaction_cfg: panto.CompactionConfig = .{
        .keep_verbatim = app_config.compaction_keep_verbatim orelse 20_000,
        .model = compaction_provider_config,
    };

    // Process-global HTTP client: one connection pool for every provider /
    // base_url the agent may switch to. Torn down at shutdown.
    panto.init(alloc, io);
    defer panto.deinit();

    const banner_model_initial: []const u8 = switch (provider_config) {
        inline else => |c| c.model,
    };
    const banner_provider_initial: []const u8 = model_ref.provider;

    // The session store: a directory-backed JSONL catalog rooted at the
    // per-cwd `session_dir` (the CLI owns the cwd→dir grouping; the store
    // is cwd-agnostic). Must outlive the agent, which appends through a
    // `Session` handle minted/resolved below.
    var session_store_impl = try panto.FileSystemJSONLStore.init(alloc, io, session_dir, cwd);
    defer session_store_impl.deinit();
    const session_store = session_store_impl.store();

    // 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 = openSession(
        session_store,
        cli_flags,
        session_dir,
        stdout,
        &stdout_file,
    ) catch |err| switch (err) {
        error.SessionNotFound, error.AmbiguousSessionId => std.process.exit(1),
        else => return err,
    };
    // `session.info` is adopted by the agent below (`Agent.init` can't fail)
    // and freed in the agent's `deinit`; no separate cleanup here.

    const is_resume = cli_flags.resume_kind != .none and session.info.message_count > 0;

    // On resume, reconstruct the conversation from the store. (The
    // dangling-prompt recovery feature was dropped in R2.) On a fresh
    // session, the agent starts with an empty conversation.
    var adopted_conversation: ?panto.conversation.Conversation = null;
    if (is_resume) {
        adopted_conversation = try session.load();
        const sid = session.info.id;
        try stdout.print(
            "resumed session {s} ({d} messages)\n",
            .{ sid[0..@min(8, sid.len)], session.info.message_count },
        );
    }

    // Assemble the active configuration snapshot and build the agent now,
    // before luarocks bootstrap and extension loading. The agent adopts the
    // session handle (for free persistence) and the resumed conversation
    // (or starts fresh). The agent owns its own tool registry (post-R1);
    // extensions register their tool source onto it *after* the agent is
    // built but before the first turn, so in-place registration is visible.
    // System-prompt seeding/reconciliation below runs *through* the agent
    // so those entries persist.
    const active_config: panto.Config = .{
        .provider = provider_config,
        .compaction = compaction_cfg,
    };
    var agent = panto.agent.Agent.init(
        alloc,
        io,
        &active_config,
        session,
        adopted_conversation,
    );
    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();

    // Source the system prompt now that the base agent tree has been
    // staged to `$PANTO_HOME/agent` (the bootstrap above writes the
    // bundled `SYSTEM.md` there). The prompt is sourced by convention
    // from SYSTEM.md / APPEND_SYSTEM.md across the base/user/project
    // layers; the base layer is `luarocks_rt.layout.agent_dir`.
    var sp_arena = std.heap.ArenaAllocator.init(alloc);
    defer sp_arena.deinit();
    if (is_resume) {
        // SYSTEM.md / APPEND_SYSTEM.md may have changed since this log was
        // created; reconcile without rewriting history.
        try system_prompt.reconcileResume(
            sp_arena.allocator(),
            io,
            init.environ_map,
            luarocks_rt.layout.agent_dir,
            &agent,
        );
    } else {
        // Fresh session — source and install the system prompt.
        try system_prompt.seedFresh(
            sp_arena.allocator(),
            io,
            init.environ_map,
            luarocks_rt.layout.agent_dir,
            &agent,
        );
    }

    // 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 — base
    // ($PANTO_HOME/agent), user ($XDG_CONFIG_HOME/panto or
    // $HOME/.config/panto), and project (./.panto). Project shadows
    // user shadows base; 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,
        .{
            .extensions = &app_config.extensions,
            .tools = &app_config.tools,
        },
    ) 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());
    }

    // Resolve the compaction system prompt (COMPACTION.md across layers,
    // last wins; built-in default otherwise) and arm automatic compaction.
    // Owned by `sp_arena`, which lives for the whole REPL.
    const compaction_prompt = try system_prompt.resolveCompaction(
        sp_arena.allocator(),
        io,
        init.environ_map,
        luarocks_rt.layout.agent_dir,
    );
    agent.compaction_system_prompt = compaction_prompt;

    const banner_base: []const u8 = switch (provider_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_renderer = CLIRenderer{
        .stdout = stdout,
        .file = &stdout_file,
        .allocator = alloc,
    };
    defer cli_renderer.deinit();

    // Build the slash-command registry and register builtins, then append
    // any commands declared by Lua extensions (harvested at load time).
    var cmd_registry = command.Registry.init(alloc);
    defer cmd_registry.deinit();
    try command_compaction.register(&cmd_registry);

    // Append slash commands declared by Lua extensions via
    // `panto.register_command`. A name collision with a builtin (or
    // between two extensions) surfaces as `error.DuplicateCommand` and
    // aborts startup, matching the tool-name collision policy.
    for (rt.commandList()) |lua_cmd| {
        cmd_registry.registerLua(
            lua_cmd.name,
            lua_cmd.description,
            lua_cmd.handler_ref,
        ) catch |err| {
            std.log.err(
                "lua: failed to register command '/{s}': {t}",
                .{ lua_cmd.name, err },
            );
            return err;
        };
    }

    var cmd_ctx: command.Context = .{
        .allocator = alloc,
        .agent = &agent,
        .stdout = stdout,
        .stdout_file = &stdout_file,
        .compaction_prompt = compaction_prompt,
        .provider_name = banner_provider_initial,
        .model_name = banner_model_initial,
        .lua_rt = rt,
    };

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

        // Slash commands are handled by the CLI, not sent to the model.
        // Any line starting with `/` is a command; an unrecognized one
        // prints an error rather than reaching the model.
        if (std.mem.startsWith(u8, line, "/")) {
            cmd_registry.dispatch(line, &cmd_ctx) catch |err| switch (err) {
                command.Error.CommandNotFound => {
                    try stdout.print("\n[unknown command: {s}]\n> ", .{line});
                    try stdout_file.flush();
                },
                else => {
                    try stdout.print("\n[command error: {s}]\n> ", .{@errorName(err)});
                    try stdout_file.flush();
                },
            };
            continue;
        }

        // Drive the turn. `run` submits + persists the user prompt and the
        // agent owns the conversation, persisting everything it generates
        // (assistant/tool-result messages, and the compaction window on
        // auto-compaction) through its session store — the CLI no longer
        // touches persistence.
        cli_renderer.beginTurn();
        driveTurn(&agent, .{ .text = line }, &cli_renderer) catch |err| {
            cli_renderer.renderError(err);
            try stdout.print("\n[error: {s}]\n", .{@errorName(err)});
        };

        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
// -----------------------------------------------------------------------------

/// Mint or resolve the `Session` to drive, against the catalog `store`.
/// Fresh sessions are created on demand; resume resolves by id or picks the
/// most recent. The returned `Session` owns its `info` (freed via the
/// agent's `deinit`, which adopts it).
fn openSession(
    store: panto.SessionStore,
    flags: AgentFlags,
    session_dir: []const u8,
    stdout: *std.Io.Writer,
    stdout_file: *std.Io.File.Writer,
) !panto.Session {
    switch (flags.resume_kind) {
        .none => return store.create(),
        .most_recent => {
            if (try store.latest()) |sess| return sess;
            try stdout.print("no sessions to resume; starting fresh.\n", .{});
            try stdout_file.flush();
            return store.create();
        },
        .by_id => {
            const id = flags.resume_id.?;
            if (try store.resolve(id)) |sess| return sess;
            try stdout.print("error: no session matching '{s}' in {s}\n", .{ id, session_dir });
            try stdout_file.flush();
            return error.SessionNotFound;
        },
    }
}