// server — the LSP server core: JSON-RPC dispatch, the initialize
// lifecycle, and the document store.
//
// This module is pure protocol logic: `handle` takes one decoded message
// body and returns what to send back (if anything) and whether to exit —
// no file descriptors, so the tests drive it with byte buffers end to
// end. The driver in main.flash owns the actual stdio loop.
//
// Lifecycle per the spec: before `initialize`, every other request is
// answered with ServerNotInitialized and every notification except
// `exit` is dropped. `shutdown` answers with a null result and arms the
// exit handshake; `exit` then ends the process — code 0 after a
// shutdown, code 1 without one, so a client that dies mid-session is
// distinguishable from a clean stop. Unknown requests get MethodNotFound;
// unknown notifications (including `$/`-prefixed ones) are ignored.
//
// Memory: each message is handled inside the caller's per-message arena —
// the response and all JSON nodes die with it. Only the document store
// lives longer: opened texts are duplicated into the server's own
// allocator and freed on replace and close.

use std
use core
use "check"

pub const Error = error{OutOfMemory}

// One open document: the uri and the full current text, both owned by
// the server's allocator.
pub const Doc = struct {
    uri []u8,
    text []u8,
}

// What `handle` decided: a message body to frame and write — a request's
// response, or a server-initiated notification such as
// publishDiagnostics; null when there is nothing to send — and an exit
// code when the client asked the process to end.
pub const Outcome = struct {
    response ?[]u8,
    exit ?u8,
}

pub const Server = struct {
    // Long-lived allocator — the document store only.
    alloc std.mem.Allocator,
    // Reported as serverInfo.version in the initialize answer.
    version []u8,
    docs core.list.List(Doc),
    initialized bool,
    shutdown_requested bool,
    // True when the client offered utf-8 position encoding and we
    // declared it; diagnostics columns depend on this.
    utf8_positions bool,

    pub fn init(alloc std.mem.Allocator, version []u8) Server {
        return .{ .alloc = alloc, .version = version, .docs = .empty, .initialized = false, .shutdown_requested = false, .utf8_positions = false }
    }

    pub fn deinit(self *mut Server) void {
        for d in self.docs.items {
            self.alloc.free(d.uri)
            self.alloc.free(d.text)
        }
        self.docs.deinit(self.alloc)
    }

    // Handle one message body. `arena` is the per-message arena: the
    // returned response and every intermediate JSON node come from it.
    pub fn handle(self *mut Server, arena std.mem.Allocator, body []u8) Error!Outcome {
        v := core.json.parse(arena, body) catch |err| switch err {
            error.Malformed => return respond(try errorResponse(arena, "null", -32700, "parse error")),
            error.OutOfMemory => return error.OutOfMemory,
        }
        method_v := v.get("method") orelse return respond(try errorResponse(arena, try idJson(arena, v), -32600, "invalid request"))
        var method []u8 = ""
        switch method_v {
            .string => |s| {
                method = s
            },
            else => return respond(try errorResponse(arena, try idJson(arena, v), -32600, "invalid request")),
        }
        is_request := v.get("id") != null

        // The pre-initialize gate: only `initialize` and `exit` pass.
        if !self.initialized && !core.mem.eql(u8, method, "initialize") && !core.mem.eql(u8, method, "exit") {
            if is_request {
                return respond(try errorResponse(arena, try idJson(arena, v), -32002, "server not initialized"))
            }
            return none()
        }

        if core.mem.eql(u8, method, "initialize") {
            self.initialized = true
            self.utf8_positions = wantsUtf8(v.get("params"))
            return respond(try self.initializeResponse(arena, try idJson(arena, v)))
        }
        if core.mem.eql(u8, method, "initialized") {
            return none()
        }
        if core.mem.eql(u8, method, "shutdown") {
            self.shutdown_requested = true
            return respond(try concat3(arena, "{\"jsonrpc\":\"2.0\",\"id\":", try idJson(arena, v), ",\"result\":null}"))
        }
        if core.mem.eql(u8, method, "exit") {
            var code u8 = 1
            if self.shutdown_requested {
                code = 0
            }
            return Outcome{ .response = null, .exit = code }
        }
        if core.mem.eql(u8, method, "textDocument/didOpen") {
            if textDocument(v) |td| {
                if stringField(td, "uri") |uri| {
                    if stringField(td, "text") |txt| {
                        try self.upsert(uri, txt)
                        return respond(try check.publishJson(arena, uri, self.text(uri).?, false))
                    }
                }
            }
            return none()
        }
        if core.mem.eql(u8, method, "textDocument/didChange") {
            changed := try self.applyChange(v)
            if changed |uri| {
                return respond(try check.publishJson(arena, uri, self.text(uri).?, false))
            }
            return none()
        }
        if core.mem.eql(u8, method, "textDocument/didSave") {
            // The on-save pass adds sema on top of the parse.
            if textDocument(v) |td| {
                if stringField(td, "uri") |uri| {
                    if self.text(uri) |src| {
                        return respond(try check.publishJson(arena, uri, src, true))
                    }
                }
            }
            return none()
        }
        if core.mem.eql(u8, method, "textDocument/didClose") {
            if textDocument(v) |td| {
                if stringField(td, "uri") |uri| {
                    self.close(uri)
                    // An empty publish clears the file's stale squiggles.
                    return respond(try check.publishJson(arena, uri, "", false))
                }
            }
            return none()
        }

        if is_request {
            return respond(try errorResponse(arena, try idJson(arena, v), -32601, "method not found"))
        }
        return none()
    }

    // The current text of `uri`, or null when it is not open.
    pub fn text(self *Server, uri []u8) ?[]u8 {
        for d in self.docs.items {
            if core.mem.eql(u8, d.uri, uri) {
                return d.text
            }
        }
        return null
    }

    fn initializeResponse(self *Server, arena std.mem.Allocator, id []u8) Error![]u8 {
        var out core.list.List(u8) = .empty
        try out.appendSlice(arena, "{\"jsonrpc\":\"2.0\",\"id\":")
        try out.appendSlice(arena, id)
        try out.appendSlice(arena, ",\"result\":{\"capabilities\":{\"textDocumentSync\":1")
        if self.utf8_positions {
            try out.appendSlice(arena, ",\"positionEncoding\":\"utf-8\"")
        }
        try out.appendSlice(arena, "},\"serverInfo\":{\"name\":\"flashd\",\"version\":\"")
        try out.appendSlice(arena, self.version)
        try out.appendSlice(arena, "\"}}}")
        return out.toOwnedSlice(arena)
    }

    // Full-document sync: replace the stored text with the last entry of
    // contentChanges and report the uri, or null when the message does
    // not carry one — we declared Full sync, so a range-only change
    // cannot be applied.
    fn applyChange(self *mut Server, v core.json.Value) Error!?[]u8 {
        td := textDocument(v) orelse return null
        uri := stringField(td, "uri") orelse return null
        params := v.get("params") orelse return null
        changes_v := params.get("contentChanges") orelse return null
        switch changes_v {
            .array => |xs| {
                if xs.len == 0 {
                    return null
                }
                new_text := stringField(xs[xs.len - 1], "text") orelse return null
                try self.upsert(uri, new_text)
                return uri
            },
            else => return null,
        }
    }

    fn upsert(self *mut Server, uri []u8, new_text []u8) Error!void {
        for d, i in self.docs.items {
            if core.mem.eql(u8, d.uri, uri) {
                dup := try self.alloc.dupe(u8, new_text)
                self.alloc.free(self.docs.items[i].text)
                self.docs.items[i].text = dup
                return
            }
        }
        u := try self.alloc.dupe(u8, uri)
        errdefer self.alloc.free(u)
        t := try self.alloc.dupe(u8, new_text)
        errdefer self.alloc.free(t)
        try self.docs.append(self.alloc, .{ .uri = u, .text = t })
    }

    fn close(self *mut Server, uri []u8) void {
        for d, i in self.docs.items {
            if core.mem.eql(u8, d.uri, uri) {
                self.alloc.free(d.uri)
                self.alloc.free(d.text)
                _ = self.docs.swapRemove(i)
                return
            }
        }
    }
}

fn respond(body []u8) Outcome {
    return .{ .response = body, .exit = null }
}

fn none() Outcome {
    return .{ .response = null, .exit = null }
}

// The request id re-serialized exactly as it arrived (number or string;
// "null" when absent) — safe to splice into a response.
fn idJson(arena std.mem.Allocator, v core.json.Value) Error![]u8 {
    idv := v.get("id") orelse return "null"
    return core.json.stringify(arena, idv) catch |err| switch err {
        error.OutOfMemory => return error.OutOfMemory,
        error.Malformed => return "null",
    }
}

fn errorResponse(arena std.mem.Allocator, id []u8, code i64, msg []u8) Error![]u8 {
    var out core.list.List(u8) = .empty
    try out.appendSlice(arena, "{\"jsonrpc\":\"2.0\",\"id\":")
    try out.appendSlice(arena, id)
    try out.appendSlice(arena, ",\"error\":{\"code\":")
    var buf [32]u8 = undefined
    try out.appendSlice(arena, core.fmt.bufPrint(buf[0..], "{d}", .{code}) catch unreachable)
    try out.appendSlice(arena, ",\"message\":\"")
    try out.appendSlice(arena, msg)
    try out.appendSlice(arena, "\"}}")
    return out.toOwnedSlice(arena)
}

fn concat3(arena std.mem.Allocator, a []u8, b []u8, c []u8) Error![]u8 {
    var out core.list.List(u8) = .empty
    try out.appendSlice(arena, a)
    try out.appendSlice(arena, b)
    try out.appendSlice(arena, c)
    return out.toOwnedSlice(arena)
}

// params.textDocument of `v`, when present.
fn textDocument(v core.json.Value) ?core.json.Value {
    params := v.get("params") orelse return null
    return params.get("textDocument")
}

fn stringField(v core.json.Value, key []u8) ?[]u8 {
    f := v.get(key) orelse return null
    switch f {
        .string => |s| return s,
        else => return null,
    }
}

// Did the client offer utf-8 position encoding (LSP 3.17,
// capabilities.general.positionEncodings)?
fn wantsUtf8(params ?core.json.Value) bool {
    p := params orelse return false
    caps := p.get("capabilities") orelse return false
    g := caps.get("general") orelse return false
    encs := g.get("positionEncodings") orelse return false
    switch encs {
        .array => |xs| {
            for x in xs {
                switch x {
                    .string => |s| {
                        if core.mem.eql(u8, s, "utf-8") {
                            return true
                        }
                    },
                    else => {},
                }
            }
        },
        else => {},
    }
    return false
}

fn contains(haystack []u8, needle []u8) bool {
    return core.mem.indexOf(u8, haystack, needle) != null
}

test "initialize answers capabilities and arms the session" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0.0.0-test")
    defer srv.deinit()

    o := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"capabilities\":{\"general\":{\"positionEncodings\":[\"utf-16\",\"utf-8\"]}}}}")
    r := o.response.?
    try std.testing.expect(contains(r, "\"id\":1"))
    try std.testing.expect(contains(r, "\"textDocumentSync\":1"))
    try std.testing.expect(contains(r, "\"positionEncoding\":\"utf-8\""))
    try std.testing.expect(contains(r, "\"name\":\"flashd\""))
    try std.testing.expect(contains(r, "\"version\":\"0.0.0-test\""))
    try std.testing.expect(srv.initialized)
    try std.testing.expect(srv.utf8_positions)
    try std.testing.expect(o.exit == null)
}

test "without a utf-8 offer the encoding stays at the protocol default" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()

    o := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"capabilities\":{}}}")
    try std.testing.expect(!contains(o.response.?, "positionEncoding"))
    try std.testing.expect(!srv.utf8_positions)
}

test "requests before initialize are refused, notifications dropped" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()

    o := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":7,\"method\":\"shutdown\"}")
    try std.testing.expect(contains(o.response.?, "-32002"))
    n := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didOpen\",\"params\":{}}")
    try std.testing.expect(n.response == null)
    try std.testing.expect(n.exit == null)
}

test "the shutdown-exit handshake reports a clean zero" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()

    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{}}")
    s := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"shutdown\"}")
    try std.testing.expect(contains(s.response.?, "\"result\":null"))
    e := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"exit\"}")
    try std.testing.expectEqual(0, e.exit.?)
}

test "exit without shutdown reports failure" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()

    e := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"exit\"}")
    try std.testing.expectEqual(1, e.exit.?)
}

test "unknown requests get MethodNotFound, unknown notifications nothing" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()

    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{}}")
    u := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":\"abc\",\"method\":\"textDocument/hover\"}")
    try std.testing.expect(contains(u.response.?, "-32601"))
    try std.testing.expect(contains(u.response.?, "\"id\":\"abc\""))
    n := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"$/cancelRequest\",\"params\":{\"id\":1}}")
    try std.testing.expect(n.response == null)
}

test "unparseable bodies answer ParseError with a null id" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()

    o := try srv.handle(a.allocator(), "{nope")
    try std.testing.expect(contains(o.response.?, "-32700"))
    try std.testing.expect(contains(o.response.?, "\"id\":null"))
}

test "didOpen, didChange, and didClose drive the document store" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()

    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{}}")
    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didOpen\",\"params\":{\"textDocument\":{\"uri\":\"file:///a.flash\",\"text\":\"fn one\"}}}")
    try std.testing.expect(core.mem.eql(u8, srv.text("file:///a.flash").?, "fn one"))

    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didChange\",\"params\":{\"textDocument\":{\"uri\":\"file:///a.flash\"},\"contentChanges\":[{\"text\":\"fn two\"}]}}")
    try std.testing.expect(core.mem.eql(u8, srv.text("file:///a.flash").?, "fn two"))

    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didClose\",\"params\":{\"textDocument\":{\"uri\":\"file:///a.flash\"}}}")
    try std.testing.expect(srv.text("file:///a.flash") == null)
    try std.testing.expectEqual(0, srv.docs.items.len)
}

test "the open-edit-save cycle publishes and clears diagnostics" {
    var a = core.arena.ArenaAllocator.init(std.testing.allocator)
    defer a.deinit()
    var srv = Server.init(std.testing.allocator, "0")
    defer srv.deinit()
    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{}}")

    // Open with a parse error: the publish carries a severity-1 entry.
    o := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didOpen\",\"params\":{\"textDocument\":{\"uri\":\"file:///d.flash\",\"text\":\"fn broken(\"}}}")
    try std.testing.expect(contains(o.response.?, "publishDiagnostics"))
    try std.testing.expect(contains(o.response.?, "\"severity\":1"))

    // Edit to clean source: the publish is empty — squiggles clear.
    c := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didChange\",\"params\":{\"textDocument\":{\"uri\":\"file:///d.flash\"},\"contentChanges\":[{\"text\":\"fn ok() void {\\n}\\n\"}]}}")
    try std.testing.expect(contains(c.response.?, "\"diagnostics\":[]"))

    // A parseable sema error stays quiet on change, loud on save.
    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didChange\",\"params\":{\"textDocument\":{\"uri\":\"file:///d.flash\"},\"contentChanges\":[{\"text\":\"fn f() usize {\\n    return 1 / 0\\n}\\n\"}]}}")
    s := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didSave\",\"params\":{\"textDocument\":{\"uri\":\"file:///d.flash\"}}}")
    try std.testing.expect(contains(s.response.?, "\"severity\":1"))

    // Close: one final empty publish.
    x := try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didClose\",\"params\":{\"textDocument\":{\"uri\":\"file:///d.flash\"}}}")
    try std.testing.expect(contains(x.response.?, "\"diagnostics\":[]"))
}

fn handleSweep(alloc std.mem.Allocator) !void {
    var a = core.arena.ArenaAllocator.init(alloc)
    defer a.deinit()
    var srv = Server.init(alloc, "0")
    defer srv.deinit()
    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{}}")
    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didOpen\",\"params\":{\"textDocument\":{\"uri\":\"file:///s.flash\",\"text\":\"x\"}}}")
    _ = try srv.handle(a.allocator(), "{\"jsonrpc\":\"2.0\",\"method\":\"textDocument/didChange\",\"params\":{\"textDocument\":{\"uri\":\"file:///s.flash\"},\"contentChanges\":[{\"text\":\"y\"}]}}")
}

test "handle survives failure at every allocation point" {
    try std.testing.checkAllAllocationFailures(std.testing.allocator, handleSweep, .{})
}