const std = @import("std");
const mem = std.mem;
const meta = std.meta;
const math = std.math;

const bytecode = @import("bytecode.zig");
const Opcode = bytecode.Opcode;
const value = @import("value.zig");
const Value = value.Value;
const Vm = @import("vm.zig");

fn recordBytecodeSize(fields: []const value.TagId) usize {
    var size: usize = 0;

    size += 1; // Opcode.field
    size += 1; // count: u8
    size += 2 * fields.len; // tags: [count]u16
    size += 1; // Opcode.return

    return size;
}

fn recordBytecode(comptime fields: []const value.TagId) [recordBytecodeSize(fields)]u8 {
    if (fields.len > 255) @compileError("too many fields");

    var code = [_]u8{undefined} ** recordBytecodeSize(fields);
    var cursor: usize = 0;

    code[cursor] = @intFromEnum(Opcode.field);
    cursor += 1;
    code[cursor] = @as(u8, @truncate(fields.len));
    cursor += 1;

    for (fields) |field| {
        const tag_id = mem.toBytes(field);
        code[cursor] = tag_id[0];
        code[cursor + 1] = tag_id[1];
        cursor += 2;
    }

    code[cursor] = @intFromEnum(Opcode.ret);
    cursor += 1;

    return code;
}

const record_dotter_bytecode = recordBytecode(&.{ .From, .To, .Num });
pub const record_dotter: bytecode.Chunk = .{ .bytecode = &record_dotter_bytecode };

pub const Context = struct {
    vm: *Vm,
    allocator: mem.Allocator,
    args: []Value,

    pub fn arg(cx: *const Context, i: usize) ?Value {
        if (i < cx.args.len) {
            return cx.args[i];
        } else {
            return null;
        }
    }
};

pub const Fn = *const fn (Context) Vm.Error!Value;
pub const FnTable = [256]Fn;

fn invalid(cx: Context) !Value {
    return cx.vm.throw("invalid system function", .{});
}

fn todo(cx: Context) !Value {
    return cx.vm.throw("not yet implemented", .{});
}

fn typeError(vm: *Vm, val: Value, i: usize, expect: []const u8) Vm.Error {
    return vm.throw("argument #{}: {s} expected, but got {s}", .{ i + 1, expect, val.typeName() });
}

// Strongly-typed type tags for Vec4 and Rgba, which are otherwise the same type.
// Used for differentiating between Vec4 and Rgba arguments.
const Vec4 = struct { value: value.Vec4 };
const Rgba = struct { value: value.Rgba };

fn fromArgument(cx: Context, comptime T: type, i: usize, val: Value) Vm.Error!T {
    // NOTE: THIS IS CURRENTLY TERRIBLY BROKEN BECAUSE i IS ITERATED EVEN WHEN IT DOESN'T NEED TO BE.

    switch (T) {
        // Context variables
        Context => return cx,
        *Vm => return cx.vm,
        mem.Allocator => return cx.allocator,

        // No conversion
        Value => return val,

        // Primitives
        f32 => {
            if (val != .number) return typeError(cx.vm, val, i, "number");
            return val.number;
        },
        Vec4 => {
            if (val != .vec4) return typeError(cx.vm, val, i, "vec4");
            return .{ .value = val.vec4 };
        },
        Rgba => {
            if (val != .rgba) return typeError(cx.vm, val, i, "rgba");
            return .{ .value = val.rgba };
        },

        // Refs
        value.List => {
            if (val != .ref or val.ref.* != .list) return typeError(cx.vm, val, i, "list");
            return val.ref.list;
        },
        *const value.Shape => {
            if (val != .ref or val.ref.* != .shape) return typeError(cx.vm, val, i, "shape");
            return &val.ref.shape;
        },
        *const value.Closure => {
            if (val != .ref or val.ref.* != .closure) return typeError(cx.vm, val, i, "function");
            return &val.ref.closure;
        },
        else => {},
    }
    @compileError(@typeName(T) ++ " is unsupported as an argument type");
}

fn intoReturn(cx: Context, any: anytype) Vm.Error!Value {
    const T = @TypeOf(any);
    return switch (T) {
        Value => any,
        bool => if (any) .true else .false,
        f32 => .{ .number = any },
        value.Ref => {
            const ref = cx.allocator.create(value.Ref) catch return cx.vm.outOfMemory();
            ref.* = any;
            return .{ .ref = ref };
        },
        else => switch (@typeInfo(T)) {
            .optional => if (any) |v| intoReturn(cx, v) else .nil,
            .error_union => intoReturn(cx, try any),
            else => @compileError(@typeName(T) ++ " is unsupported as a return type"),
        },
    };
}

/// Erase a well-typed function into a function that operates on raw values.
/// The erased function performs all the necessary conversions automatically.
///
/// Note that the argument order is important---function arguments go first, then context (such as
/// the VM or the allocator.) Otherwise argument indices will not match up.
fn erase(comptime func: anytype) Fn {
    return Erased(func).call;
}

fn Erased(comptime func: anytype) type {
    return struct {
        fn call(cx: Context) Vm.Error!Value {
            const param_count = @typeInfo(@TypeOf(func)).@"fn".params.len;
            if (cx.args.len != param_count) {
                return cx.vm.throw("function expects {} arguments, but it received {}", .{ param_count, cx.args.len });
            }

            const Args = meta.ArgsTuple(@TypeOf(func));
            var args: Args = undefined;
            inline for (meta.fields(Args), 0..) |field, i| {
                @field(args, field.name) = try fromArgument(cx, field.type, i, cx.args[i]);
            }

            const result = @call(.auto, func, args);
            return intoReturn(cx, result);
        }
    };
}

const SparseFn = struct { u8, Fn };
const SparseFnTable = []const SparseFn;

fn makeFnTable(init: SparseFnTable) FnTable {
    var table = [_]Fn{invalid} ** @typeInfo(FnTable).array.len;
    for (init) |entry| {
        const index, const func = entry;
        table[index] = func;
    }
    return table;
}

pub const fns = makeFnTable(&[_]SparseFn{
    // NOTE: The indices here must match those defined in system.rs.

    // Once the rest of the compiler is rewritten in Zig, it would be a good idea to rework this
    // system _not_ to hardcode the indices here.

    .{ 0x00, erase(add) }, // +
    .{ 0x01, erase(sub) }, // -
    .{ 0x02, erase(mul) }, // *
    .{ 0x03, erase(div) }, // /
    .{ 0x04, erase(neg) }, // -_
    .{ 0x10, erase(floor) }, // floor
    .{ 0x11, erase(ceil) }, // ceil
    .{ 0x12, erase(round) }, // round
    .{ 0x13, erase(abs) }, // abs
    .{ 0x14, erase(mod) }, // mod
    .{ 0x15, erase(pow) }, // pow
    .{ 0x16, erase(sqrt) }, // sqrt
    .{ 0x17, erase(cbrt) }, // cbrt
    .{ 0x18, erase(exp) }, // exp
    .{ 0x19, erase(exp2) }, // exp2
    .{ 0x1a, erase(ln) }, // ln
    .{ 0x1b, erase(log2) }, // log2
    .{ 0x1c, erase(log10) }, // log10
    .{ 0x1d, erase(hypot) }, // hypot
    .{ 0x1e, erase(sin) }, // sin
    .{ 0x1f, erase(cos) }, // cos
    .{ 0x20, erase(tan) }, // tan
    .{ 0x21, erase(asin) }, // asin
    .{ 0x22, erase(acos) }, // acos
    .{ 0x23, erase(atan) }, // atan
    .{ 0x24, erase(atan2) }, // atan2
    .{ 0x25, erase(expMinus1) }, // expMinus1
    .{ 0x26, erase(ln1Plus) }, // ln1Plus
    .{ 0x27, erase(sinh) }, // sinh
    .{ 0x28, erase(cosh) }, // cosh
    .{ 0x29, erase(tanh) }, // tanh
    .{ 0x2a, erase(asinh) }, // asinh
    .{ 0x2b, erase(acosh) }, // acosh
    .{ 0x2c, erase(atanh) }, // atanh
    .{ 0x2d, erase(min) }, // min
    .{ 0x2e, erase(max) }, // max
    .{ 0x30, erase(lerp) }, // lerp
    .{ 0x40, erase(not) }, // !
    .{ 0x41, erase(Value.eql) }, // ==
    .{ 0x42, erase(notEql) }, // !=
    .{ 0x43, erase(less) }, // <
    .{ 0x44, erase(lessOrEqual) }, // <=
    .{ 0x45, erase(greater) }, // >
    .{ 0x46, erase(greaterOrEqual) }, // >=
    .{ 0x80, vec }, // vec
    .{ 0x81, erase(vecX) }, // vecX
    .{ 0x82, erase(vecY) }, // vecY
    .{ 0x83, erase(vecZ) }, // vecZ
    .{ 0x84, erase(vecW) }, // vecW
    .{ 0x85, rgba }, // rgba
    .{ 0x86, erase(rgbaR) }, // rgbaR
    .{ 0x87, erase(rgbaG) }, // rgbaG
    .{ 0x88, erase(rgbaB) }, // rgbaB
    .{ 0x89, erase(rgbaA) }, // rgbaA
    .{ 0x90, erase(listLen) }, // len
    .{ 0x91, erase(listIndex) }, // index
    .{ 0x92, erase(range) }, // range
    .{ 0x93, erase(map) }, // map
    .{ 0x94, erase(filter) }, // filter
    .{ 0x95, erase(reduce) }, // reduce
    .{ 0x96, erase(flatten) }, // flatten
    .{ 0xc0, erase(valueToShape) }, // toShape
    .{ 0xc1, erase(line) }, // line
    .{ 0xc2, erase(rect) }, // rect
    .{ 0xc3, erase(circle) }, // circle
    .{ 0xe0, erase(stroke) }, // stroke
    .{ 0xe1, erase(fill) }, // fill
    .{ 0xf0, erase(withDotter) }, // withDotter
});

fn add(a: Value, b: Value, vm: *Vm) Vm.Error!Value {
    if (meta.activeTag(a) != meta.activeTag(b)) {
        return vm.throw("arguments must be of the same type", .{});
    }

    return switch (a) {
        .number => .{ .number = a.number + b.number },
        .vec4 => .{ .vec4 = a.vec4 + b.vec4 },
        .rgba => .{ .rgba = a.rgba + b.rgba },
        else => vm.throw("number, vec4, or rgba arguments expected, but got {s}", .{a.typeName()}),
    };
}

fn sub(a: Value, b: Value, vm: *Vm) Vm.Error!Value {
    if (meta.activeTag(a) != meta.activeTag(b)) {
        return vm.throw("arguments must be of the same type", .{});
    }

    return switch (a) {
        .number => .{ .number = a.number - b.number },
        .vec4 => .{ .vec4 = a.vec4 - b.vec4 },
        .rgba => .{ .rgba = a.rgba - b.rgba },
        else => vm.throw("number, vec4, or rgba arguments expected, but got {s}", .{a.typeName()}),
    };
}

fn mul(a: Value, b: Value, vm: *Vm) Vm.Error!Value {
    if (meta.activeTag(a) != meta.activeTag(b)) {
        return vm.throw("arguments must be of the same type", .{});
    }

    return switch (a) {
        .number => .{ .number = a.number * b.number },
        .vec4 => .{ .vec4 = a.vec4 * b.vec4 },
        .rgba => .{ .rgba = a.rgba * b.rgba },
        else => vm.throw("number, vec4, or rgba arguments expected, but got {s}", .{a.typeName()}),
    };
}

fn div(a: Value, b: Value, vm: *Vm) Vm.Error!Value {
    if (meta.activeTag(a) != meta.activeTag(b)) {
        return vm.throw("arguments must be of the same type", .{});
    }

    return switch (a) {
        .number => .{ .number = a.number * b.number },
        .vec4 => .{ .vec4 = a.vec4 * b.vec4 },
        .rgba => .{ .rgba = a.rgba * b.rgba },
        else => vm.throw("number, vec4, or rgba arguments expected, but got {s}", .{a.typeName()}),
    };
}

fn neg(a: Value, vm: *Vm) Vm.Error!Value {
    return switch (a) {
        .number => .{ .number = -a.number },
        .vec4 => .{ .vec4 = -a.vec4 },
        else => vm.throw("number or vec4 argument expected, but got {s}", .{a.typeName()}),
    };
}

fn floor(a: f32) f32 {
    return @floor(a);
}

fn ceil(a: f32) f32 {
    return @ceil(a);
}

fn round(a: f32) f32 {
    return @round(a);
}

fn abs(a: f32) f32 {
    return @abs(a);
}

fn mod(a: f32, b: f32) f32 {
    return @mod(a, b);
}

fn pow(a: f32, b: f32) f32 {
    return math.pow(f32, a, b);
}

fn sqrt(a: f32) f32 {
    return @sqrt(a);
}

fn cbrt(a: f32) f32 {
    return math.cbrt(a);
}

fn exp(a: f32) f32 {
    return @exp(a);
}

fn exp2(a: f32) f32 {
    return @exp2(a);
}

fn ln(a: f32) f32 {
    return @log(a);
}

fn log2(a: f32) f32 {
    return @log2(a);
}

fn log10(a: f32) f32 {
    return @log10(a);
}

fn hypot(a: f32, b: f32) f32 {
    return math.hypot(a, b);
}

fn sin(a: f32) f32 {
    return @sin(a);
}

fn cos(a: f32) f32 {
    return @cos(a);
}

fn tan(a: f32) f32 {
    return @tan(a);
}

fn asin(a: f32) f32 {
    return math.asin(a);
}

fn acos(a: f32) f32 {
    return math.acos(a);
}

fn atan(a: f32) f32 {
    return math.atan(a);
}

fn atan2(y: f32, x: f32) f32 {
    return math.atan2(y, x);
}

fn expMinus1(a: f32) f32 {
    return math.expm1(a);
}

fn ln1Plus(a: f32) f32 {
    return math.log1p(a);
}

fn sinh(a: f32) f32 {
    return math.sinh(a);
}

fn cosh(a: f32) f32 {
    return math.cosh(a);
}

fn tanh(a: f32) f32 {
    return math.tanh(a);
}

fn asinh(a: f32) f32 {
    return math.asinh(a);
}

fn acosh(a: f32) f32 {
    return math.acosh(a);
}

fn atanh(a: f32) f32 {
    return math.atanh(a);
}

fn min(a: Value, b: Value, vm: *Vm) Vm.Error!Value {
    if (meta.activeTag(a) != meta.activeTag(b)) {
        return vm.throw("arguments must be of the same type", .{});
    }

    return switch (a) {
        .number => .{ .number = @min(a.number, b.number) },
        .vec4 => .{ .vec4 = @min(a.vec4, b.vec4) },
        .rgba => .{ .rgba = @min(a.rgba, b.rgba) },
        else => if (a.lt(b).?) a else b,
    };
}

fn max(a: Value, b: Value, vm: *Vm) Vm.Error!Value {
    if (meta.activeTag(a) != meta.activeTag(b)) {
        return vm.throw("arguments must be of the same type", .{});
    }

    return switch (a) {
        .number => .{ .number = @max(a.number, b.number) },
        .vec4 => .{ .vec4 = @max(a.vec4, b.vec4) },
        .rgba => .{ .rgba = @max(a.rgba, b.rgba) },
        else => if (a.gt(b).?) a else b,
    };
}

fn lerp(a: Value, b: Value, t: f32, vm: *Vm) Vm.Error!Value {
    if (meta.activeTag(a) != meta.activeTag(b)) {
        return vm.throw("arguments must be of the same type", .{});
    }

    return switch (a) {
        .number => .{ .number = math.lerp(a.number, b.number, t) },
        .vec4 => .{ .vec4 = math.lerp(a.vec4, b.vec4, @as(value.Vec4, @splat(t))) },
        .rgba => .{ .rgba = math.lerp(a.rgba, b.rgba, @as(value.Rgba, @splat(t))) },
        else => vm.throw("number, vec4, or rgba expected, but got {s}", .{a.typeName()}),
    };
}

fn not(a: Value) bool {
    return !a.isTruthy();
}

// Value.eql is used as-is

fn notEql(a: Value, b: Value) bool {
    return !a.eql(b);
}

fn less(a: Value, b: Value, vm: *Vm) Vm.Error!bool {
    return a.lt(b) orelse vm.throw("{s} and {s} cannot be compared", .{ a.typeName(), b.typeName() });
}

fn greater(a: Value, b: Value, vm: *Vm) Vm.Error!bool {
    return a.gt(b) orelse vm.throw("{s} and {s} cannot be compared", .{ a.typeName(), b.typeName() });
}

fn lessOrEqual(a: Value, b: Value, vm: *Vm) Vm.Error!bool {
    const isGreater = try greater(a, b, vm);
    return !isGreater;
}

fn greaterOrEqual(a: Value, b: Value, vm: *Vm) Vm.Error!bool {
    const isLess = try less(a, b, vm);
    return !isLess;
}

fn vec(cx: Context) Vm.Error!Value {
    if (cx.args.len > 4) return cx.vm.throw("function expects 1 to 4 arguments, but it received {}", .{cx.args.len});
    for (cx.args) |arg| {
        if (arg != .number) return cx.vm.throw("number expected, but got {s}", .{arg.typeName()});
    }

    const zero: Value = .{ .number = 0 };
    return .{ .vec4 = .{
        (cx.arg(0) orelse zero).number,
        (cx.arg(1) orelse zero).number,
        (cx.arg(2) orelse zero).number,
        (cx.arg(3) orelse zero).number,
    } };
}

fn vecX(v: Vec4) f32 {
    return v.value[0];
}

fn vecY(v: Vec4) f32 {
    return v.value[1];
}

fn vecZ(v: Vec4) f32 {
    return v.value[2];
}

fn vecW(v: Vec4) f32 {
    return v.value[3];
}

fn rgba(cx: Context) Vm.Error!Value {
    if (cx.args.len > 4) return cx.vm.throw("function expects 1 to 4 arguments, but it received {}", .{cx.args.len});
    for (cx.args) |arg| {
        if (arg != .number) return cx.vm.throw("number expected, but got {s}", .{arg.typeName()});
    }

    const zero: Value = .{ .number = 0 };
    return .{ .rgba = .{
        (cx.arg(0) orelse zero).number,
        (cx.arg(1) orelse zero).number,
        (cx.arg(2) orelse zero).number,
        (cx.arg(3) orelse zero).number,
    } };
}

fn rgbaR(v: Rgba) f32 {
    return v.value[0];
}

fn rgbaG(v: Rgba) f32 {
    return v.value[1];
}

fn rgbaB(v: Rgba) f32 {
    return v.value[2];
}

fn rgbaA(v: Rgba) f32 {
    return v.value[3];
}

fn listLen(list: value.List) f32 {
    return @floatFromInt(list.len);
}

/// `index`
fn listIndex(list: value.List, index: f32, vm: *Vm) Vm.Error!Value {
    const i: usize = @intFromFloat(index);
    if (i >= list.len) return vm.throw("list index out of bounds. length is {}, index is {}", .{ list.len, i });
    return list[i];
}

fn range(fstart: f32, fend: f32, vm: *Vm, a: mem.Allocator) Vm.Error!value.Ref {
    const start: u32 = @intFromFloat(fstart);
    const end: u32 = @intFromFloat(fend);

    // Careful here. We don't want someone to generate a list that's so long it DoSes the server.
    // Therefore generating a list consumes fuel, in addition to bulk memory.
    // The cost is still much cheaper than doing it manually.
    const count = @max(start, end) - @min(start, end);
    try vm.consumeFuel(&vm.fuel, count);
    const list = a.alloc(Value, count) catch return vm.outOfMemory();

    if (start < end) {
        var i = start;
        for (list) |*element| {
            element.* = .{ .number = @floatFromInt(i) };
            i += 1;
        }
    } else {
        var i = end;
        for (list) |*element| {
            element.* = .{ .number = @floatFromInt(i) };
            i -= 1;
        }
    }

    return .{ .list = list };
}

fn map(list: value.List, f: *const value.Closure, vm: *Vm, a: mem.Allocator) Vm.Error!value.Ref {
    if (f.param_count != 1) {
        return vm.throw("function passed to map must have a single parameter (\\x -> x), but it has {}", .{f.param_count});
    }

    const mapped_list = a.dupe(Value, list) catch return vm.outOfMemory();
    for (list, mapped_list) |src, *dst| {
        const bottom = vm.stack_top;
        try vm.push(src);
        try vm.run(a, f, bottom);
        dst.* = try vm.pop();
    }

    return .{ .list = mapped_list };
}

fn filter(list: value.List, f: *const value.Closure, vm: *Vm, a: mem.Allocator) Vm.Error!value.Ref {
    if (f.param_count != 1) {
        return vm.throw("function passed to filter must have a single parameter (\\x -> True), but it has {}", .{f.param_count});
    }

    // Implementing filter is a bit tricky to do without resizable arrays.
    // There are a few paths one could take, but the simplest is to duplicate the list and truncate
    // its length. This wastes a lot of memory, but it probably isn't going to matter in practice.
    //
    // Serves you right for wanting to waste cycles on generating a list and then filtering it down
    // instead of just generating the list you want, I guess.
    //
    // In the future we could try allocating a bit map for the filter's results, but I'm not sure
    // it's worth it.
    const filtered_list = a.alloc(Value, list.len) catch return vm.outOfMemory();
    var len: usize = 0;
    for (list) |val| {
        const bottom = vm.stack_top;
        try vm.push(val);
        try vm.run(a, f, bottom);
        const condition = try vm.pop();
        if (condition.isTruthy()) {
            filtered_list[len] = val;
            len += 1;
        }
    }

    return .{ .list = filtered_list };
}

fn reduce(list: value.List, init: Value, f: *const value.Closure, vm: *Vm, a: mem.Allocator) Vm.Error!Value {
    if (f.param_count != 2) {
        return vm.throw("function passed to reduce must have two parameters (\\acc, x -> acc), but it has {}", .{f.param_count});
    }

    var accumulator = init;
    for (list) |val| {
        const bottom = vm.stack_top;
        try vm.push(accumulator);
        try vm.push(val);
        try vm.run(a, f, bottom);
        accumulator = try vm.pop();
    }

    return accumulator;
}

fn flatten(list: value.List, vm: *Vm, a: mem.Allocator) Vm.Error!value.Ref {
    var len: usize = 0;
    for (list) |val| {
        if (val == .ref and val.ref.* == .list) {
            len += val.ref.list.len;
        } else {
            len += 1;
        }
    }

    const flattened_list = a.alloc(Value, len) catch return vm.outOfMemory();
    var i: usize = 0;
    for (list) |val| {
        if (val == .ref and val.ref.* == .list) {
            @memcpy(flattened_list[i..][0..val.ref.list.len], val.ref.list);
            i += val.ref.list.len;
        } else {
            flattened_list[i] = val;
            i += 1;
        }
    }

    return .{ .list = flattened_list };
}

fn toShape(val: value.Value) ?value.Shape {
    return switch (val) {
        .nil, .false, .true, .tag, .number, .rgba => null,
        .vec4 => |v| .{ .point = value.vec2From4(v) },
        .ref => |r| if (r.* == .shape) r.shape else null,
    };
}

/// `toShape`
fn valueToShape(val: value.Value) ?value.Ref {
    if (toShape(val)) |shape| {
        return .{ .shape = shape };
    } else {
        return null;
    }
}

fn line(start: Vec4, end: Vec4) value.Ref {
    return .{ .shape = .{ .line = .{
        .start = value.vec2From4(start.value),
        .end = value.vec2From4(end.value),
    } } };
}

fn rect(top_left: Vec4, size: Vec4) value.Ref {
    return .{ .shape = .{ .rect = .{
        .top_left = value.vec2From4(top_left.value),
        .size = value.vec2From4(size.value),
    } } };
}

fn circle(center: Vec4, radius: f32) value.Ref {
    return .{ .shape = .{ .circle = .{
        .center = value.vec2From4(center.value),
        .radius = radius,
    } } };
}

fn stroke(thickness: f32, color: Rgba, shape: *const value.Shape) value.Ref {
    return .{ .scribble = .{
        .shape = shape.*,
        .action = .{ .stroke = .{
            .thickness = thickness,
            .color = color.value,
        } },
    } };
}

fn fill(color: Rgba, shape: *const value.Shape) value.Ref {
    return .{ .scribble = .{
        .shape = shape.*,
        .action = .{ .fill = .{
            .color = color.value,
        } },
    } };
}

fn withDotter(cont: *const value.Closure, vm: *Vm) Vm.Error!value.Ref {
    if (cont.param_count != 1) {
        return vm.throw("function passed to withDotter must have a single parameter (\\d -> _), but it has {}", .{cont.param_count});
    }
    return .{ .reticle = .{ .dotter = .{
        .draw = cont,
    } } };
}