//! SSH server that serves a matcha `Model` — one fresh model instance per
//! connection, wired to the SSH channel through matcha's injectable I/O seam.
//! Mirrors charmbracelet/wish + its bubbletea middleware, collapsed into a
//! single "wrap a model" entry point.
//!
//! Architecture: the connection fiber owns libssh (it alone calls
//! `ssh_event_dopoll` / `ssh_channel_write`). It runs the handshake, then
//! builds a `matcha.Program` reading from / writing to the `Session`'s byte
//! buffers and runs it on separate fibers. Input pushed by libssh callbacks
//! flows to the program's input fiber; rendered frames flow back through
//! `out_buf`, which the connection fiber flushes each poll.

const std = @import("std");
const Io = std.Io;
const Allocator = std.mem.Allocator;
const Value = std.atomic.Value;
const c = @import("libssh");
const matcha = @import("matcha");
const Options = @import("Options.zig");
const Session = @import("Session.zig");
const Pty = @import("Pty.zig").Pty;
const Window = @import("Pty.zig").Window;

const log = std.log.scoped(.dream);

const SshBind = c.ssh_bind_struct;
const SshSession = c.ssh_session_struct;
const SshChannel = c.ssh_channel_struct;
const SshServerCallbacks = c.ssh_server_callbacks_struct;
const SshChannelCallbacks = c.ssh_channel_callbacks_struct;
const SshKey = c.ssh_key_struct;

/// Build a server that serves `Model` over SSH. `interface` is the name of the
/// `matcha.Model` field embedded in `Model` (usually `"model"`). Each accepted
/// connection default-constructs a fresh `Model` and runs it.
pub fn serve(
    comptime Model: type,
    comptime interface: []const u8,
    process: std.process.Init,
    options: Options,
) !Server(Model, interface) {
    return Server(Model, interface).init(process, options);
}

pub fn Server(comptime Model: type, comptime interface: []const u8) type {
    comptime {
        if (!@hasField(Model, interface)) @compileError(
            "dream.serve: '" ++ @typeName(Model) ++ "' has no field '" ++ interface ++ "'",
        );
        if (@FieldType(Model, interface) != matcha.Model) @compileError(
            "dream.serve: field '" ++ interface ++ "' of '" ++ @typeName(Model) ++
                "' must be a matcha.Model",
        );
    }
    return struct {
        const Self = @This();

        process: std.process.Init,
        options: Options,
        bind: *SshBind,
        running: Value(bool) = .init(false),

        pub fn init(process: std.process.Init, options: Options) !Self {
            if (c.ssh_init() != c.SSH_OK) return error.LibsshInit;
            const bind = c.ssh_bind_new() orelse return error.OutOfMemory;
            errdefer c.ssh_bind_free(bind);

            var host_buf: [64]u8 = undefined;
            var port_buf: [8]u8 = undefined;
            try setBindOption(bind, c.SSH_BIND_OPTIONS_BINDADDR, try options.hostZ(&host_buf));
            try setBindOption(bind, c.SSH_BIND_OPTIONS_BINDPORT_STR, try options.portZ(&port_buf));
            if (options.banner.len > 0)
                try setBindOption(bind, c.SSH_BIND_OPTIONS_BANNER, options.banner);

            if (options.host_key_pem) |pem| {
                try setBindOption(bind, c.SSH_BIND_OPTIONS_IMPORT_KEY_STR, pem);
            } else if (options.host_key_path) |path| {
                try setBindOption(bind, c.SSH_BIND_OPTIONS_HOSTKEY, path);
            } else return error.HostKeyRequired;

            return .{ .process = process, .options = options, .bind = bind };
        }

        pub fn deinit(self: *Self) void {
            c.ssh_bind_free(self.bind);
        }

        pub fn shutdown(self: *Self) void {
            self.running.store(false, .release);
        }

        pub fn listenAndServe(self: *Self, io: Io) !void {
            if (c.ssh_bind_listen(self.bind) < 0) {
                log.err("ssh_bind_listen: {s}", .{c.ssh_get_error(self.bind)});
                return error.ListenFailed;
            }
            self.running.store(true, .release);

            var group: Io.Group = .init;
            defer group.await(io) catch {};

            while (self.running.load(.acquire)) {
                const ssh_session = c.ssh_new() orelse return error.OutOfMemory;
                if (c.ssh_bind_accept(self.bind, ssh_session) != c.SSH_OK) {
                    log.warn("ssh_bind_accept: {s}", .{c.ssh_get_error(self.bind)});
                    c.ssh_free(ssh_session);
                    continue;
                }
                group.concurrent(io, handleConnection, .{ self, io, ssh_session }) catch |e| {
                    log.err("dispatch failed: {t}", .{e});
                    c.ssh_free(ssh_session);
                };
            }
        }

        fn handleConnection(self: *Self, io: Io, ssh_session: *SshSession) void {
            defer {
                c.ssh_disconnect(ssh_session);
                c.ssh_free(ssh_session);
            }
            self.runConnection(io, ssh_session) catch |e| {
                log.warn("connection error: {t}", .{e});
            };
        }

        fn runConnection(self: *Self, io: Io, ssh_session: *SshSession) !void {
            const gpa = self.process.gpa;
            var state: ConnState = .{
                .allocator = gpa,
                .io = io,
                .ssh_session = ssh_session,
            };
            defer if (state.user_name.len > 0) gpa.free(state.user_name);
            defer if (state.pty_term) |t| gpa.free(t);
            defer if (state.command_line) |cmd| gpa.free(cmd);
            defer if (state.command_slots) |slots| gpa.free(slots);
            defer if (state.session) |*s| s.deinit();

            initServerCallbacks(&state);
            if (c.ssh_set_server_callbacks(ssh_session, &state.server_cb) != c.SSH_OK)
                return error.CallbacksFailed;

            if (c.ssh_handle_key_exchange(ssh_session) != c.SSH_OK) {
                log.warn("key exchange failed: {s}", .{c.ssh_get_error(ssh_session)});
                return error.KeyExchangeFailed;
            }

            c.ssh_set_auth_methods(
                ssh_session,
                c.SSH_AUTH_METHOD_PASSWORD | c.SSH_AUTH_METHOD_PUBLICKEY | c.SSH_AUTH_METHOD_NONE,
            );

            const event = c.ssh_event_new() orelse return error.OutOfMemory;
            defer c.ssh_event_free(event);
            if (c.ssh_event_add_session(event, ssh_session) != c.SSH_OK)
                return error.EventAddFailed;

            // Drive the handshake until a shell/exec request lands us in
            // `.running` (or the client bails out to `.done`).
            while (state.phase != .running and state.phase != .done) {
                if (c.ssh_event_dopoll(event, 200) == c.SSH_ERROR) {
                    log.warn("ssh_event_dopoll: {s}", .{c.ssh_get_error(ssh_session)});
                    return error.PollError;
                }
            }
            if (state.phase == .done) return;

            try self.runModel(io, event, &state.session.?);
        }

        /// Wrap a fresh `Model` in a `matcha.Program` bound to the session, run
        /// it, and pump libssh until it exits.
        fn runModel(self: *Self, io: Io, event: c.ssh_event, session: *Session) !void {
            var model: Model = .{};
            const iface: *matcha.Model = &@field(model, interface);

            const win: ?matcha.Terminal.Size = if (session.pty()) |p| .{
                .width = @intCast(p.window.width),
                .height = @intCast(p.window.height),
            } else null;

            var program = try matcha.Program.init(self.process, iface, .{
                .input = session.reader(),
                .output = session.writer(),
                .window_size = win,
            });
            defer program.deinit();

            // Forward resize events; quit the program when the session closes.
            var watcher = try io.concurrent(watchSession, .{ session, &program });

            // Run the program on its own task so this fiber keeps pumping libssh.
            var prog_done: Value(bool) = .init(false);
            var run_task = try io.concurrent(runProgram, .{ &program, &prog_done });

            while (!prog_done.load(.acquire)) {
                session.drainOutput();
                if (c.ssh_event_dopoll(event, 50) == c.SSH_ERROR) {
                    log.warn("dopoll mid-session: {s}", .{c.ssh_get_error(session.ssh_session)});
                    session.close();
                    break;
                }
            }
            run_task.await(io);
            session.drainOutput(); // final frame(s)

            session.stopWindowEvents();
            watcher.await(io);

            _ = c.ssh_channel_request_send_exit_status(session.channel, @intCast(session.exit_status));
            _ = c.ssh_channel_send_eof(session.channel);
            _ = c.ssh_channel_close(session.channel);
        }
    };
}

// ─── program tasks (non-generic) ────────────────────────────────────────────

fn runProgram(program: *matcha.Program, done: *Value(bool)) void {
    defer done.store(true, .release);
    program.run() catch |e| log.warn("program exited with error: {t}", .{e});
}

fn watchSession(session: *Session, program: *matcha.Program) void {
    while (session.nextWindow()) |win| {
        program.send(.{ .window_size = .{
            .width = @intCast(win.width),
            .height = @intCast(win.height),
        } });
    }
    // Window queue closed → session ending → tell the program to quit.
    program.quit();
}

// ─── connection state + libssh callbacks (non-generic) ──────────────────────

const Phase = enum { authenticating, await_channel, await_shell, running, done };

const ConnState = struct {
    allocator: Allocator,
    io: Io,

    ssh_session: *SshSession,
    channel: ?*SshChannel = null,
    user_name: []const u8 = "",

    pty_term: ?[]u8 = null,
    command_line: ?[]u8 = null,
    command_slots: ?[][]const u8 = null,

    phase: Phase = .authenticating,

    session: ?Session = null,

    server_cb: SshServerCallbacks = std.mem.zeroes(SshServerCallbacks),
    channel_cb: SshChannelCallbacks = std.mem.zeroes(SshChannelCallbacks),
};

fn initServerCallbacks(state: *ConnState) void {
    state.server_cb.size = @sizeOf(SshServerCallbacks);
    state.server_cb.userdata = state;
    state.server_cb.auth_password_function = cbAuthPassword;
    state.server_cb.auth_pubkey_function = cbAuthPubkey;
    state.server_cb.auth_none_function = cbAuthNone;
    state.server_cb.channel_open_request_session_function = cbChannelOpen;

    state.channel_cb.size = @sizeOf(SshChannelCallbacks);
    state.channel_cb.userdata = state;
    state.channel_cb.channel_data_function = cbChannelData;
    state.channel_cb.channel_eof_function = cbChannelEof;
    state.channel_cb.channel_close_function = cbChannelClose;
    state.channel_cb.channel_pty_request_function = cbPtyRequest;
    state.channel_cb.channel_shell_request_function = cbShellRequest;
    state.channel_cb.channel_exec_request_function = cbExecRequest;
    state.channel_cb.channel_pty_window_change_function = cbWindowChange;
}

fn castState(userdata: ?*anyopaque) *ConnState {
    return @ptrCast(@alignCast(userdata.?));
}

fn cbAuthNone(_: ?*SshSession, user: [*c]const u8, userdata: ?*anyopaque) callconv(.c) c_int {
    const state = castState(userdata);
    rememberUser(state, user);
    return c.SSH_AUTH_DENIED;
}

fn cbAuthPassword(
    _: ?*SshSession,
    user: [*c]const u8,
    _: [*c]const u8,
    userdata: ?*anyopaque,
) callconv(.c) c_int {
    const state = castState(userdata);
    rememberUser(state, user);
    state.phase = .await_channel;
    return c.SSH_AUTH_SUCCESS;
}

fn cbAuthPubkey(
    _: ?*SshSession,
    user: [*c]const u8,
    _: ?*SshKey,
    sig_state: u8,
    userdata: ?*anyopaque,
) callconv(.c) c_int {
    const state = castState(userdata);
    rememberUser(state, user);
    if (sig_state == c.SSH_PUBLICKEY_STATE_NONE) return c.SSH_AUTH_SUCCESS;
    if (sig_state == c.SSH_PUBLICKEY_STATE_VALID) {
        state.phase = .await_channel;
        return c.SSH_AUTH_SUCCESS;
    }
    return c.SSH_AUTH_DENIED;
}

fn cbChannelOpen(session: ?*SshSession, userdata: ?*anyopaque) callconv(.c) ?*SshChannel {
    const state = castState(userdata);
    if (state.channel != null) return null;
    const ch = c.ssh_channel_new(session) orelse return null;
    state.channel = ch;
    state.phase = .await_shell;

    state.session = Session.init(state.io, state.allocator, state.ssh_session, ch) catch {
        c.ssh_channel_free(ch);
        state.channel = null;
        return null;
    };
    state.session.?.user_name = state.user_name;

    if (c.ssh_set_channel_callbacks(ch, &state.channel_cb) != c.SSH_OK)
        log.warn("ssh_set_channel_callbacks failed", .{});
    return ch;
}

fn cbChannelData(
    _: ?*SshSession,
    _: ?*SshChannel,
    data: ?*anyopaque,
    len: u32,
    _: c_int, // is_stderr — server-side, client data is always stdin
    userdata: ?*anyopaque,
) callconv(.c) c_int {
    const state = castState(userdata);
    if (state.session == null) return 0;
    const bytes: [*]const u8 = @ptrCast(data.?);
    // `&state.session.?` addresses the stored Session in place — not a copy of
    // the unwrapped optional, which would drop the pushed input on the floor.
    const n = state.session.?.pushInput(bytes[0..len]) catch return 0;
    return @intCast(n);
}

fn cbChannelEof(_: ?*SshSession, _: ?*SshChannel, userdata: ?*anyopaque) callconv(.c) void {
    if (castState(userdata).session) |*s| s.markEof();
}

fn cbChannelClose(_: ?*SshSession, _: ?*SshChannel, userdata: ?*anyopaque) callconv(.c) void {
    const state = castState(userdata);
    if (state.session) |*s| s.close();
    state.phase = .done;
}

fn cbPtyRequest(
    _: ?*SshSession,
    _: ?*SshChannel,
    term: [*c]const u8,
    width: c_int,
    height: c_int,
    pxwidth: c_int,
    pxheight: c_int,
    userdata: ?*anyopaque,
) callconv(.c) c_int {
    const state = castState(userdata);
    if (state.session) |*s| {
        const term_src = stringOrEmpty(term);
        if (term_src.len > 0) state.pty_term = state.allocator.dupe(u8, term_src) catch null;
        s.pty_value = .{
            .term = state.pty_term orelse "",
            .window = .{
                .width = @intCast(@max(0, width)),
                .height = @intCast(@max(0, height)),
                .pixel_width = @intCast(@max(0, pxwidth)),
                .pixel_height = @intCast(@max(0, pxheight)),
            },
        };
        s.last_window = s.pty_value.?.window;
    }
    return c.SSH_OK;
}

fn cbWindowChange(
    _: ?*SshSession,
    _: ?*SshChannel,
    width: c_int,
    height: c_int,
    pxwidth: c_int,
    pxheight: c_int,
    userdata: ?*anyopaque,
) callconv(.c) c_int {
    const state = castState(userdata);
    if (state.session) |*s| s.setWindow(.{
        .width = @intCast(@max(0, width)),
        .height = @intCast(@max(0, height)),
        .pixel_width = @intCast(@max(0, pxwidth)),
        .pixel_height = @intCast(@max(0, pxheight)),
    });
    return c.SSH_OK;
}

fn cbShellRequest(_: ?*SshSession, _: ?*SshChannel, userdata: ?*anyopaque) callconv(.c) c_int {
    castState(userdata).phase = .running;
    return c.SSH_OK;
}

fn cbExecRequest(
    _: ?*SshSession,
    _: ?*SshChannel,
    command: [*c]const u8,
    userdata: ?*anyopaque,
) callconv(.c) c_int {
    const state = castState(userdata);
    if (state.session) |*s| {
        const cmd = stringOrEmpty(command);
        if (cmd.len > 0) {
            const owned = state.allocator.dupe(u8, cmd) catch return c.SSH_ERROR;
            const slots = state.allocator.alloc([]const u8, 1) catch {
                state.allocator.free(owned);
                return c.SSH_ERROR;
            };
            slots[0] = owned;
            state.command_line = owned;
            state.command_slots = slots;
            s.command_argv = slots;
        }
    }
    state.phase = .running;
    return c.SSH_OK;
}

fn rememberUser(state: *ConnState, user: [*c]const u8) void {
    if (state.user_name.len != 0) return;
    const u = stringOrEmpty(user);
    if (u.len == 0) return;
    state.user_name = state.allocator.dupe(u8, u) catch "";
    if (state.session) |*s| s.user_name = state.user_name;
}

fn setBindOption(bind: *SshBind, opt: c_uint, value: [:0]const u8) !void {
    if (c.ssh_bind_options_set(bind, opt, value.ptr) < 0) {
        log.err("ssh_bind_options_set({}): {s}", .{ opt, c.ssh_get_error(bind) });
        return error.BindOptionFailed;
    }
}

fn stringOrEmpty(s: ?[*:0]const u8) []const u8 {
    const p = s orelse return "";
    return std.mem.span(p);
}