from __future__ import annotations

import functools
from abc import abstractmethod, ABC
from typing import Any, Iterable

from . import ir, semantic, lexer
from .errors import SemanticAnalysisError, OverrideMandatoryError
from .logger import Logger
from .source import SourceLocation
from .typechecking import typecheck

logger = Logger(__name__)


class Node:

    def __init__(self):
        self.context: semantic.Context | None = None
        self._repr_guard: bool = False

    @abstractmethod
    def _values(self) -> list[Node | Any]:
        raise OverrideMandatoryError()

    @functools.cache
    def location(self) -> SourceLocation:
        locations = [v.location() for v in self._values()]
        begin = min([loc.begin for loc in locations])
        end = max([loc.end for loc in locations])
        loc = SourceLocation(begin=begin, end=end)
        return loc

    def __repr__(self):
        if self._repr_guard:
            return self.__class__.__name__
        self._repr_guard = True
        vals = self._values()
        if type(vals) == list:
            vals = ", ".join(repr(val) for val in vals)
        else:
            vals = repr(vals)
        self._repr_guard = False
        return f"{self.__class__.__name__}({vals})"

    def pprint(self, depth: int | None = None, indent: str = "\t"):
        print("\n".join(self._pprint(depth=depth, indent=indent)))

    def _pprint(self, depth: int | None, indent: str, _depth: int = 0) -> list[str]:
        if depth is not None and _depth >= depth:
            return [f"{indent}{self.__class__.__name__} {{ ... }}"]

        vals = self._values()
        try:
            vals = (val for val in vals)
        except TypeError:
            vals = (vals,)

        result = [f"{self.__class__.__name__} {{"]
        for val in vals:
            if isinstance(val, Node):
                result += val._pprint(depth=depth, indent=indent, _depth=_depth + 1)
            else:
                result += [f"{indent}{repr(val)}"]
        result += [f"}} // {self.__class__.__name__}"]
        for i, line in enumerate(result):
            result[i] = indent + line

        return result

    @abstractmethod
    def intermediate_representation(self) -> list[ir.IRItem]:
        raise OverrideMandatoryError()

    def semantic_analysis(self, context: semantic.Context):
        logger.debug(f"Doing semantic analysis in {self}")
        for value in self._values():
            if isinstance(value, Node):
                value.semantic_analysis(context)
        self.context = context

    @staticmethod
    def _prepare_sources_ir(result: list[ir.IRAction], values: Iterable[Value]) -> list[ir.IRValue]:
        vals = [value.intermediate_representation() for value in values]
        for value in vals:
            result += value

        return [value[-1].destination() for value in vals]


class Literal(Node, ABC):
    def __init__(self, location: SourceLocation, value: Any):
        super().__init__()
        self.value = value
        self.loc = location

    def location(self) -> SourceLocation:
        return self.loc

    def _values(self) -> list[Node | Any]:
        return [self.value]

    def _pprint(self, depth: int | None, indent: str, _depth: int = 0) -> list[str]:
        return [f"{indent}{repr(self)}"]

    def intermediate_representation(self) -> list[ir.IRItem]:
        dest = ir.IRRegister(location=self.location())
        immediate = ir.IRImmediate(location=self.location(), value=self.value)
        result = [ir.IRMove(location=self.location(), dest=dest, source=immediate)]
        return result


class PseudoNode(Literal):
    """
    Only used for better diagnostics
    """

    def __init__(self, token: lexer.Token):
        super().__init__(token.loc, token.value)
        self.token = token

    def intermediate_representation(self) -> list[ir.IRItem]:
        return []

    def semantic_analysis(self, context: semantic.Context):
        pass


class Sum(Node):
    def intermediate_representation(self) -> list[ir.IRItem]:
        result: list[ir.IRAction] = []

        values_results = Node._prepare_sources_ir(result=result, values=self.values)

        dest = ir.IRRegister(location=self.location())
        result += [ir.IRAdd(self.location(), dest, *values_results)]
        return result

    def __init__(self, *values: Value):
        super().__init__()
        self.values = values

    def _values(self) -> list[Node | Any]:
        return list(self.values)


class Sub(Node):
    def intermediate_representation(self) -> list[ir.IRItem]:
        result: list[ir.IRAction] = []

        first_val = self.first_value.intermediate_representation()
        result += first_val

        values_results = Node._prepare_sources_ir(result=result, values=self.values)

        for i, value_result in enumerate(values_results):
            d = ir.IRRegister(location=self.location())
            result += [ir.IRNegation(location=self.location(), dest=d, source=value_result)]
            values_results[i] = result[-1].destination()

        values_results += [first_val[-1].destination()]

        dest = ir.IRRegister(location=self.location())
        result += [ir.IRAdd(self.location(), dest, *values_results)]
        return result

    def __init__(self, first_value: Value, *values: Value):
        super().__init__()
        self.first_value = first_value
        self.values = values

    def _values(self) -> list[Node | Any]:
        return [self.first_value] + list(self.values)


class Product(Node):
    def intermediate_representation(self) -> list[ir.IRItem]:
        result: list[ir.IRAction] = []

        values_results = Node._prepare_sources_ir(result=result, values=self.values)
        dest = ir.IRRegister(location=self.location())
        result += [ir.IRMul(self.location(), dest, *values_results)]
        return result

    def __init__(self, *values: Value):
        super().__init__()
        self.values = values

    def _values(self) -> list[Node | Any]:
        return list(self.values)


class Division(Node):
    def intermediate_representation(self) -> list[ir.IRItem]:
        result: list[ir.IRAction] = []

        first_val = self.first_value.intermediate_representation()
        result += first_val

        values_results = Node._prepare_sources_ir(result=result, values=self.values)

        for i, value_result in enumerate(values_results):
            d = ir.IRRegister(location=self.location())
            result += [ir.IRInvert(location=self.location(), dest=d, source=value_result)]
            values_results[i] = result[-1].destination()

        values_results += [first_val[-1].destination()]

        dest = ir.IRRegister(location=self.location())
        result += [ir.IRMul(self.location(), dest, *values_results)]
        return result

    def __init__(self, first_value: Value, *values: Value):
        super().__init__()
        self.first_value = first_value
        self.values = values

    def _values(self) -> list[Node | Any]:
        return [self.first_value] + list(self.values)


BinaryOperation = Sum | Sub | Product | Division


@typecheck
class Float(Literal):
    def __init__(self, location: SourceLocation, value: float):
        super().__init__(location, value)


@typecheck
class Integer(Literal):

    def __init__(self, location: SourceLocation, value: int):
        super().__init__(location, value)


class Expression(Node):
    def intermediate_representation(self) -> list[ir.IRItem]:
        return self.node.intermediate_representation()

    def __init__(self, node: Node):
        super().__init__()
        self.node = node

    def _values(self) -> list[Node | Any]:
        return [self.node]

    def semantic_analysis(self, context: semantic.Context):
        return self.node.semantic_analysis(context)

    def location(self) -> SourceLocation:
        return self.node.location()


class Statement(Node):

    def __init__(self, *nodes: Node):
        super().__init__()
        self.nodes = list(nodes)

    def _values(self) -> list[Node | Any]:
        return self.nodes

    def intermediate_representation(self) -> list[ir.IRItem]:
        result: list[ir.IRItem] = []
        for node in self.nodes:
            result += node.intermediate_representation()
        return result


class Block(Statement):
    def __init__(self, name: str, *nodes: Node):
        super().__init__(*nodes)
        self.name = name

    def semantic_analysis(self, context: semantic.Context | None):
        child_context = semantic.Context(name=self.name, parent=context)
        if context is not None:
            context.add_context(child_context)
        logger.debug(f"Created new context: {child_context}")
        super().semantic_analysis(child_context)


class Identifier(Literal):
    def __init__(self, location: SourceLocation, name: str):
        super().__init__(location, name)
        self.value: str


class Variable(Node):
    def __init__(self, identifier: Identifier):
        super().__init__()
        self.value: semantic.Variable | None = None
        self.identifier = identifier

    def _values(self) -> list[Node | Any]:
        return [self.identifier]

    def semantic_analysis(self, context: semantic.Context):
        variable = context.get_variable(self.identifier.value)
        if variable is None:
            raise SemanticAnalysisError(location=self.location(), message=f"Unknown variable '{self.identifier.value}'")

        self.value = variable
        logger.debug(f"Linked variable reference to var {variable}")

    def intermediate_representation(self) -> list[ir.IRItem]:
        assert self.value is not None

        dest = ir.IRRegister(location=self.location())
        immediate = ir.IRVariable(location=self.location(), fq_identifier=self.value.fully_qualified_name())
        result = [ir.IRMove(location=self.location(), dest=dest, source=immediate)]
        return result


class Assignment(Node):
    def intermediate_representation(self) -> list[ir.IRItem]:
        assert self.variable is not None
        assert self.value is not None

        result: list[ir.IRItem] = []
        value = self.value.intermediate_representation()
        result += value

        dest = ir.IRVariable(location=self.location(), fq_identifier=self.variable.fully_qualified_name())
        result += [ir.IRMove(location=self.location(), dest=dest.destination(), source=value[-1].destination())]

        return result

    def __init__(self, identifier: Identifier, value: Value | None):
        super().__init__()
        self.identifier = identifier
        self.value = value
        self.variable: semantic.Variable | None = None

    def _values(self) -> list[Node | Any]:
        return [self.identifier, self.value]

    def semantic_analysis(self, context: semantic.Context):
        super(Assignment, self).semantic_analysis(context)
        name = self.identifier.value
        variable = context.set_variable(name, value=self.value)
        self.variable = variable
        logger.debug(f"Assigned variable {variable.fully_qualified_name()}")


class Definition(Assignment):

    def __init__(self, identifier: Identifier, type_identifier: Identifier, value: Value | None,
                 *pseudo_nodes: PseudoNode):
        super().__init__(identifier, value)
        self.type_identifier = type_identifier
        self.pseudo_nodes = list(pseudo_nodes)

    def intermediate_representation(self) -> list[ir.IRItem]:
        if self.value is None:
            return [ir.IRVoid(self.location())]

        return super(Definition, self).intermediate_representation()

    def _values(self) -> list[Node | Any]:
        v = [self.identifier, self.type_identifier]
        if self.value is not None:
            v += [self.value]
        v += self.pseudo_nodes
        return v

    @property
    def pure(self) -> bool:
        return self.value is None


Number = Float | Integer
Value = BinaryOperation | Number | Variable | Expression