import ast
import os
from jaspyx.ast_util import ast_load, ast_store, ast_call
from jaspyx.visitor import BaseVisitor


class Import(BaseVisitor):
    import_path = ['.']

    def load_module(self, pieces):
        module_name = '.'.join(pieces)
        if module_name in self.registry:
            return

        if len(pieces) > 1:
            parent = self.registry['.'.join(pieces[:-1])]
            import_path = [os.path.split(parent.path)[0]]
        else:
            import_path = self.import_path

        for path in import_path:
            module_path = os.path.join(path, pieces[-1], '__init__.jpx')
            if os.path.exists(module_path):
                break
            module_path = os.path.join(path, pieces[-1]) + '.jpx'
            if os.path.isfile(module_path):
                break
        else:
            raise ImportError('module %s not found' % module_name)

        c = ast.parse(open(module_path).read(), module_path)
        self.registry[module_name] = v = self.__class__(module_path, self.registry)
        v.import_path = self.import_path
        v.visit(c)

    def init_module(self, pieces):
        module_path = []
        for piece in pieces:
            module_path.append(piece)
            module_name = '.'.join(module_path)

            self.load_module(module_path)

            self.visit(
                ast.If(
                    ast.Compare(
                        ast_call(
                            ast_load('type'),
                            ast.Subscript(
                                ast_load('__modules__'),
                                ast.Index(ast.Str(module_name)),
                                ast.Load(),
                            )
                        ),
                        [ast.Eq()],
                        [ast.Str('undefined')]
                    ),
                    [
                        ast.Assign(
                            [ast.Subscript(
                                ast_load('__modules__'),
                                ast.Index(ast.Str(module_name)),
                                ast.Store(),
                            )],
                            ast.Dict(
                                keys=[
                                    ast.Str('__modules__'),
                                    ast.Str('__registry__'),
                                ],
                                values=[
                                    ast_load('__modules__'),
                                    ast_load('__registry__'),
                                ]
                            )
                        ),
                        ast.Expr(
                            ast_call(
                                ast.Subscript(
                                    ast_load('__registry__'),
                                    ast.Index(ast.Str(module_name)),
                                    ast.Load(),
                                ),
                                ast.Subscript(
                                    ast_load('__modules__'),
                                    ast.Index(ast.Str(module_name)),
                                    ast.Load(),
                                ),
                            )
                        )
                    ] + ([
                        ast.Assign(
                            [
                                ast.Attribute(
                                    ast.Subscript(
                                        ast_load('__modules__'),
                                        ast.Index(ast.Str('.'.join(module_path[:-1]))),
                                        ast.Load()
                                    ),
                                    module_path[-1],
                                    ast.Store()
                                )
                            ],
                            ast.Subscript(
                                ast_load('__modules__'),
                                ast.Index(ast.Str(module_name)),
                                ast.Load()
                            )
                        )
                    ] if len(module_path) > 1 else []),
                    []
                )
            )

    def visit_Import(self, node):
        for name in node.names:
            module_path = name.name.split('.')

            self.init_module(module_path)

            if not name.asname:
                self.visit(
                    ast.Assign(
                        [ast_store(module_path[0])],
                        ast_load('__modules__', module_path[0])
                    )
                )
            else:
                self.visit(
                    ast.Assign(
                        [ast_store(name.asname)],
                        ast_load('__modules__', *module_path)
                    )
                )

    def visit_ImportFrom(self, node):
        if node.level:
            raise NotImplementedError('Relative imports are not supported')

        module_path = node.module.split('.')

        self.init_module(module_path)

        for name in node.names:
            asname = name.asname if name.asname else name.name
            self.visit(
                ast.Assign(
                    [ast_store(asname)],
                    ast_load('__modules__', *(module_path + [name.name]))
                )
            )

@@ -5144,32 +5144,751 @@
  in node.names:%0A
+            if name.name == '*':%0A                if self.stack%5B-1%5D.scope.prefix != %5B'__module__'%5D:%0A                    raise NotImplementedError('from x import * only implemented at module level')%0A%0A                self.visit(ast.For(%0A                    ast_store('$tmp'),%0A                    ast_load('__modules__', node.module),%0A                    %5B%0A                        ast.Expr(ast_call(%0A                            ast_load('eval'),%0A                            ast.Str('__module__%5B__module__.$tmp%5D = __module__.__modules__.%25s%5B__module__.$tmp%5D' %25%0A                                    node.module),%0A                        ))%0A                    %5D,%0A                    %5B%5D%0A                ))%0A            else:%0A    
             asna
@@ -5938,32 +5938,36 @@
 ame%0A            
+    
 self.visit(%0A    
@@ -5970,32 +5970,36 @@
                 
+   
+ 
 ast.Assign(%0A    
@@ -5986,32 +5986,36 @@
     ast.Assign(%0A
+    
                 
@@ -6051,32 +6051,36 @@
                 
+    
 ast_load('__modu
@@ -6091,25 +6091,21 @@
 _', 
-*(
+node.
 module
-_path + %5B
+, 
 name
@@ -6113,12 +6113,13 @@
 name
-%5D)
 )%0A
+   
     
@@ -6122,32 +6122,37 @@
                 
+ 
 )%0A
+    
             )%0A

Support from x import * at module level.