from pickle import dump, load

from ..core import PPQ_CONFIG
from ..IR import BaseGraph, GraphExporter
from ..IR.base.graph import GraphBuilder


class NativeExporter(GraphExporter):
    def __init__(self) -> None:
        super().__init__()

    def export(self, file_path: str, graph: BaseGraph, config_path: str = None, dump_value: bool = True):
        def dump_elements_to_file(file, elements: list):
            for element in elements:
                dump(element, file)

        with open(file_path, "wb") as file:
            dump_elements_to_file(
                file,
                elements=[
                    "PPQ GRAPH DEFINITION",  # PPQ Signature.
                    PPQ_CONFIG.VERSION,  # PPQ Signature.
                    graph,
                ],
            )


class NativeImporter(GraphBuilder):
    def __init__(self) -> None:
        super().__init__()

    def build(self, file_path: str, **kwargs) -> BaseGraph:
        def load_elements_from_file(file, num_of_elements: int) -> list:
            try:
                return [load(file) for _ in range(num_of_elements)]
            except EOFError as e:
                raise Exception("File format parsing error. Unexpected EOF found.")

        with open(file_path, "rb") as file:
            signature, version, graph = load_elements_from_file(file, 3)
            if signature != "PPQ GRAPH DEFINITION":
                raise Exception("File format parsing error. Graph Signature has been damaged.")
            if str(version) > PPQ_CONFIG.VERSION:
                print(
                    f"\033[31mWarning: Dump file is created by PPQ({str(version)}), "
                    f"however you are using PPQ({PPQ_CONFIG.VERSION}).\033[0m"
                )

        assert isinstance(graph, BaseGraph), "File format parsing error. Graph Definition has been damaged."
        try:
            for op in graph.operations.values():
                input_copy, _ = op.inputs.copy(), op.inputs.clear()
                for name in input_copy:
                    op.inputs.append(graph.variables[name])
                output_copy, _ = op.outputs.copy(), op.outputs.clear()
                for name in output_copy:
                    op.outputs.append(graph.variables[name])

            for var in graph.variables.values():
                dest_copy, _ = var.dest_ops.copy(), var.dest_ops.clear()
                for name in dest_copy:
                    var.dest_ops.append(graph.operations[name])
                if var.source_op is not None:
                    var.source_op = graph.operations[var.source_op]

            graph._graph_inputs = {name: graph.variables[name] for name in graph._graph_inputs}
            graph._graph_outputs = {name: graph.variables[name] for name in graph._graph_outputs}
        except Exception as e:
            raise Exception("File format parsing error. Graph Definition has been damaged.")
        return graph
