"""
This module contains the code to convert parsed Ourlang into WebAssembly code
"""
from typing import Generator, Tuple

from . import ourlang
from . import wasm

Statements = Generator[wasm.Statement, None, None]

def type_(inp: ourlang.OurType) -> wasm.OurType:
    if isinstance(inp, ourlang.OurTypeInt32):
        return wasm.OurTypeInt32()

    if isinstance(inp, ourlang.OurTypeInt64):
        return wasm.OurTypeInt64()

    if isinstance(inp, ourlang.OurTypeFloat32):
        return wasm.OurTypeFloat32()

    if isinstance(inp, ourlang.OurTypeFloat64):
        return wasm.OurTypeFloat64()

    raise NotImplementedError(type_, inp)

OPERATOR_MAP = {
    '+': 'add',
    '-': 'sub',
}

I32_OPERATOR_MAP = { # TODO: Introduce UInt32 type
    '<': 'lt_s',
    '>': 'gt_s',
    '<=': 'le_s',
    '>=': 'ge_s',
}

def expression(inp: ourlang.Expression) -> Statements:
    if isinstance(inp, ourlang.ConstantInt32):
        yield wasm.Statement('i32.const', str(inp.value))
        return

    if isinstance(inp, ourlang.ConstantInt64):
        yield wasm.Statement('i64.const', str(inp.value))
        return

    if isinstance(inp, ourlang.ConstantFloat32):
        yield wasm.Statement('f32.const', str(inp.value))
        return

    if isinstance(inp, ourlang.ConstantFloat64):
        yield wasm.Statement('f64.const', str(inp.value))
        return

    if isinstance(inp, ourlang.VariableReference):
        yield wasm.Statement('local.get', '${}'.format(inp.name))
        return

    if isinstance(inp, ourlang.BinaryOp):
        yield from expression(inp.left)
        yield from expression(inp.right)

        if isinstance(inp.type, ourlang.OurTypeInt32):
            if operator := OPERATOR_MAP.get(inp.operator, None):
                yield wasm.Statement(f'i32.{operator}')
                return
            if operator := I32_OPERATOR_MAP.get(inp.operator, None):
                yield wasm.Statement(f'i32.{operator}')
                return
        if isinstance(inp.type, ourlang.OurTypeInt64):
            if operator := OPERATOR_MAP.get(inp.operator, None):
                yield wasm.Statement(f'i64.{operator}')
                return
        if isinstance(inp.type, ourlang.OurTypeFloat32):
            if operator := OPERATOR_MAP.get(inp.operator, None):
                yield wasm.Statement(f'f32.{operator}')
                return
        if isinstance(inp.type, ourlang.OurTypeFloat64):
            if operator := OPERATOR_MAP.get(inp.operator, None):
                yield wasm.Statement(f'f64.{operator}')
                return

        raise NotImplementedError(expression, inp.type, inp.operator)

    raise NotImplementedError(expression, inp)

def statement_return(inp: ourlang.StatementReturn) -> Statements:
    yield from expression(inp.value)
    yield wasm.Statement('return')

def statement_if(inp: ourlang.StatementIf) -> Statements:
    yield from expression(inp.test)

    yield wasm.Statement('if')

    for stat in inp.statements:
        yield from statement(stat)

    if inp.else_statements:
        yield wasm.Statement('else')
        for stat in inp.else_statements:
            yield from statement(stat)

    yield wasm.Statement('end')

def statement(inp: ourlang.Statement) -> Statements:
    if isinstance(inp, ourlang.StatementReturn):
        yield from statement_return(inp)
        return

    if isinstance(inp, ourlang.StatementIf):
        yield from statement_if(inp)
        return

    raise NotImplementedError(statement, inp)

def function_argument(inp: Tuple[str, ourlang.OurType]) -> wasm.Param:
    return (inp[0], type_(inp[1]), )

def function(inp: ourlang.Function) -> wasm.Function:
    return wasm.Function(
        inp.name,
        inp.exported,
        [
            function_argument(x)
            for x in inp.posonlyargs
        ],
        [], # TODO
        type_(inp.returns),
        [
            x
            for y in inp.statements
            for x in statement(y)
        ]
    )

def module(inp: ourlang.Module) -> wasm.Module:
    result = wasm.Module()

    result.functions = [*map(
        function, inp.functions.values(),
    )]

    return result