# Copyright (c) OpenMMLab. All rights reserved.
import ast
import copy
import os
import os.path as osp
from importlib import import_module
from importlib.util import find_spec
from unittest import TestCase

import numpy
import numpy.compat
import numpy.linalg as linalg

import mmengine
from mmengine.config import Config
from mmengine.config.lazy import LazyAttr, LazyObject
from mmengine.config.utils import ImportTransformer, _gather_abs_import_lazyobj
from mmengine.fileio import LocalBackend, PetrelBackend


class TestImportTransformer(TestCase):

    @classmethod
    def setUpClass(cls) -> None:
        cls.data_dir = osp.join(  # type: ignore
            osp.dirname(__file__), '..', 'data', 'config',
            'lazy_module_config')
        super().setUpClass()

    def test_lazy_module(self):
        cfg_path = osp.join(self.data_dir, 'test_ast_transform.py')
        with open(cfg_path) as f:
            codestr = f.read()
        codeobj = ast.parse(codestr)
        global_dict = {
            'LazyObject': LazyObject,
        }
        base_dict = {
            '._base_.default_runtime': {
                'default_scope': 'test_config'
            },
            '._base_.scheduler': {
                'val_cfg': {}
            },
        }
        codeobj = ImportTransformer(global_dict, base_dict).visit(codeobj)
        codeobj, _ = _gather_abs_import_lazyobj(codeobj)
        codeobj = ast.fix_missing_locations(codeobj)

        exec(compile(codeobj, cfg_path, mode='exec'), global_dict, global_dict)
        # 1. absolute import
        # 1.1 import module as LazyObject
        lazy_numpy = global_dict['numpy']
        self.assertIsInstance(lazy_numpy, LazyObject)

        # 1.2 getattr as LazyAttr
        self.assertIsInstance(lazy_numpy.linalg, LazyAttr)
        self.assertIsInstance(lazy_numpy.compat, LazyAttr)

        # 1.3 Build module from LazyObject. amp and functional can be accessed
        imported_numpy = lazy_numpy.build()
        self.assertIs(imported_numpy.linalg, linalg)
        self.assertIs(imported_numpy.compat, numpy.compat)

        # 1.4 Build module from LazyAttr
        imported_linalg = lazy_numpy.linalg.build()
        imported_compat = lazy_numpy.compat.build()
        self.assertIs(imported_compat, numpy.compat)
        self.assertIs(imported_linalg, linalg)

        # 1.5 import ... as, and build module from LazyObject
        lazy_linalg = global_dict['linalg']
        self.assertIsInstance(lazy_linalg, LazyObject)
        self.assertIs(lazy_linalg.build(), linalg)
        self.assertIsInstance(lazy_linalg.norm, LazyAttr)
        self.assertIs(lazy_linalg.norm.build(), linalg.norm)

        # 1.6 import built in module
        imported_os = global_dict['os']
        self.assertIs(imported_os, os)

        # 2. Relative import
        # 2.1 from ... import ...
        lazy_local_backend = global_dict['local']
        self.assertIsInstance(lazy_local_backend, LazyObject)
        self.assertIs(lazy_local_backend.build(), LocalBackend)

        # 2.2 from ... import ... as ...
        lazy_petrel_backend = global_dict['PetrelBackend']
        self.assertIsInstance(lazy_petrel_backend, LazyObject)
        self.assertIs(lazy_petrel_backend.build(), PetrelBackend)

        # 2.3 from ... import builtin module or obj from `mmengine.Config`
        self.assertIs(global_dict['find_module'], find_spec)
        self.assertIs(global_dict['Config'], Config)

        # 3 test import base config
        # 3.1 simple from ... import and from ... import ... as
        self.assertEqual(global_dict['scope'], 'test_config')
        self.assertDictEqual(global_dict['val_cfg'], {})

        # 4. Error catching
        cfg_path = osp.join(self.data_dir,
                            'test_ast_transform_error_catching1.py')
        with open(cfg_path) as f:
            codestr = f.read()
        codeobj = ast.parse(codestr)
        global_dict = {'LazyObject': LazyObject}
        with self.assertRaisesRegex(
                RuntimeError,
                r'Illegal syntax in config! `from xxx import \*`'):
            codeobj = ImportTransformer(global_dict).visit(codeobj)


class TestLazyObject(TestCase):

    def test_init(self):
        LazyObject('mmengine')
        LazyObject('mmengine.fileio')
        LazyObject('mmengine.fileio', 'LocalBackend')

        # module must be str
        with self.assertRaises(TypeError):
            LazyObject(1)

        # imported must be a sequence of string or None
        with self.assertRaises(TypeError):
            LazyObject('mmengine', ['error_type'])

    def test_build(self):
        lazy_mmengine = LazyObject('mmengine')
        self.assertIs(lazy_mmengine.build(), mmengine)

        lazy_mmengine_fileio = LazyObject('mmengine.fileio')
        self.assertIs(lazy_mmengine_fileio.build(),
                      import_module('mmengine.fileio'))

        lazy_local_backend = LazyObject('mmengine.fileio', 'LocalBackend')
        self.assertIs(lazy_local_backend.build(), LocalBackend)

        # TODO: The commented test is required, we need to test the built
        # LazyObject can access the `mmengine.dataset`. We need to clean the
        # environment to make sure the `dataset` is not imported before, and
        # it is triggered by lazy_mmengine.build(). However, if we simply
        # pop the `mmengine.dataset` will lead to other tests failed, of which
        # reason is still unknown. We need to figure out the reason and fix it
        # in the latter

        # sys.modules.pop('mmengine.config')
        # sys.modules.pop('mmengine.fileio')
        # sys.modules.pop('mmengine')
        # lazy_mmengine = LazyObject(['mmengine', 'mmengine.dataset'])
        # self.assertIs(lazy_mmengine.build().dataset,
        #               import_module('mmengine.config'))
        copied = copy.deepcopy(lazy_local_backend)
        self.assertDictEqual(copied.__dict__, lazy_local_backend.__dict__)

        with self.assertRaises(RuntimeError):
            lazy_mmengine()

        with self.assertRaises(ImportError):
            LazyObject('unknown').build()


class TestLazyAttr(TestCase):
    # Since LazyAttr should only be built from LazyObect, we only test
    # the build method here.
    def test_build(self):
        lazy_mmengine = LazyObject('mmengine')
        local_backend = lazy_mmengine.fileio.LocalBackend
        self.assertIs(local_backend.build(), LocalBackend)

        copied = copy.deepcopy(local_backend)
        self.assertDictEqual(copied.__dict__, local_backend.__dict__)

        with self.assertRaises(RuntimeError):
            local_backend()

        with self.assertRaisesRegex(
                ImportError,
                'Failed to import mmengine.fileio.LocalBackend.unknown'):
            local_backend.unknown.build()