const std = @import("std");
const warn = std.debug.warn;
const TypeId = @import("builtin").TypeId;

fn Iterator(comptime T: type) type {
    return struct {
        const Self = this;
        const ResetError = error {
            NotResetable,
        };

        const IterType = T;

        next_fn: fn(self: &Self) ?T,
        reset_fn: ?fn(self: &Self) ResetError!void,

        fn next(self: &Self) ?T {
            return self.next_fn(self);
        }

        fn reset(self: &Self) !void {
            const rfn = self.reset_fn ?? return error.NotResetable;
            return rfn(self);
        }
    };
}

fn UnwrapArg(comptime ArgType: type) type {
    if (@typeId(ArgType) == TypeId.Pointer) {
        return ArgType.Child;
    }

    return ArgType;
}

fn Filter(comptime filter_fn: var) type {
    const FnType = @typeOf(filter_fn);

    if (@typeId(FnType) != TypeId.Fn) {
        @compileError("Filter expects a function parameter.");
    }

    if (FnType.arg_count != 1) {
        @compileError("Filter expects a function with one argument.");
    }

    if (FnType.ReturnType != bool) {
        @compileError("Filter expects a function that returns bool.");
    }

    const T = UnwrapArg(@ArgType(FnType, 0));
    const It = Iterator(T);

    return struct {
        const Self = this;
        it: It,
        input_it: &It,

        fn init(input_it: &It) Self {
            return Self {
                .it = It {
                    .next_fn = next,
                    .reset_fn = reset,
                },
                .input_it = input_it,
            };
        }

        fn next(it: &It) ?T {
            const self = @fieldParentPtr(Self, "it", it);

            while (self.input_it.next()) |*val| {
                if (filter_fn(*val)) {
                    return *val;
                }
            }

            return null;
        }

        fn reset(it: &It) !void {
            const self = @fieldParentPtr(Self, "it", it);
            return self.input_it.reset();
        }
    };
}

fn Map(comptime map_fn: var) type {
    const FnType = @typeOf(map_fn);

    if (@typeId(FnType) != TypeId.Fn) {
        @compileError("Map expects a function parameter.");
    }

    if (FnType.arg_count != 1) {
        @compileError("Map expects a function with one argument.");
    }

    const T = UnwrapArg(@ArgType(FnType, 0));
    const V = FnType.ReturnType;

    const InputIt = Iterator(T);
    const OutputIt = Iterator(V);

    return struct {
        const Self = this;
        it: OutputIt,
        input_it: &InputIt,

        fn init(input_it: &InputIt) Self {
            return Self {
                .it = OutputIt {
                    .next_fn = next,
                    .reset_fn = reset,
                },
                .input_it = input_it,
            };
        }

        fn next(it: &OutputIt) ?V {
            const self = @fieldParentPtr(Self, "it", it);

            const val = self.input_it.next() ?? return null;
            return map_fn(val);
        }

        fn reset(it: &OutputIt) !void {
            const self = @fieldParentPtr(Self, "it", it);
            return self.input_it.reset();
        }
    };
}

fn Range(comptime start: usize, comptime end: usize) type {
    if (start > end) {
        @compileError("Invalid range.");
    }

    const It = Iterator(usize);

    return struct {
        const Self = this;
        it: It,
        current: usize,

        fn init() Self {
            return Self {
                .it = It {
                    .next_fn = next,
                    .reset_fn = reset,
                },
                .current = start,
            };
        }

        fn next(it: &It) ?usize {
            const self = @fieldParentPtr(Self, "it", it);

            if (self.current >= end) {
                return null;
            }

            var ret = self.current;
            self.current += 1;
            return ret;
        }

        fn reset(it: &It) It.ResetError!void {
            const self = @fieldParentPtr(Self, "it", it);

            self.current = start;
        }
    };
}

fn Repeat(comptime val: var, comptime times: i64) type {
    const T = @typeOf(val);
    const It = Iterator(T);

    return struct {
        const Self = this;
        it: It,
        current: if (times == -1) void else if (times > 0) i64 else unreachable,

        fn init() Self {
            return Self {
                .it = It {
                    .next_fn = next,
                    .reset_fn = if (times == -1) null else reset,
                },
                .current = if (times == -1) {} else 0,
            };
        }

        fn next(it: &It) ?T {
            if (times > 0) {
                const self = @fieldParentPtr(Self, "it", it);

                if (self.current >= times) {
                    return null;
                }

                self.current += 1;
            }

            return val;
        }

        fn reset(it: &It) !void {
            if (times > 0) {
                const self = @fieldParentPtr(Self, "it", it);
                self.current = 0;
            }

            return error.NotResetable;
        }
    };
}

fn addOne(x: usize) usize {
    return x + 1;
}

fn isEven(x: usize) bool {
    return x % 2 == 0;
}

fn consume(it: var) void {
    while (it.next()) |*val| {
        warn("Next value: {}\n", *val);
    }
}

const Foo = struct {
    data: usize,
};

const FooIterator = struct {
    const Self = this;
    const It = Iterator(Foo);

    it: It,
    current: usize,

    fn init() Self {
        return Self {
            .it = It {
                .next_fn = next,
                .reset_fn = null,
            },
            .current = 0,
        };
    }

    fn next(it: &It) ?Foo {
        const self = @fieldParentPtr(Self, "it", it);
        const current = self.current;
        self.current += 1;

        if (self.current > 30) {
            return null;
        }

        return Foo {
            .data = current,
        };
    }
};

fn filter_foo(foo: &const Foo) bool {
    return isEven(foo.data);
}

pub fn main() !void {
    var map_it = &Map(addOne).init(&Range(0, 10).init().it).it;

    consume(map_it);

    var filter_it = &Filter(isEven).init(map_it).it;

    // Reset iterator
    try filter_it.reset();

    consume(filter_it);

    var filter_foos = &Filter(filter_foo).init(&FooIterator.init().it).it;

    while (filter_foos.next()) |*val| {
        warn("Foo data: {}\n", val.data);
    }
}