const std = @import("std");
const dbmod = @import("db.zig");
const txnmod = @import("txn.zig");
const rowmod = @import("row.zig");
const idx = @import("index.zig");
const querymod = @import("query.zig");
const txntablemod = @import("txn_table.zig");
const collectionsmod = @import("collections.zig");
const collttlmod = @import("collttl.zig");
const ttlmod = @import("ttl.zig");
const exprmod = @import("expr.zig");
const joinmod = @import("join.zig");
const aggregatemod = @import("aggregate.zig");
const windowmod = @import("window.zig");

const Allocator = std.mem.Allocator;
const Db = dbmod.Db;
const Transaction = txnmod.Transaction;
const Table = querymod.Table;
const TableTxn = txntablemod.TableTxn;
const Value = rowmod.Value;
const Field = rowmod.Field;
const Row = rowmod.Row;
const CollKind = collectionsmod.Kind;

/// Errors `Repl.dispatch` can raise on top of whatever the underlying `Db`/
/// `Transaction`/`Table` call it wraps returns (those propagate as-is, e.g.
/// `error.Conflict` from a failed commit).
pub const ReplError = error{
    UnknownCommand,
    WrongArgCount,
    NoActiveTransaction,
    TransactionAlreadyActive,
    InvalidField,
    UnknownCollectionKind,
    UnknownCollectionOp,
    InvalidArgument,
    NoSuchKey,
};

/// Interactive session bound to one already-open `*Db`. The one-shot CLI in
/// main.zig opens a fresh `Db` per process per command, and its `txn
/// <script-file>` command runs a whole transaction non-interactively from a
/// static file. `Repl` instead keeps the store open across many commands and
/// lets `begin`/`commit`/`rollback` bracket an interactive `Transaction`:
/// a `get` between `begin` and `commit` can steer what the *next* `put`
/// does (e.g. "read balance, then decide how much to add"), which a script
/// file evaluated top to bottom cannot express.
///
/// `table put/get/del` follow the active transaction automatically: with no
/// transaction open they go straight through `query.Table` (each call is
/// independently durable the moment it returns, exactly like `table put` on
/// the plain CLI); with a transaction open they route through a `TableTxn`
/// bound to it instead, so multiple `table put` calls across multiple
/// tables in the same `begin`/`commit` bracket commit atomically as one
/// `WriteBatch`, or none of them do on `error.Conflict` - the interactive
/// equivalent of `txn <script-file>`.
///
/// `scan` always reads the live, already-committed `Db` state regardless of
/// any open transaction: `Transaction` has no scan/range primitive (only
/// `get`, which resolves the transaction's own staged writes first), so a
/// staged-but-uncommitted put is invisible to `scan` until `commit` lands it
/// for real. This is a deliberate, documented limitation, not an oversight.
/// `table select`/`table join` are the same story one level up: `TableTxn`
/// only ever brackets single-row put/get/delete, so both always run against
/// `query.Table.open(self.db, ...)` directly, never through the active
/// transaction, exactly like `scan`.
pub const Repl = struct {
    db: *Db,
    allocator: Allocator,
    txn: ?Transaction = null,
    tableTxns: std.StringHashMap(TableTxn),

    pub fn init(allocator: Allocator, db: *Db) Repl {
        return .{
            .db = db,
            .allocator = allocator,
            .tableTxns = std.StringHashMap(TableTxn).init(allocator),
        };
    }

    pub fn deinit(self: *Repl) void {
        self.discardTxnState();
        self.tableTxns.deinit();
    }

    pub fn inTransaction(self: *const Repl) bool {
        return self.txn != null;
    }

    fn discardTxnState(self: *Repl) void {
        var it = self.tableTxns.iterator();
        while (it.next()) |entry| {
            self.allocator.free(entry.key_ptr.*);
            entry.value_ptr.deinit();
        }
        self.tableTxns.clearRetainingCapacity();
        if (self.txn) |*t| {
            t.deinit();
            self.txn = null;
        }
    }

    /// Runs one line of input, writing command output or `error: <name>` to
    /// `writer`. Returns `false` on `exit`/`quit` (caller should stop
    /// reading), `true` otherwise - including after a reported error, so one
    /// bad line never kills the session.
    pub fn runLine(self: *Repl, line: []const u8, writer: anytype) !bool {
        const trimmed = std.mem.trim(u8, line, " \t\r\n");
        if (trimmed.len == 0 or trimmed[0] == '#') return true;

        var tokens = std.ArrayList([]const u8).init(self.allocator);
        defer tokens.deinit();
        var it = std.mem.tokenizeAny(u8, trimmed, " \t");
        while (it.next()) |tok| try tokens.append(tok);
        if (tokens.items.len == 0) return true;

        const cmd = tokens.items[0];
        if (std.mem.eql(u8, cmd, "exit") or std.mem.eql(u8, cmd, "quit")) return false;

        self.dispatch(cmd, tokens.items[1..], writer) catch |err| {
            try writer.print("error: {s}\n", .{@errorName(err)});
        };
        return true;
    }

    /// Reads lines from `reader` until EOF, running each through `runLine`.
    /// Writes a `cobalt>`/`cobalt(txn)>` prompt before each read when
    /// `prompt` is set (interactive stdin use); script files pass `false`.
    pub fn runFromReader(self: *Repl, reader: anytype, writer: anytype, prompt: bool) !void {
        var buf: [8192]u8 = undefined;
        while (true) {
            if (prompt) {
                try writer.writeAll(if (self.inTransaction()) "cobalt(txn)> " else "cobalt> ");
            }
            const line = (try reader.readUntilDelimiterOrEof(&buf, '\n')) orelse break;
            if (!try self.runLine(line, writer)) break;
        }
    }

    fn dispatch(self: *Repl, cmd: []const u8, args: []const []const u8, writer: anytype) !void {
        if (std.mem.eql(u8, cmd, "begin")) return self.cmdBegin(args, writer);
        if (std.mem.eql(u8, cmd, "commit")) return self.cmdCommit(args, writer);
        if (std.mem.eql(u8, cmd, "rollback")) return self.cmdRollback(args, writer);
        if (std.mem.eql(u8, cmd, "put")) return self.cmdPut(args, writer);
        if (std.mem.eql(u8, cmd, "get")) return self.cmdGet(args, writer);
        if (std.mem.eql(u8, cmd, "del")) return self.cmdDel(args, writer);
        if (std.mem.eql(u8, cmd, "scan")) return self.cmdScan(args, writer);
        if (std.mem.eql(u8, cmd, "stats")) return self.cmdStats(args, writer);
        if (std.mem.eql(u8, cmd, "flush")) return self.cmdFlush(args, writer);
        if (std.mem.eql(u8, cmd, "compact")) return self.cmdCompact(args, writer);
        if (std.mem.eql(u8, cmd, "table")) return self.cmdTable(args, writer);
        if (std.mem.eql(u8, cmd, "coll")) return self.cmdColl(args, writer);
        if (std.mem.eql(u8, cmd, "expire")) return self.cmdExpire(args, writer);
        if (std.mem.eql(u8, cmd, "persist")) return self.cmdPersist(args, writer);
        if (std.mem.eql(u8, cmd, "ttl")) return self.cmdTtl(args, writer);
        if (std.mem.eql(u8, cmd, "expire-sweep")) return self.cmdExpireSweep(args, writer);
        if (std.mem.eql(u8, cmd, "help")) return self.cmdHelp(writer);
        return ReplError.UnknownCommand;
    }

    fn cmdBegin(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 0) return ReplError.WrongArgCount;
        if (self.txn != null) return ReplError.TransactionAlreadyActive;
        self.txn = Transaction.begin(self.db);
        try writer.writeAll("OK\n");
    }

    fn cmdCommit(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 0) return ReplError.WrongArgCount;
        if (self.txn == null) return ReplError.NoActiveTransaction;
        try self.txn.?.commit();
        self.discardTxnState();
        try writer.writeAll("OK\n");
    }

    fn cmdRollback(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 0) return ReplError.WrongArgCount;
        if (self.txn == null) return ReplError.NoActiveTransaction;
        self.discardTxnState();
        try writer.writeAll("OK\n");
    }

    fn cmdPut(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 2) return ReplError.WrongArgCount;
        if (self.txn) |*t| {
            try t.put(args[0], args[1]);
        } else {
            try self.db.put(args[0], args[1]);
        }
        try writer.writeAll("OK\n");
    }

    fn cmdGet(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 1) return ReplError.WrongArgCount;
        const value = if (self.txn) |*t|
            try t.get(self.allocator, args[0])
        else
            try self.db.get(self.allocator, args[0]);
        defer if (value) |v| self.allocator.free(v);
        if (value) |v| {
            try writer.print("{s}\n", .{v});
        } else {
            try writer.writeAll("(nil)\n");
        }
    }

    fn cmdDel(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 1) return ReplError.WrongArgCount;
        if (self.txn) |*t| {
            try t.delete(args[0]);
        } else {
            try self.db.delete(args[0]);
        }
        try writer.writeAll("OK\n");
    }

    fn cmdScan(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len > 1) return ReplError.WrongArgCount;
        const prefix = if (args.len == 1) args[0] else "";
        var iter = try self.db.scanPrefix(self.allocator, prefix);
        defer iter.deinit();
        while (try iter.next()) |entry| {
            try writer.print("{s}\t{s}\n", .{ entry.key, entry.value });
        }
    }

    /// `expire`/`persist`/`ttl`/`expire-sweep` mirror the one-shot CLI's
    /// same-named commands (main.zig) over `ttl.zig`'s flat-key TTL layer.
    /// Like `scan` above, these always act on the live, already-committed
    /// `Db` regardless of any open transaction: a TTL marker lives under its
    /// own keyspace tag, set/cleared directly against `self.db`, never
    /// staged through `Transaction` - there is no OCC-safe way to make
    /// "expire this key in N seconds" conditional on a commit that hasn't
    /// happened yet, so (deliberately, like `coll`'s ops) this bypasses the
    /// active transaction entirely rather than pretending to participate in
    /// it.
    fn cmdExpire(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 2) return ReplError.WrongArgCount;
        const seconds = std.fmt.parseInt(u64, args[1], 10) catch return ReplError.InvalidArgument;
        const existing = try self.db.get(self.allocator, args[0]) orelse return ReplError.NoSuchKey;
        self.allocator.free(existing);
        try ttlmod.setExpireSeconds(self.db, self.allocator, args[0], std.time.nanoTimestamp(), seconds);
        try writer.writeAll("OK\n");
    }

    fn cmdPersist(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 1) return ReplError.WrongArgCount;
        try ttlmod.clearExpire(self.db, self.allocator, args[0]);
        try writer.writeAll("OK\n");
    }

    fn cmdTtl(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 1) return ReplError.WrongArgCount;
        switch (try ttlmod.remaining(self.db, self.allocator, args[0], std.time.nanoTimestamp())) {
            .noKey => try writer.writeAll("-2\n"),
            .noExpiry => try writer.writeAll("-1\n"),
            .seconds => |s| try writer.print("{d}\n", .{s}),
        }
    }

    fn cmdExpireSweep(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 0) return ReplError.WrongArgCount;
        const reaped = try ttlmod.sweepExpired(self.db, self.allocator, std.time.nanoTimestamp());
        try writer.print("reaped: {d}\n", .{reaped});
    }

    fn cmdStats(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 0) return ReplError.WrongArgCount;
        try self.db.stats().print(writer);
    }

    fn cmdFlush(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 0) return ReplError.WrongArgCount;
        try self.db.flush();
        try writer.writeAll("OK\n");
    }

    fn cmdCompact(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len != 0) return ReplError.WrongArgCount;
        try self.db.compact();
        try writer.writeAll("OK\n");
    }

    fn cmdTable(self: *Repl, args: []const []const u8, writer: anytype) !void {
        if (args.len < 1) return ReplError.WrongArgCount;
        const op = args[0];
        const rest = args[1..];

        if (std.mem.eql(u8, op, "select")) return self.cmdTableSelect(rest, writer);
        if (std.mem.eql(u8, op, "explain")) return self.cmdTableExplain(rest, writer);
        if (std.mem.eql(u8, op, "update")) return self.cmdTableUpdate(rest, writer);
        if (std.mem.eql(u8, op, "delete")) return self.cmdTableDelete(rest, writer);
        if (std.mem.eql(u8, op, "join")) return self.cmdTableJoin(rest, writer);
        if (std.mem.eql(u8, op, "aggregate")) return self.cmdTableAggregate(rest, writer);
        if (std.mem.eql(u8, op, "window")) return self.cmdTableWindow(rest, writer);

        if (rest.len < 2) return ReplError.WrongArgCount;
        const tableName = rest[0];
        const pk = try Value.parseCli(rest[1]);

        if (std.mem.eql(u8, op, "put")) {
            var fields = std.ArrayList(Field).init(self.allocator);
            defer fields.deinit();
            for (rest[2..]) |raw| try fields.append(try parseField(raw));
            try self.tablePut(tableName, pk, fields.items);
            try writer.writeAll("OK\n");
        } else if (std.mem.eql(u8, op, "get")) {
            if (rest.len != 2) return ReplError.WrongArgCount;
            var row = try self.tableGet(tableName, pk) orelse {
                try writer.writeAll("(nil)\n");
                return;
            };
            defer row.deinit();
            try printRow(writer, row);
        } else if (std.mem.eql(u8, op, "del")) {
            if (rest.len != 2) return ReplError.WrongArgCount;
            try self.tableDel(tableName, pk);
            try writer.writeAll("OK\n");
        } else {
            return ReplError.UnknownCommand;
        }
    }