from dataclasses import dataclass
from functools import singledispatch
from typing import Dict, List

from .ast import ImplementationFile, FunctionDefinition, ExpressionStatement, FunctionCallExpression, \
    VariableExpression, ConstantExpression, ReturnStatement, BasicType


@dataclass
class SsaResult:
    data: List[str]
    code: List[str]


@dataclass
class CompileContext:
    next_data: int = 0
    next_temp: int = 0


def build_ssa(file: ImplementationFile) -> str:
    result = compile_to_ssa(file, CompileContext())
    data = '\n'.join(result.data)
    code = '\n'.join(result.code)
    return data + '\n\n' + code


@singledispatch
def compile_to_ssa(target, context: CompileContext) -> SsaResult:
    raise NotImplementedError('unannotated compile on ' + str(type(target)))


@compile_to_ssa.register
def _(target: ImplementationFile, context: CompileContext):
    data = []
    code = []
    for target in target.contents:
        result = compile_to_ssa(target, context)
        data += result.data
        code += result.code
    return SsaResult(data, code)


@compile_to_ssa.register
def _(target: FunctionDefinition, context: CompileContext) -> SsaResult:
    data = []
    code = []
    for statement in target.body:
        result = compile_to_ssa(statement, context)
        data += result.data
        code += result.code
    code = ['    ' + instr for instr in code]
    assert len(target.args) == 0
    assert target.return_type == BasicType('int32')
    code = [f"export function w ${target.name}() {{", "@start", *code, "}"]
    return SsaResult(data, code)


@compile_to_ssa.register
def _(target: ExpressionStatement, context: CompileContext) -> SsaResult:
    return compile_to_ssa(target.body, context)


@compile_to_ssa.register
def _(target: FunctionCallExpression, context: CompileContext) -> SsaResult:
    assert isinstance(target.function, VariableExpression)
    data = []
    code = []
    args = []
    for i, expr in enumerate(target.arguments):
        arg_dest = context.next_temp
        result = compile_to_ssa(expr, context)
        data += result.data
        code += result.code
        args += [f"l %t{arg_dest}"]
    code += [f"call ${target.function.name}({','.join(args)}, ...)"]
    return SsaResult(data, code)


@compile_to_ssa.register
def _(target: ConstantExpression, context: CompileContext) -> SsaResult:
    if target.value.startswith('"'):
        data_dest = context.next_data
        context.next_data += 1
        data = [f"data $data{data_dest} = {{ b {target.value}, b 0 }}"]
        temp = context.next_temp
        context.next_temp += 1
        code = [f"%t{temp} =l copy $data{data_dest}"]
    else:
        assert not target.value.startswith('0b')
        assert not target.value.startswith('0B')
        assert not target.value.startswith('0o')
        assert not target.value.startswith('0x')
        assert not target.value.startswith('0X')
        assert not target.value.startswith('0f')
        assert not target.value.startswith('0F')
        assert '.' not in target.value
        assert not target.value.startswith("'")
        data = []
        temp = context.next_temp
        context.next_temp += 1
        code = [f"%t{temp} =w copy {target.value}"]
    return SsaResult(data, code)


@compile_to_ssa.register
def _(target: ReturnStatement, context: CompileContext) -> SsaResult:
    if target.body is None:
        return SsaResult([], ['ret'])
    ret_val_dest = context.next_temp
    result = compile_to_ssa(target.body, context)
    result.code.append(f"ret %t{ret_val_dest}")
    return result