// Selection
// SPDX-FileCopyrightText: 2025 Nguyễn Gia Phong
// SPDX-License-Identifier: GPL-3.0-or-later

const assert = std.debug.assert;
const eql = std.meta.eql;
const order = std.math.order;
const std = @import("std");

const Node = tree_sitter.Node;
const Tree = tree_sitter.Tree;
const tree_sitter = @import("tree-sitter");

const Grapheme = vaxis.Graphemes.Grapheme;
const vaxis = @import("vaxis");

const Buffer = @import("Buffer.zig");
const graphemes = &@import("root").graphemes;

const Selection = @This();
head: Unit,
tail: Unit,

pub fn startByte(selection: Selection) u32 {
    return @min(selection.head.startByte(), selection.tail.startByte());
}

pub fn endByte(selection: Selection) u32 {
    return @max(selection.head.endByte(), selection.tail.endByte());
}

pub fn inHead(selection: Selection, grapheme: Grapheme) bool {
    return selection.head.contains(grapheme);
}

pub fn inBody(selection: Selection, grapheme: Grapheme) bool {
    return (!selection.inHead(grapheme)
            and selection.startByte() <= grapheme.offset
            and grapheme.offset + grapheme.len <= selection.endByte());
}

// TODO: nesting
pub const Span = struct {
    start: u32 = 0,
    end: u32 = 0,
    tree: *const Tree,

    pub fn root(text_len: u32, tree: *const Tree) Span {
        return .{ .start = 0, .end = text_len, .tree = tree };
    }

    fn parent(span: Span, text_len: u32) Unit {
        const root_span = Span.root(text_len, span.tree);
        if (eql(root_span, span))
            return .{ .span = span };
        const root_node = span.tree.rootNode();
        const root_node_start = root_node.startByte();
        const root_node_end = root_node.endByte();
        if (root_node_start == span.start and root_node_end == span.end)
            return .{ .span = root_span };
        if (root_node_start > span.start) {
            const first_span = Span{
                .start = 0,
                .end = root_node_start,
                .tree = span.tree,
            };
            return .{
                .span = if (eql(first_span, span)) root_span else first_span
            };
        }
        if (root_node_end < span.end) {
            const last_span = Span{
                .start = root_node_start,
                .end = text_len,
                .tree = span.tree,
            };
            return .{
                .span = if (eql(last_span, span)) root_span else last_span
            };
        }
        const parent_node = root_node.descendantForByteRange(span.start,
                                                             span.end).?;
        const parent_node_start = parent_node.startByte();
        assert(parent_node_start <= span.start);
        const parent_node_end = parent_node.endByte();
        assert(parent_node_end >= span.end);
        const parent_span = grow: {
            if (parent_node.firstChildForByte(span.end)) |next_node| {
                break :grow if (next_node.prevSibling()) |prev_node| Span{
                    .start = prev_node.endByte(),
                    .end = next_node.startByte(),
                    .tree = span.tree,
                } else Span{
                    .start = parent_node_start,
                    .end = next_node.startByte(),
                    .tree = span.tree,
                };
            } else {
                const siblings = parent_node.childCount();
                if (siblings == 0)
                    return .{ .node = parent_node };
                const prev_node = parent_node.child(siblings - 1).?;
                break :grow Span{
                    .start = prev_node.endByte(),
                    .end = parent_node_end,
                    .tree = span.tree,
                };
            }
        };
        return if (eql(parent_span, span)) .{
            .node = parent_node
        } else .{
            .span = parent_span
        };
    }
};

const GraphemeUnit = struct {
    slice: Grapheme,
    tree: *const Tree,

    fn startByte(unit: GraphemeUnit) u32 {
        return unit.slice.offset;
    }

    fn endByte(unit: GraphemeUnit) u32 {
        return unit.slice.offset + unit.slice.len;
    }

    fn parent(unit: GraphemeUnit, text_len: u32) Unit {
        return Span.parent(.{
            .start = unit.startByte(),
            .end = unit.endByte(),
            .tree = unit.tree,
        }, text_len);
    }

    fn at(unit: GraphemeUnit, buffer: Buffer, index: u32) GraphemeUnit {
        return .{
            .slice = buffer.graphemeAt(index),
            .tree = unit.tree,
        };
    }

    fn before(unit: GraphemeUnit, buffer: Buffer, index: u32) ?GraphemeUnit {
        return if (buffer.graphemeBefore(index)) |grapheme| .{
            .slice = grapheme,
            .tree = unit.tree,
        } else null;
    }
};

// TODO: undo
pub const Unit = union(enum) {
    node: Node,
    span: Span,
    grapheme: GraphemeUnit,

    fn startByte(unit: Unit) u32 {
        return switch (unit) {
            .node => |node| node.startByte(),
            .span => |span| span.start,
            .grapheme => |grapheme| grapheme.startByte(),
        };
    }

    fn endByte(unit: Unit) u32 {
        return switch (unit) {
            .node => |node| node.endByte(),
            .span => |span| span.end,
            .grapheme => |grapheme| grapheme.endByte(),
        };
    }

    fn tree(unit: Unit) *const Tree {
        return switch (unit) {
            .node => |node| node.tree,
            .span => |span| span.tree,
            .grapheme => |grapheme| grapheme.tree,
        };
    }

    fn equal(unit: Unit, other: Unit) bool {
        return (unit.startByte() == other.startByte()
                and unit.endByte() == other.endByte());
    }

    pub fn up(unit: Unit, buffer: Buffer) Unit {
        const parent = switch (unit) {
            .node => |node| if (node.parent()) |parent| Unit{
                .node = parent
            } else Span.parent(.{
                .start = node.startByte(),
                .end = node.endByte(),
                .tree = node.tree,
            }, buffer.length()),
            .span => |span| span.parent(buffer.length()),
            .grapheme => |grapheme| grapheme.parent(buffer.length()),
        };
        return if (!parent.equal(unit) or parent.equal(.{
            .span = .root(buffer.length(), unit.tree())
        })) parent else parent.up(buffer);
    }

    pub fn right(unit: Unit, buffer: Buffer) Unit {
        switch (unit) {
            .node => |node| {
                const prev = node.endByte();
                const root_node = node.tree.rootNode();
                if (root_node.eql(node)) {
                    const span = Span{
                        .start = prev,
                        .end = buffer.length(),
                        .tree = node.tree,
                    };
                    if (span.start < span.end)
                        return .{ .span = span };
                } else if (node.nextSibling()) |sibling| {
                    const next = sibling.startByte();
                    return switch (order(prev, next)) {
                        .lt => .{
                            .span = .{
                                .start = prev,
                                .end = next,
                                .tree = node.tree,
                            }
                        },
                        .eq => .{ .node = sibling },
                        .gt => unreachable,
                    };
                }
            },
            .span => |span| switch (unit.up(buffer)) {
                .node => |parent_node| {
                    if (parent_node.firstChildForByte(span.end)) |next_node|
                        return .{ .node = next_node };
                },
                .span => |parent_span| {
                    const root_node = span.tree.rootNode();
                    if (root_node.startByte() == span.end)
                        return .{ .node = root_node };
                    const next_span = Span{
                        .start = span.end,
                        .end = parent_span.end,
                        .tree = span.tree,
                    };
                    if (next_span.start < next_span.end)
                        return .{ .span = next_span };
                },
                .grapheme => unreachable,
            },
            .grapheme => |grapheme| {
                const next_grapheme = grapheme.at(buffer, grapheme.endByte());
                const len = buffer.length();
                if (next_grapheme.parent(len).equal(grapheme.parent(len)))
                    return .{ .grapheme = next_grapheme };
                return .right(.{
                    .span = .{
                        .start = grapheme.startByte(),
                        .end = grapheme.endByte(),
                        .tree = grapheme.tree,
                    }
                }, buffer);
            },
        }
        return unit.up(buffer);
    }

    pub fn left(unit: Unit, buffer: Buffer) Unit {
        switch (unit) {
            .node => |node| {
                const next = node.startByte();
                const root_node = node.tree.rootNode();
                if (root_node.eql(node)) {
                    const span = Span{
                        .start = 0,
                        .end = next,
                        .tree = node.tree,
                    };
                    if (span.start < span.end)
                        return .{ .span = span };
                } else if (node.prevSibling()) |sibling| {
                    const prev = sibling.endByte();
                    return switch (order(prev, next)) {
                        .lt => .{
                            .span = .{
                                .start = prev,
                                .end = next,
                                .tree = node.tree,
                            }
                        },
                        .eq => .{ .node = sibling },
                        .gt => unreachable,
                    };
                }
            },
            .span => |span| switch (unit.up(buffer)) {
                .node => |parent_node| {
                    if (parent_node.firstChildForByte(span.end)) |next_node| {
                        if (next_node.prevSibling()) |prev_node|
                            return .{ .node = prev_node };
                    } else {
                        const siblings = parent_node.childCount();
                        if (siblings > 0)
                            return .{
                                .node =  parent_node.child(siblings - 1).?
                            };
                    }
                },
                .span => |parent_span| {
                    const root_node = span.tree.rootNode();
                    if (root_node.endByte() == span.start)
                        return .{ .node = root_node };
                    const prev_span = Span{
                        .start = span.start,
                        .end = parent_span.start,
                        .tree = span.tree,
                    };
                    if (prev_span.start < prev_span.end)
                        return .{ .span = prev_span };
                },
                .grapheme => unreachable,
            },
            .grapheme => |grapheme| {
                const start = grapheme.startByte();
                if (grapheme.before(buffer, start)) |prev_grapheme| {
                    const len = buffer.length();
                    if (prev_grapheme.parent(len).equal(grapheme.parent(len)))
                        return .{ .grapheme = prev_grapheme };
                }
                return .left(.{
                    .span = .{
                        .start = start,
                        .end = grapheme.endByte(),
                        .tree = grapheme.tree,
                    }
                }, buffer);
            },
        }
        return unit.up(buffer);
    }

    pub fn down(unit: Unit, buffer: Buffer) Unit {
        switch (unit) {
            .node => |node| {
                const start = node.startByte();
                const end = node.endByte();
                // TODO: history
                return if (node.child(0)) |child|
                    if (child.startByte() == start and child.endByte() == end)
                        .down(.{ .node = child }, buffer)
                    else
                        .{ .node = child }
                else .down(.{
                    .span = .{ .start = start, .end = end, .tree = node.tree }
                }, buffer);
            },
            .span => |span| return .{
                .grapheme = .{
                    // TODO: history
                    .slice = buffer.graphemeAt(span.start),
                    .tree = span.tree,
                }
            },
            .grapheme => return unit,
        }
    }

    pub fn contains(unit: Unit, other: Grapheme) bool {
        return switch (unit) {
            .node => |node| (node.startByte() <= other.offset
                             and other.offset + other.len <= node.endByte()),
            .span => |span| (span.start <= other.offset
                             and other.offset + other.len <= span.end),
            .grapheme => |grapheme| eql(grapheme.slice, other),
        };
    }
};