from typing import DefaultDict, List, Dict, Callable, Set

import collections
import linecache


from avroc.codegen.graph import find_recursive_types
from avroc.util import SchemaType, rand_str
from avroc.schema import expand_names, gather_named_types

import ast

try:
    from ast import unparse
except ImportError:
    from astunparse import unparse


class Compiler:
    schema: SchemaType
    variable_name_counts: DefaultDict[str, int]

    # List of schemas which are defined recursively. This is needed to generate
    # separate functions for each of these types so that they can call
    # themselves recursively when reading.
    recursive_types: List[Dict]

    # Types by absolute name
    named_types: Dict[str, SchemaType]

    types_referenced_by_name: Set[str]

    entrypoint_name: str

    debug: bool = False

    def __init__(self, schema: SchemaType, entrypoint_name: str):
        self.schema = expand_names(schema)
        self.named_types = gather_named_types(self.schema)
        self.types_referenced_by_name = set()

        self.entrypoint_name = entrypoint_name
        self.recursive_types = find_recursive_types(self.schema)

        self.variable_name_counts = collections.defaultdict(int)

    def new_variable(self, name: str) -> str:
        """
        Returns a new name for a variable which is guaranteed to be unique.
        """
        count = self.variable_name_counts[name]
        self.variable_name_counts[name] = count + 1
        if count == 0:
            return f"{name}"
        return f"{name}{count}"

    def generate_source_code(self) -> str:
        """
        Generates Python source code for a module which can handle the given schema.
        """
        module = self.generate_module()
        return unparse(module)

    def compile(self) -> Callable:
        """
        Compiles the given schema into a callable which can encode or decode avro
        messages.
        """

        # Generate the AST:
        module = self.generate_module()

        # Set any unset fields on AST nodes to 'None'. We do this so nothing is
        # undefine when using the third-party astunparse library, which expects
        # all fields to be set on AST node instances, even if set to None.
        for node in ast.walk(module):
            for field in node._fields:
                if not hasattr(node, field):
                    setattr(node, field, None)

                if isinstance(node, ast.Constant):
                    # See https://github.com/simonpercivall/astunparse/issues/60
                    if not hasattr(node, "kind"):
                        node.kind = None

        # Produce source code which is equivalent to the AST:
        source_code = unparse(module)
        if self.debug:
            print(source_code)

        # Insert the source code into linecache.cache so that stack traces and
        # other tools can browse the generated code:
        filename = f"autogenerated_file_{rand_str(8)}.py"
        compiled = compile(source_code, filename, mode="exec")
        source_lines = source_code.splitlines()
        linecache.cache[filename] = (  # type: ignore
            len(source_lines),
            None,
            source_lines,
            filename,
        )

        # Compile the AST:
        namespace = {}  # type: ignore
        exec(compiled, namespace)

        # Return the callable entrypoint:
        return namespace[self.entrypoint_name]  # type: ignore

    def generate_module(self):
        raise NotImplementedError("should be implemented by subclasses")
