#!/usr/bin/env python
# -*- coding: utf-8 -*-

from itertools import product, chain, zip_longest, tee


IF, ELSE_IF, ELSE, END_IF = 0, 1, 2, 3

NL = "\\n"
dash = lambda ws: [" ", "-"][ws]


def cond_kind(i, n):
    i += 1
    if i == 1:
        return IF # if
    elif (i == n) and (i > 1):
        return ELSE # else
    elif i > n:
        return END_IF # endif
    else:
        return ELSE_IF # else if


# From: https://docs.python.org/3/library/itertools.html#itertools-recipes
def pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


def write_cond(conds, active_branch):
    n = len(conds) - 1

    lits = []
    for i in range(1, n + 2 + 1):
        ws1 = "\\n" * i
        ws2 = "\\r\\n" * i
        lits.append((ws1, str(i), ws2))

    conds = list(conds)
    for i, (pws, nws) in enumerate(conds):
        kind = cond_kind(i, n)
        b = str(i == active_branch).lower()
        cond = [f"if {b}", f"else if {b}", "else", "endif"][kind]
        cond = f"{{%{dash(pws)} {cond} {dash(nws)}%}}"
        conds[i] = cond

    it = map("".join, lits)
    it = filter(None, chain.from_iterable(zip_longest(it, conds)))
    code = "".join(it)

    expected = f"{lits[0][0]}{lits[0][1]}"
    for i, (cond, (before, after)) in enumerate(zip(conds, pairwise(lits))):
        kind = cond_kind(i, n)
        pws = cond.startswith("{%-")
        nws = cond.endswith("-%}")

        cond = i == active_branch
        prev_cond = i == (active_branch + 1)

        if prev_cond or (kind == IF):
            expected += before[2] * (not pws)
        if cond or (kind == END_IF):
            expected += after[0] * (not nws)
            expected += after[1]

    # FIXME: Askama does not include whitespace before eof
    # expected += lits[-1][2]

    return code, expected


if __name__ == "__main__":
    # The amount of branches to generate
    n = 2 # branches


    with open("ws.rs", "w") as f:
        f.write("""\
// This file is auto generated by gen_ws_tests.py

use askama::Template;

macro_rules! test_template {
    ($source:literal, $rendered:expr) => {{
        #[derive(Template)]
        #[template(source = $source, ext = "txt")]
        struct CondWS;

        assert_eq!(CondWS.render().unwrap(), $rendered);
    }};
}

#[rustfmt::skip]
#[test]
fn test_cond_ws() {
""")

        for branches in range(1, n + 1):
            for x in product([False, True], repeat=(branches+1)*2):
                # it = iter(x)
                # conds = list(zip(it, it))
                conds = list(zip(x[::2], x[1::2]))

                for i in range(branches):
                    code, expected = write_cond(conds, i)
                    f.write(f'    test_template!("{code}", "{expected}");\n')

            if branches != n:
                f.write("\n")
        f.write("}\n")