# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


from textwrap import dedent
from typing import Callable

import libcst as cst
from libcst.metadata import ExperimentalReentrantCodegenProvider, MetadataWrapper
from libcst.testing.utils import UnitTest, data_provider


class ExperimentalReentrantCodegenProviderTest(UnitTest):
    @data_provider(
        {
            "simple_top_level_statement": {
                "old_module": (
                    """\
                    import math
                    c = math.sqrt(a*a + b*b)
                    """
                ),
                "new_module": (
                    """\
                    import math
                    c = math.hypot(a, b)
                    """
                ),
                "old_node": lambda m: m.body[1],
                "new_node": cst.parse_statement("c = math.hypot(a, b)"),
            },
            "replacement_inside_block": {
                "old_module": (
                    """\
                    import math
                    def do_math(a, b):
                        c = math.sqrt(a*a + b*b)
                        return c
                    """
                ),
                "new_module": (
                    """\
                    import math
                    def do_math(a, b):
                        c = math.hypot(a, b)
                        return c
                    """
                ),
                "old_node": lambda m: m.body[1].body.body[0],
                "new_node": cst.parse_statement("c = math.hypot(a, b)"),
            },
            "missing_trailing_newline": {
                "old_module": "old_fn()",  # this module has no trailing newline
                "new_module": "new_fn()",
                "old_node": lambda m: m.body[0],
                "new_node": cst.parse_statement("new_fn()\n"),
            },
            "nested_blocks_with_missing_trailing_newline": {
                "old_module": (
                    """\
                    if outer:
                        if inner:
                            old_fn()"""  # this module has no trailing newline
                ),
                "new_module": (
                    """\
                    if outer:
                        if inner:
                            new_fn()"""
                ),
                "old_node": lambda m: m.body[0].body.body[0].body.body[0],
                "new_node": cst.parse_statement("new_fn()\n"),
            },
        }
    )
    def test_provider(
        self,
        old_module: str,
        new_module: str,
        old_node: Callable[[cst.Module], cst.CSTNode],
        new_node: cst.BaseStatement,
    ) -> None:
        old_module = dedent(old_module)
        new_module = dedent(new_module)

        mw = MetadataWrapper(cst.parse_module(old_module))
        codegen_partial = mw.resolve(ExperimentalReentrantCodegenProvider)[
            old_node(mw.module)
        ]

        self.assertEqual(codegen_partial.get_original_module_code(), old_module)
        self.assertEqual(codegen_partial.get_modified_module_code(new_node), new_module)

    def test_byte_conversion(
        self,
    ) -> None:
        module_bytes = "fn()\n".encode("utf-16")
        mw = MetadataWrapper(
            cst.parse_module("fn()\n", cst.PartialParserConfig(encoding="utf-16"))
        )
        codegen_partial = mw.resolve(ExperimentalReentrantCodegenProvider)[
            mw.module.body[0]
        ]
        self.assertEqual(codegen_partial.get_original_module_bytes(), module_bytes)
        self.assertEqual(
            codegen_partial.get_modified_module_bytes(cst.parse_statement("fn2()\n")),
            "fn2()\n".encode("utf-16"),
        )
