aish/main.lua

-- main.lua — entry point
-- Phase 0: arg parsing, config load, REPL start.
-- See docs/PHASE0.md §4, §10. -p one-shot mode lands per issue #4.

-- Resolve modules + vendored dkjson relative to this script's directory,
-- not cwd. Packaged install puts main.lua at /usr/share/lua/5.1/aish/ and
-- the /usr/bin/aish wrapper execs `luajit /usr/share/lua/5.1/aish/main.lua`
-- from whatever cwd the user is in — siblings must still resolve. Dev mode
-- (`luajit main.lua` from repo root) keeps working because arg[0] is then
-- "main.lua" with no "/" — _dir falls back to "./".
local _dir = arg[0]:match("(.*/)") or "./"
package.path = _dir .. "?.lua;" .. _dir .. "vendor/?.lua;" .. package.path

local USAGE = [[
aish — AI-augmented conversational shell.

Usage:
  aish [--config <path>] [--help]            -- interactive REPL
  aish -p "<prompt>" [--config <path>]       -- one-shot, print + exit

In -p mode, if stdin is not a TTY it's read as additional context and
prepended to the prompt as a fenced block — composes with Unix pipes:
  tail app.log | aish -p "any anomalies?"

Config resolution order (PHASE0.md §10):
  1. --config <path>
  2. $AISH_CONFIG
  3. ~/.config/aish/config.lua
  4. ./config.lua
]]

local function parse_args(argv)
    local out = {}
    local i = 1
    while i <= #argv do
        local a = argv[i]
        if a == "--config" then
            out.config = argv[i + 1]
            i = i + 2
        elseif a == "--help" or a == "-h" then
            out.help = true
            i = i + 1
        elseif a == "-p" or a == "--prompt" then
            out.prompt = argv[i + 1]
            if not out.prompt then
                io.stderr:write("aish: -p requires a prompt argument\n")
                os.exit(2)
            end
            i = i + 2
        else
            io.stderr:write("aish: unrecognized argument: " .. a .. "\n")
            os.exit(2)
        end
    end
    return out
end

local function load_config(opts)
    -- --config is explicit: use exactly that path or fail. No silent fallback.
    if opts.config then
        local f = io.open(opts.config, "r")
        if not f then
            error("aish: --config " .. opts.config .. ": cannot open")
        end
        f:close()
        return dofile(opts.config), opts.config
    end

    local home = os.getenv("HOME") or ""
    local candidates = {}
    local function push(p) if p and p ~= "" then candidates[#candidates + 1] = p end end
    push(os.getenv("AISH_CONFIG"))
    push(home .. "/.config/aish/config.lua")
    push("./config.lua")

    for _, path in ipairs(candidates) do
        local f = io.open(path, "r")
        if f then f:close(); return dofile(path), path end
    end
    error("aish: no config.lua found (tried: "
          .. table.concat(candidates, ", ") .. ")")
end

-- ---------------------------------------------------------------- Phase 9 project overlay

-- Walk-up from libc.getcwd() looking for .aish.lua. Stops at $HOME
-- OR filesystem root (whichever comes first). Returns the first
-- found path or nil. Per R1 (review fold-in), uses a proper-prefix
-- check (NOT bare bytes-prefix) to avoid false positive when HOME
-- is "/home/user" and cwd is "/home/user2/...".
local function _find_project_config()
    local home = os.getenv("HOME")
    if not home or home == "" then return nil end
    -- Lazy-require so the existing load_config path stays untouched
    -- when no project overlay considered.
    local libc_ok, libc = pcall(require, "ffi.libc")
    if not libc_ok then return nil end
    local dir = libc.getcwd()
    if not dir then return nil end
    -- R1: proper prefix (dir == home OR dir starts with home .. "/")
    if dir ~= home and dir:sub(1, #home + 1) ~= home .. "/" then
        return nil
    end
    while dir and #dir > 0 do
        local candidate = dir .. "/.aish.lua"
        local f = io.open(candidate, "rb")
        if f then f:close(); return candidate end
        if dir == home or dir == "/" then return nil end
        -- Walk up one level
        dir = dir:gsub("/[^/]*$", "")
        if dir == "" then dir = "/" end
    end
    return nil
end

local function _trust_file_path()
    return os.getenv("AISH_TRUST_FILE")
        or ((os.getenv("HOME") or "") .. "/.aish/trusted-projects")
end

-- Interactive trust prompt. R2: caller must NOT invoke this in
-- one-shot (-p) mode — io.read or rl.readline would consume piped
-- stdin. Returns true on user accept (and persists the trust).
-- Per A8, uses rl.readline; if it misbehaves at this early call
-- site, the function returns false (skip overlay) rather than
-- falling back to io.read.
local function _check_and_maybe_prompt(project_path, history)
    local sha = history._sha256_file(project_path)
    if not sha then
        io.stderr:write("aish: project config " .. project_path
                        .. ": sha256 failed; skipping\n")
        return false
    end
    local tpath = _trust_file_path()
    if history.is_trusted(tpath, project_path, sha) then
        return true
    end
    -- Trust prompt.
    io.stderr:write("aish: project config found: " .. project_path .. "\n")
    io.stderr:write("aish: UNTRUSTED. Loading it runs arbitrary Lua code.\n")
    local rl_ok, rl = pcall(require, "ffi.readline")
    if not rl_ok then
        io.stderr:write("aish: readline unavailable; declining trust prompt\n")
        return false
    end
    local ans = rl.readline("[aish] trust this project config? [y/N] ")
    if ans and ans:lower():sub(1, 1) == "y" then
        history.add_trusted(tpath, project_path, sha)
        return true
    end
    return false
end

-- Wrap load_config with a project-overlay step. Always-on (no
-- config flag); overlay activates only when a trusted .aish.lua
-- is found in/above cwd. In one-shot (-p) mode the trust prompt
-- is SKIPPED to avoid io consuming piped stdin (R2) — only pre-
-- trusted overlays load in -p.
local function load_config_with_overlay(opts)
    local user_cfg, user_path = load_config(opts)
    local sources = {}
    for k, _ in pairs(user_cfg) do sources[k] = "user" end

    local proj_path = _find_project_config()
    if not proj_path then
        user_cfg._sources = sources
        return user_cfg, user_path, nil
    end

    local history_ok, history = pcall(require, "history")
    if not history_ok then
        user_cfg._sources = sources
        return user_cfg, user_path, nil
    end

    -- R2: skip trust prompt in -p mode.
    local trusted
    if opts.prompt then
        local sha = history._sha256_file(proj_path)
        local tpath = _trust_file_path()
        trusted = sha and history.is_trusted(tpath, proj_path, sha)
        if not trusted then
            io.stderr:write("aish: project config " .. proj_path
                .. " skipped in -p mode (untrusted; run aish interactively to trust)\n")
        end
    else
        trusted = _check_and_maybe_prompt(proj_path, history)
    end
    if not trusted then
        user_cfg._sources = sources
        return user_cfg, user_path, nil
    end

    local ok, proj_cfg = pcall(dofile, proj_path)
    if not ok or type(proj_cfg) ~= "table" then
        io.stderr:write("aish: project config " .. proj_path
                        .. " load failed: " .. tostring(proj_cfg) .. "\n")
        user_cfg._sources = sources
        return user_cfg, user_path, nil
    end
    -- Shallow merge: project replaces user at top level. Update sources map.
    for k, v in pairs(proj_cfg) do
        user_cfg[k] = v
        sources[k] = "project"
    end
    user_cfg._sources = sources
    return user_cfg, user_path, proj_path
end

-- One-shot mode: read non-TTY stdin (if any), compose prompt, stream
-- broker reply to stdout, exit. Bypasses repl.lua entirely — no REPL,
-- no MCP, no tool loop, no Norris. The model's reply is printed
-- verbatim (including any "CMD:" lines, which are NOT executed in
-- this mode by design — the user can pipe-grep them as they wish).
local function run_one_shot(config, user_prompt)
    local libc   = require("ffi.libc")
    local broker = require("broker")

    local composed = user_prompt
    if not libc.isatty(0) then
        local piped = io.read("*a") or ""
        if piped ~= "" then
            composed = "```\n" .. piped .. "\n```\n\n" .. user_prompt
        end
    end

    local model_name = config.default_model
    local model_cfg  = config.models and config.models[model_name]
    if not model_cfg then
        io.stderr:write(("aish: default_model '%s' not found in models{}\n")
                        :format(tostring(model_name)))
        os.exit(2)
    end

    local messages = { { role = "user", content = composed } }
    local got_any = false
    local ok, err = broker.chat_stream(model_cfg, messages,
        function(kind, payload)
            if kind == "text" and payload and payload ~= "" then
                io.write(payload); io.flush()
                got_any = true
            end
        end)
    if not ok then
        if got_any then io.write("\n") end
        io.stderr:write("aish: broker error: " .. tostring(err) .. "\n")
        os.exit(1)
    end
    if got_any then io.write("\n") end
end

local function main(argv)
    local opts = parse_args(argv or {})
    if opts.help then io.write(USAGE); return end

    local config, config_path, project_path = load_config_with_overlay(opts)
    io.stderr:write(("aish: loaded config from %s\n"):format(config_path))
    if project_path then
        io.stderr:write(("aish: project config: %s (overlaid on %s)\n")
                        :format(project_path, config_path))
    end

    if opts.prompt then
        run_one_shot(config, opts.prompt)
        return
    end

    local repl = require("repl")
    repl.run(config)
end

main(arg)