# Copyright 2021-2022, Barcelona Supercomputing Center (BSC), Spain
#
# This software was partially supported by the EuroHPC-funded project ADMIRE
#   (Project ID: 956748, https://www.admire-eurohpc.eu).
#
# This file is part of rpcc.
#
# rpcc is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
#
# rpcc is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
# of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License along with rpcc.  If not, see
# <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: GPL-3.0-or-later

import pathlib
from typing import List, Tuple, Optional, Dict

import lark
from lark import v_args, Discard
from loguru import logger

import rpcc.parser
from rpcc.exceptions import RpcRedefinition
from rpcc.transformers.cxx17.formatter import FileFormatter
from rpcc.meta import FilePosition
from rpcc.parser import Parser
from rpcc.ir import RemoteProcedure, Argument, ReturnValue, IRTree, ArgumentList, ReturnValueList, ConditionalDefinition
from rpcc.transformers import FileTransformer
from .types import Bool, Double, Float, Integer, String, ExposedBuffer, CSize


class _RpcInfo:
    pos_in_text: FilePosition

    def __init__(self, pos_in_text: FilePosition):
        self.pos_in_text = pos_in_text


class _CxxCompilerDefines(ConditionalDefinition):
    _symbol: str

    def __init__(self, symbol):
        self._symbol = symbol

    @property
    def start_keyword(self):
        return "#ifdef"

    @property
    def end_keyword(self):
        return "#endif"

    @property
    def symbol(self):
        return self._symbol

    @symbol.setter
    def symbol(self, value):
        self._symbol = value


class _CxxCompilerNotDefines(ConditionalDefinition):
    _symbol: str

    def __init__(self, symbol):
        self._symbol = symbol

    @property
    def start_keyword(self):
        return "#ifndef"

    @property
    def end_keyword(self):
        return "#endif"

    @property
    def symbol(self):
        return self._symbol

    @symbol.setter
    def symbol(self, value):
        self._symbol = value


class _AstTransformer(lark.Transformer):
    input_file: pathlib.Path
    text: str
    rpc_info: Dict[str, _RpcInfo]

    def __init__(self, input_file: pathlib.Path, text: str):
        super().__init__()
        self.input_file = input_file
        self.text = text
        self.rpc_info = dict()

    @v_args(inline=True)
    def start(self, package: Optional[str], rpcs: List[RemoteProcedure]) -> IRTree:
        """Transform the AST tree generated by Lark into an `IR` instance, that we can use to generate the output files
        """

        return IRTree(package, rpcs)

    @v_args(inline=True)
    def package(self, name: str) -> str:
        return name

    # rpc_list can be converted directly into list
    rpc_list = list

    @v_args(inline=True, meta=True)
    def rpc(self, meta: lark.tree.Meta, name: str, id: Optional[int], include_if_expr: Optional[ConditionalDefinition],
            args: Optional[ArgumentList], retvals: Optional[ReturnValueList]) -> RemoteProcedure:
        """Transform a `rpc` node from the AST into a `RemoteProcedure` object from a name, a list of `Argument`s and a
        `ReturnVariable`.

        :param meta: Metainformation about the current Token such as line, column, start_pos, etc.
        :param name: A name for the remote procedure.
        :param id: An optional numeric id the remote procedure.
        :param args: An `ArgumentList` describing the input arguments for the remote procedure. `ArgumentList` may be
        empty if the input file defined an empty `arguments` clause, or `None` if not defined.
        :param retvals: A `ReturnValueList` describing the remote procedure's return values. `ReturnValueList` may be
        empty if the input file defined an empty `returns` clause, or `None` if not defined.
        :param include_if_expr: A C/C++ preprocessor expression specifying when this rpc should be included in
        compilation.
        :return: A `RemoteProcedure` object describing the remote procedure.
        """

        if not id:
            id = RemoteProcedure.generate_id(margo_compatible=False)

        if not args:
            args = ArgumentList([])

        if not retvals:
            retvals = ReturnValueList([])

        pos = FilePosition(self.input_file, self.text, meta.line, meta.column, meta.start_pos)
        if name in self.rpc_info:
            prev_pos = self.rpc_info[name].pos_in_text
            raise RpcRedefinition(pos, name, prev_pos)
        else:
            self.rpc_info[name] = _RpcInfo(pos)

        return RemoteProcedure(id, name, args, retvals, include_if_expr)

    @v_args(inline=True)
    def rpc_name(self, name: str) -> str:
        """Transform a `rpc_name` node from the AST into its corresponding string. The string is extracted from the AST
        nodes' value.

        :param name: A string value from the `rpc_name` node.
        :return: A name for the remote procedure.
        """
        return name

    @v_args(inline=True)
    def rpc_id(self, id: int) -> int:
        """Transform a `rpc_id` node from the AST into its corresponding id. The id is extracted from the AST
        nodes' value.

        :param id: An integer value from the `rpc_id` node.
        :return: An id for the remote procedure.
        """
        return id

    @v_args(inline=True)
    def include_if(self, expr) -> ConditionalDefinition:
        return expr

    @v_args(inline=True)
    def ifdef(self, symbol) -> ConditionalDefinition:
        return _CxxCompilerDefines(symbol)

    @v_args(inline=True)
    def ifndef(self, symbol) -> ConditionalDefinition:
        return _CxxCompilerNotDefines(symbol)

    def rpc_args(self, children: List[Tuple[str, str]]) -> ArgumentList:
        """Transform a list of tuples returned by Lark into a list of `rpc.Arguments`. The tuples are implicitly
        constructed by Lark by transforming all the `var` nodes in the AST that have this particular `rpc_args` node
        as a parent.

        :param children: A list of tuples describing the rpc arguments.
        :return: A list of `rpc.Arguments`
        """
        return ArgumentList([Argument(id, typeinfo) for (typeinfo, id) in children])

    def rpc_retvals(self, children: List[Tuple[str, str]]) -> ReturnValueList:
        """Transform a tuple returned by Lark into a `ReturnVariable`. The tuple is implictly constructed by Lark by
        transforming the `var` node under this particular `rpc_return` node.

        :param children: A list of tuples describing the rpc return values.
        :return: A list of `rpc.ReturnVariables`
        """
        return ReturnValueList([ReturnValue(id, typeinfo) for (typeinfo, id) in children])

    # variables can be converted directly into tuples
    var = tuple

    def bool(self, _) -> Bool:
        return Bool()

    def double(self, _) -> Double:
        return Double()

    def float(self, _) -> Float:
        return Float()

    def int8(self, _) -> Integer:
        return Integer(8, False)

    def int16(self, _) -> Integer:
        return Integer(16, False)

    def int32(self, _) -> Integer:
        return Integer(32, False)

    def int64(self, _) -> Integer:
        return Integer(64, False)

    def uint8(self, _) -> Integer:
        return Integer(8, True)

    def uint16(self, _) -> Integer:
        return Integer(16, True)

    def uint32(self, _) -> Integer:
        return Integer(32, True)

    def uint64(self, _) -> Integer:
        return Integer(64, True)

    def csize(self, _) -> CSize:
        return CSize()

    def string(self, _) -> String:
        return String()

    def exposed_buffer(self, _) -> ExposedBuffer:
        return ExposedBuffer()

    # variable names can be coverted directly into strings
    NAME = str
    DOTTED_NAME = str

    def EOS(self, children):
        return Discard


class Cxx17FileTransformer(FileTransformer):
    input_file: pathlib.Path
    copyright_file: pathlib.Path
    header_outdir: pathlib.Path
    impl_outdir: pathlib.Path
    parser: rpcc.parser.Parser
    ast_transformer: _AstTransformer

    def _get_output_paths(self) -> Tuple[pathlib.Path, pathlib.Path]:
        prefix = self.input_file.stem.replace('.', '_')
        return (pathlib.Path(self.header_outdir, f"{prefix}-rpcc.hpp"),
                pathlib.Path(self.impl_outdir, f"{prefix}-rpcc.cpp"))

    def __init__(self, input_file: pathlib.Path, copyright_file: pathlib.Path, header_outdir: pathlib.Path,
                 impl_outdir: pathlib.Path):
        self.input_file = input_file
        self.copyright_file = copyright_file
        self.header_outdir = header_outdir
        self.impl_outdir = impl_outdir
        self.parser = Parser()

    def transform(self) -> None:

        ast, text = self.parser.parse(self.input_file)

        try:
            ir = _AstTransformer(self.input_file, text).transform(ast)
        except lark.exceptions.VisitError as e:
            raise e.orig_exc

        header_file, impl_file = self._get_output_paths()
        header_text, impl_text = FileFormatter(self.input_file, self.copyright_file, header_file, impl_file).format(ir)

        try:
            with open(header_file, "w") as f:
                f.write(header_text)
        except OSError as e:
            raise OSError(f"Error writing header file:\n  {e}") from None

        try:
            with open(impl_file, "w") as f:
                f.write(impl_text)
        except OSError as e:
            raise OSError(f"Error writing implementation file:\n  {e}") from None
