//! ----------------------------------------------------
//! ----------------------------------------------------

const App = @import("app");
const World = @import("world");

const std = @import("std");
const gtl = @import("gtl");
const math = @import("math");
const wgpu = @import("wgpu");

const PipelineCache = @import("pipeline_cache.zig");

const Handle = gtl.Handle;
const Self = @This();

/// ----------------------------------------------------
/// ----------------------------------------------------
pub const Node = struct {
    name: []const u8,
    inputs: std.ArrayList(Handle(Slot)) = .empty,
    outputs: std.ArrayList(Handle(Slot)) = .empty,
    run: *const fn (*RenderCTX) anyerror!void,
};

/// ----------------------------------------------------
/// ----------------------------------------------------
pub const Edge = union(enum) {
    data: struct {
        producer: Handle(Slot),
        consumer: Handle(Slot),
    },
    order: struct {
        from: Handle(Node),
        to: Handle(Node),
    },
};

/// ----------------------------------------------------
/// ----------------------------------------------------
pub const Slot = struct {
    pub const Value = union(enum) {
        buffer: *wgpu.WGPUBufferImpl,
        sampler: *wgpu.WGPUSamplerImpl,
        texture_view: *wgpu.WGPUTextureViewImpl,
    };

    name: []const u8,
    value: ?Value = null,
};

/// ----------------------------------------------------
/// ----------------------------------------------------
pub const RenderCTX = struct {
    // --- CORE ---
    app: *App,
    world: *World,
    graph: *Self,
    node: ?Handle(Node) = null,

    // --- RAW ---
    device: *wgpu.WGPUDeviceImpl,
    encoder: *wgpu.WGPUCommandEncoderImpl,

    /// ----------------------------------------------------
    /// ----------------------------------------------------
    pub fn input(self: *RenderCTX, i: usize) !?Slot.Value {
        const n = self.node orelse return error.NodeNotSet;
        const node = self.graph.nodes.get(n) orelse return error.UnknownNode;
        if (i >= node.inputs.items.len) return error.SlotOutOfBounds;
        const slot = self.graph.slots.get(node.inputs.items[i]) orelse return error.UnknownSlot;
        return slot.value;
    }
};

/// ----------------------------------------------------
/// TODO: create Encoder type instead that will return RenderPass
/// ----------------------------------------------------
pub const RenderPass = struct {
    raw: *wgpu.WGPURenderPassEncoderImpl,

    /// ----------------------------------------------------
    /// ----------------------------------------------------
    pub fn begin(
        view: *wgpu.WGPUTextureViewImpl,
        encoder: *wgpu.WGPUCommandEncoderImpl,
        color: math.Color,
    ) !RenderPass {
        return .{
            .raw = wgpu.wgpuCommandEncoderBeginRenderPass(
                encoder,
                &wgpu.WGPURenderPassDescriptor{
                    .colorAttachmentCount = 1,
                    .colorAttachments = &wgpu.WGPURenderPassColorAttachment{
                        .view = view,
                        .resolveTarget = null,
                        .loadOp = wgpu.WGPULoadOp_Clear,
                        .storeOp = wgpu.WGPUStoreOp_Store,
                        .clearValue = .{
                            .r = color.r,
                            .g = color.g,
                            .b = color.b,
                            .a = color.a,
                        },
                        .depthSlice = wgpu.WGPU_DEPTH_SLICE_UNDEFINED,
                    },
                },
            ) orelse return error.WGPURenderPassFailed,
        };
    }

    /// ----------------------------------------------------
    /// ----------------------------------------------------
    pub fn end(self: *const RenderPass) void {
        wgpu.wgpuRenderPassEncoderEnd(self.raw);
        wgpu.wgpuRenderPassEncoderRelease(self.raw);
    }

    /// ----------------------------------------------------
    /// ----------------------------------------------------
    pub fn bindPipeline(self: *const RenderPass, pipe: *const PipelineCache.GPUPipeline) void {
        wgpu.wgpuRenderPassEncoderSetPipeline(self.raw, pipe.raw);
    }

    /// ----------------------------------------------------
    /// ----------------------------------------------------
    pub fn draw(
        self: *const RenderPass,
        vertex_count: u32,
        instance_count: u32,
    ) void {
        wgpu.wgpuRenderPassEncoderDraw(self.raw, vertex_count, instance_count, 0, 0);
    }
};

//
// FIELDS
//

nodes: gtl.SlotMap(Node),
edges: gtl.SlotMap(Edge),
slots: gtl.SlotMap(Slot),

order: ?[]Handle(Node) = null,
alloc: std.mem.Allocator,

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn init(alloc: std.mem.Allocator) Self {
    return .{
        .nodes = .init(alloc),
        .edges = .init(alloc),
        .slots = .init(alloc),
        .alloc = alloc,
    };
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn deinit(self: *Self) void {
    var it = self.nodes.valueIterator();
    while (it.next()) |node| {
        node.inputs.deinit(self.alloc);
        node.outputs.deinit(self.alloc);
    }

    self.nodes.deinit();
    self.edges.deinit();
    self.slots.deinit();

    if (self.order) |order| {
        self.alloc.free(order);
    }
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn run(
    self: *Self,
    ctx: *RenderCTX,
) !void {
    if (self.order) |order| {
        for (order) |handle| {
            ctx.node = handle;
            const node = self.nodes.get(handle) orelse return error.UnknownNode;
            try node.run(ctx);
        }
    }
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn addNode(self: *Self, node: Node) !Handle(Node) {
    self.order = null;
    return try self.nodes.put(node);
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn addEdge(self: *Self, edge: Edge) !Handle(Edge) {
    self.order = null;
    return try self.edges.put(edge);
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn addSlot(self: *Self, slot: Slot) !Handle(Slot) {
    self.order = null;
    return try self.slots.put(slot);
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn addOrderEdge(self: *Self, from: Handle(Node), to: Handle(Node)) !void {
    _ = try self.addEdge(.{ .order = .{ .from = from, .to = to } });
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn addDataEdge(self: *Self, producer: Handle(Slot), consumer: Handle(Slot)) !void {
    _ = try self.addEdge(.{ .data = .{ .producer = producer, .consumer = consumer } });
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn addInput(self: *Self, node: Handle(Node), slot: Slot) !Handle(Slot) {
    const h = try self.slots.put(slot);
    try self.nodes.get(node).?.inputs.append(self.alloc, h);
    return h;
}

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn compile(self: *Self) !void {
    if (self.order != null) return;

    { // --- VALIDATE EDGES ---
        var it = self.edges.valueIterator();
        while (it.next()) |edge| {
            switch (edge.*) {
                .data => |v| {
                    if (!self.slots.contains(v.producer)) return error.UnknownSlot;
                    if (!self.slots.contains(v.consumer)) return error.UnknownSlot;
                    if (v.producer.eql(v.consumer)) return error.SelfLoop;
                },
                .order => |v| {
                    if (!self.nodes.contains(v.from)) return error.UnknownNode;
                    if (!self.nodes.contains(v.to)) return error.UnknownNode;
                },
            }
        }
    }

    // --- SLOT OWNER ---
    var slot_owner = try self.buildSlotOwners();
    defer slot_owner.deinit();

    // --- TOPO SORT ---
    var indegree: std.AutoHashMap(Handle(Node), u32) = .init(self.alloc);
    defer indegree.deinit();

    var adj: std.AutoHashMap(Handle(Node), std.ArrayList(Handle(Node))) = .init(self.alloc);
    defer {
        var it = adj.valueIterator();
        while (it.next()) |l| l.deinit(self.alloc);
        adj.deinit();
    }

    {
        var it = self.nodes.handleIterator();
        while (it.next()) |handle| {
            try indegree.put(handle, 0);
            try adj.put(handle, .empty);
        }
    }

    {
        var it = self.edges.iterator();
        while (it.next()) |entry| {
            const from = switch (entry.value.*) {
                .data => |v| slot_owner.get(v.producer).?,
                .order => |v| v.from,
            };

            const to = switch (entry.value.*) {
                .data => |v| slot_owner.get(v.consumer).?,
                .order => |v| v.to,
            };

            const d = indegree.getPtr(to).?;
            d.* += 1;
            try adj.getPtr(from).?.append(self.alloc, to);
        }
    }

    // --- NERDY SHIT ---

    var order: std.ArrayList(Handle(Node)) = .empty;
    defer order.deinit(self.alloc);

    var ready: std.ArrayList(Handle(Node)) = .empty;
    defer ready.deinit(self.alloc);

    {
        var it = indegree.iterator();
        while (it.next()) |entry| {
            if (entry.value_ptr.* == 0) try ready.append(self.alloc, entry.key_ptr.*);
        }
    }

    while (ready.pop()) |v| {
        try order.append(self.alloc, v);
        for (adj.get(v).?.items) |m| {
            const d = indegree.getPtr(m).?;
            d.* -= 1;
            if (d.* == 0) try ready.append(self.alloc, m);
        }
    }

    if (order.items.len != self.nodes.len()) return error.CycleDetected;

    // --- RESULT ---
    self.order = try order.toOwnedSlice(self.alloc);
}

/// ----------------------------------------------------
/// ----------------------------------------------------
fn buildSlotOwners(self: *Self) !std.AutoHashMap(Handle(Slot), Handle(Node)) {
    // --- OWNERS ---
    var owners: std.AutoHashMap(Handle(Slot), Handle(Node)) = .init(self.alloc);
    errdefer owners.deinit();

    { // --- PORTS CHECK ---
        var it = self.nodes.iterator();
        while (it.next()) |entry| {
            // --- # ---
            const node_handle = entry.handle;
            const node: *const Node = entry.value;

            // --- ENSURE STUFF ---
            const lists = [_][]const Handle(Slot){ node.inputs.items, node.outputs.items };
            for (lists) |list| for (list) |slot| {
                if (!self.slots.contains(slot)) return error.UnknownSlot;
                const gop = try owners.getOrPut(slot);
                if (gop.found_existing) return error.SlotHasMultipleOwners;
                gop.value_ptr.* = node_handle;
            };
        }
    }

    { // --- ORPHAN CHECK 🥀 ---
        var it = self.slots.handleIterator();
        while (it.next()) |handle| {
            if (!owners.contains(handle)) return error.OrphanSlot;
        }
    }

    return owners;
}