import json
from pathlib import Path
from zipfile import ZipInfo
import traceback

import ibm_watsonx_data_integration.services.datastage.models.flow_json_model as models
from ibm_watsonx_data_integration.services.datastage.codegen.code_generator import (
    FlowCodeGenerator,
    # FunctionLibraryCodeGenerator,
    # JavaLibraryCodeGenerator,
    MasterCodeGenerator,
    # MatchSpecificationCodeGenerator,
    # MessageHandlerCodeGenerator,
    # ParamSetCodeGenerator,
)
from ibm_watsonx_data_integration.services.datastage.codegen.dag_generator import (
    ConnectionGenerator,
    DAGGenerator,
    # FunctionLibraryGenerator,
    # JavaLibraryGenerator,
    # MatchSpecificationGenerator,
)
from ibm_watsonx_data_integration.services.datastage.codegen.exporters.util import (
    _check_create_dir,
    _format_code,
    _generate_flow_name,
    _autogenerated_header,
)
from ibm_watsonx_data_integration.services.datastage.codegen.importers import ZipImporter


class FlowFileExporter:
    def __init__(
        self,
        zip_importer: ZipImporter,
        output_path: str | None,
        *,
        offline: bool = False,
        create_job: bool = True,
        run_job: bool = True,
        use_flow_name: bool = True,
        api_key: str = "<TODO: insert your api_key>",
        project_id: str = "<TODO: insert your project_id>",
    ):
        self.zip_importer = zip_importer
        if output_path is not None and output_path.strip() != "":
            self.output_path = str(Path(output_path).absolute())
            self.write_to_output = True
        else:
            self.output_path = ""
            self.write_to_output = False
        self.offline = offline
        self.create_job = create_job
        self.run_job = run_job
        self.use_flow_name = use_flow_name
        self.api_key = api_key
        self.project_id = project_id
        self.job_objs = set()

    # def _create_subflows(
    #     self,
    #     master_gen: MasterCodeGenerator,
    #     parent_dag: DAG,
    #     paramset_objs: list[ParameterSet],
    #     fc: DataStageFlow,
    # ):
    #     replace_nodes: dict[Node, Node] = {}

    #     for node in parent_dag.nodes():
    #         if isinstance(node, SuperNodeRef):
    #             subflow_name = node.name
    #             subflow_json = json.loads(self.zip_importer.find_subflow(subflow_name))
    #             subflow_model = models.Flow(**subflow_json)
    #             subflow_dag_gen = DAGGenerator(subflow_model)
    #             sfc = subflow_dag_gen.generate()
    #             dag = sfc._dag

    #             self._create_subflows(
    #                 master_gen=master_gen,
    #                 parent_dag=dag,
    #                 paramset_objs=paramset_objs,
    #                 fc=fc,
    #             )

    #             sfc_paramsets: list[ParameterSet] = []
    #             for paramset in sfc.parameter_sets:
    #                 for paramset_obj in paramset_objs:
    #                     if paramset_obj.name == paramset.name:
    #                         sfc_paramsets.append(paramset_obj)

    #             subflow = Subflow(dag=dag, name=subflow_name, is_local=False)
    #             super_node = SuperNode(
    #                 parent_dag=parent_dag,
    #                 subflow_dag=subflow.dag,
    #                 entry_nodes=subflow.entry_nodes,
    #                 exit_nodes=subflow.exit_nodes,
    #                 name=subflow_name,
    #                 label=node.label,
    #                 parameter_sets=sfc_paramsets,
    #                 local_parameters=sfc.local_parameters,
    #                 rcp=node.rcp,
    #             )
    #             replace_nodes[node] = super_node

    #     for node in replace_nodes:
    #         parent_dag.replace_node(node, replace_nodes[node])

    def _export_flow(self, master_gen: MasterCodeGenerator, f_content: bytes, f_info: ZipInfo):
        all_code: list[str] = FlowCodeGenerator().generate_setup(self.api_key, self.project_id)

        flow_json = json.loads(f_content)
        try:
            flow_model = models.Flow(**flow_json)
        except Exception:
            try:
                flow_json = flow_json["attachments"]
                flow_model = models.Flow(**flow_json)
            except Exception:
                try:
                    flow_json = flow_json[0]
                    flow_model = models.Flow(**flow_json)
                except Exception as e:
                    raise ValueError(f"Bad flow json: {e}")

        # paramset_objs: list[ParameterSet] = []
        # if flow_model.external_paramsets:
        #     for paramset in flow_model.external_paramsets:
        #         ps_content = self.zip_importer.find_paramset(paramset.name)
        #         ps_data = json.loads(ps_content)["parameter_set"]
        #         param_set = ParameterSet.from_dict(ps_data)
        #         paramset_objs.append(param_set)
        #         ps_code_gen = ParamSetCodeGenerator(param_set, master_gen)
        #         ps_code = ps_code_gen.generate_code()
        #         var_name = master_gen.get_object_var(param_set)
        #         ps_code += f"\nflow.use_paramset({var_name})\n"
        #         formatted = _format_code(ps_code)
        #         all_code.append(f"{formatted}\n\n")

        for conn_info, conn_content in self.zip_importer.connections:
            conn_data = json.loads(conn_content)
            conn_gen = ConnectionGenerator(conn_data)
            conn_model = conn_gen.create_connection_model()
            conn_model
            # conn_code_gen = ConnectionCodeGenerator(conn_model, master_gen)
            # conn_code = conn_code_gen.generate_code()
            # formatted = _format_code(conn_code)
            # all_code.append(f"{formatted}\n\n")

        # for f_info, f_content in self.zip_importer.data_definitions:
        #     json_data = json.loads(f_content)
        #     data_def = DataDefinition.from_dict(json_data)
        #     data_def_code_gen = DataDefinitionCodeGenerator(data_def, master_gen)
        #     data_def_code = data_def_code_gen.generate_code()
        #     formatted = _format_code(data_def_code)
        #     all_code.append(f"{formatted}\n\n")

        # for mh_info, mh_content in self.zip_importer.message_handlers:
        #     mh_data = json.loads(mh_content)
        #     properties = mh_data["entity"]
        #     properties["name"] = mh_data["name"]
        #     properties["description"] = mh_data["description"]
        #     message_handler = MessageHandler(**properties)
        #     mh_code_gen = MessageHandlerCodeGenerator(message_handler, master_gen)
        #     mh_code = mh_code_gen.generate_code()
        #     formatted = _format_code(mh_code)
        #     all_code.append(f"{formatted}\n\n")

        # for jl_info, jl_content in self.zip_importer.java_libraries:
        #     jl_data = json.loads(jl_content)
        #     jl_gen = JavaLibraryGenerator(jl_data)
        #     jl_model = jl_gen.create_java_library_model(self.output_path)
        #     jl_code_gen = JavaLibraryCodeGenerator(jl_model, self.output_path, master_gen)
        #     jl_code = jl_code_gen.generate_code()
        #     formatted = _format_code(jl_code)
        #     all_code.append(f"{formatted}\n\n")

        # for fl_info, fl_content in self.zip_importer.function_libraries:
        #     fl_data = json.loads(fl_content)
        #     fl_gen = FunctionLibraryGenerator(fl_data)
        #     fl_model = fl_gen.create_function_library_model(self.output_path)
        #     fl_code_gen = FunctionLibraryCodeGenerator(fl_model, self.output_path, master_gen)
        #     fl_code = fl_code_gen.generate_code()
        #     formatted = _format_code(fl_code)
        #     all_code.append(f"{formatted}\n\n")

        # for ms_info, ms_content in self.zip_importer.match_specifications:
        #     ms_data = json.loads(ms_content)
        #     ms_gen = MatchSpecificationGenerator(ms_data)
        #     ms_model = ms_gen.create_match_specification_model(self.output_path)
        #     ms_code_gen = MatchSpecificationCodeGenerator(ms_model, self.output_path, master_gen)
        #     ms_code = ms_code_gen.generate_code()
        #     formatted = _format_code(ms_code)
        #     all_code.append(f"{formatted}\n\n")

        dag_gen = DAGGenerator(flow_model)
        fc = dag_gen.generate()

        # self._create_subflows(
        #     master_gen=master_gen,
        #     parent_dag=fc._dag,
        #     paramset_objs=paramset_objs,
        #     fc=fc,
        # )

        # replace_nodes = {}
        # for node in fc._dag.nodes():
        # if isinstance(node, SuperNodeRef):
        #     subflow_name = node.name
        #     subflow_json = json.loads(self.zip_importer.find_subflow(subflow_name))
        #     subflow_model = models.Flow(**subflow_json)
        #     subflow_dag_gen = DAGGenerator(subflow_model)
        #     sfc = subflow_dag_gen.generate()
        #     dag = sfc._dag

        #     # sfc_paramsets: list[ParameterSet] = []
        #     # for paramset in sfc.parameter_sets:
        #     #     for paramset_obj in paramset_objs:
        #     #         if paramset_obj.name == paramset.name:
        #     #             sfc_paramsets.append(paramset_obj)

        #     subflow = Subflow(dag=dag, name=subflow_name, is_local=False)
        #     super_node = SuperNode(
        #         parent_dag=fc._dag,
        #         subflow_dag=subflow.dag,
        #         entry_nodes=subflow.entry_nodes,
        #         exit_nodes=subflow.exit_nodes,
        #         name=subflow_name,
        #         label=node.label,
        #         parameter_sets=sfc_paramsets,
        #         local_parameters=sfc.local_parameters,
        #     )
        #     replace_nodes[node] = super_node

        # elif isinstance(node, BuildStageStage):
        #     build_stage_name = node.configuration.op_name
        #     build_stage_json = json.loads(self.zip_importer.find_build_stage(build_stage_name))
        #     build_stage_asset = BuildStage.from_dict(build_stage_json)
        #     node.build_stage = build_stage_asset

        # elif isinstance(node, WrappedStageStage):
        #     wrapped_stage_name = node.configuration.op_name
        #     wrapped_stage_json = json.loads(self.zip_importer.find_wrapped_stage(wrapped_stage_name))
        #     wrapped_stage_asset = WrappedStage.from_dict(wrapped_stage_json)
        #     node.wrapped_stage = wrapped_stage_asset

        # elif isinstance(node, CustomStageStage):
        #     custom_stage_name = node.configuration.op_name
        #     custom_stage_json = json.loads(self.zip_importer.find_custom_stage(custom_stage_name))
        #     attachment_info, attachment_content = self.zip_importer.find_custom_stage_attachment(f"lib_{custom_stage_name}")
        #     if self.write_to_output:
        #         _check_create_dir(Path(self.output_path) / "attachments")
        #         attachment_path = Path(self.output_path) / "attachments" / (Path(attachment_info.filename).stem + ".so")
        #         with open(attachment_path, "w") as f:
        #             f.write(str(base64.b64encode(attachment_content))[2:-1])
        #     if "entity" in custom_stage_json:
        #         custom_stage_json["entity"]["library_path"] = str(attachment_path)
        #     else:
        #         custom_stage_json["library_path"] = str(attachment_path)
        #     custom_stage_asset = CustomStage.from_dict(custom_stage_json)
        #     node.custom_stage = custom_stage_asset

        # for node in replace_nodes:
        #     fc._dag.replace_node(node, replace_nodes[node])

        if self.use_flow_name:
            flow_name = Path(f_info.filename).stem
        else:
            flow_name = _generate_flow_name()

        code_gen = FlowCodeGenerator(flow_name, fc, master_gen, offline=self.offline, skip_subflows=True)
        code = code_gen.generate_all_zip_one_file()
        all_code.append(code)

        all_code.append(f"\n\nproject.update_flow({code_gen.composer})")

        if self.create_job:
            job_obj = object()
            self.job_objs.add(job_obj)
            job_var = master_gen.reserve_var(f"{flow_name}_job", job_obj)
            all_code.append(f'\n\n{job_var} = project.create_job(name="{flow_name}_job", flow={code_gen.composer})')
            if self.run_job:
                job_run_obj = object()
                self.job_objs.add(job_run_obj)
                job_run_var = master_gen.reserve_var(f"{flow_name}_job_run", job_run_obj)
                all_code.append(f'\n\n{job_run_var} = {job_var}.start(name="{flow_name}_job_run", description="")')

        return _format_code("\n".join(all_code))

    def run(self):
        generated_code = {}
        errors = {}
        if self.write_to_output:
            _check_create_dir(self.output_path)

        master_gen = MasterCodeGenerator()
        for f_info, f_content in self.zip_importer.flows:
            try:
                code = _autogenerated_header() + self._export_flow(master_gen, f_content, f_info)
                file_name = Path(self.output_path) / (Path(f_info.filename).stem + ".py")
                generated_code[str(file_name)] = code
                if self.write_to_output:
                    with open(file_name, "w") as f:
                        f.write(code)
            except Exception:
                errors[Path(f_info.filename).stem] = traceback.format_exc()

        return generated_code, errors
