import json
from pathlib import Path
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,
    # JobSettingsCodeGenerator,
    MasterCodeGenerator,
    # MatchSpecificationCodeGenerator,
    # MessageHandlerCodeGenerator,
    # ParamSetCodeGenerator,
    # SubflowCodeGenerator,
    # TestCaseCodeGenerator,
)
from ibm_watsonx_data_integration.services.datastage.codegen.dag_generator import (
    ConnectionGenerator,
    DAGGenerator,
)
from ibm_watsonx_data_integration.services.datastage.codegen.exporters.util import (
    _format_code,
    _generate_flow_name,
    _autogenerated_header,
)
from ibm_watsonx_data_integration.services.datastage.codegen.importers import ZipImporter


class SingleFileExporter:
    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())
            if not self.output_path.endswith(".py"):
                self.output_path += ".py"
            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 generate_flow_name(self, f_info):
        if self.use_flow_name:
            flow_name = Path(f_info.filename).stem
        else:
            flow_name = _generate_flow_name()
        return flow_name

    def run(self):
        master_gen = MasterCodeGenerator()
        # job_vars: list[str] = []
        all_code: list[str] = FlowCodeGenerator().generate_setup(self.api_key, self.project_id)
        errors = dict()
        formatted_all_code = _autogenerated_header()

        try:
            # for f_info, f_content in self.zip_importer.paramsets:
            #     json_data = json.loads(f_content)["parameter_set"]
            #     param_set = ParameterSet.from_dict(json_data)
            #     ps_code_gen = ParamSetCodeGenerator(param_set, master_gen)
            #     ps_code = ps_code_gen.generate_code()
            #     var_name = master_gen.get_object_var(param_set)
            #     all_code.append(ps_code)
            #     all_code.append(f"flow.use_paramset({var_name})")

            for f_info, f_content in self.zip_importer.connections:
                json_data = json.loads(f_content)
                conn_gen = ConnectionGenerator(json_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()
                # all_code.append(conn_code)

            # 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()
            #     all_code.append(data_def_code)

            # for f_info, f_content in self.zip_importer.message_handlers:
            #     json_data = json.loads(f_content)
            #     properties = json_data["entity"]
            #     properties["name"] = json_data["name"]
            #     properties["description"] = json_data["description"]
            #     message_handler = MessageHandler(**properties)
            #     mh_code_gen = MessageHandlerCodeGenerator(message_handler, master_gen)
            #     mh_code = mh_code_gen.generate_code()
            #     var_name = master_gen.get_object_var(message_handler)
            #     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, os.getcwd(), master_gen)
            #     jl_code = jl_code_gen.generate_code()
            #     formatted = _format_code(jl_code)
            #     all_code.append(f"{formatted}\n\n")

            for f_info, f_content in self.zip_importer.flows:
                flow_name = self.generate_flow_name(f_info)

                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}")

                dag_gen = DAGGenerator(flow_model)
                fc = dag_gen.generate()
                # replace_nodes = {}

                # KEEP TRACK OF USED SUBFLOWS AND CREATE OTHERS AFTER
                # 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)
                #         dag = subflow_dag_gen.generate()._dag
                #         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,
                #         )
                #         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))
                #         if self.write_to_output:
                #             attachment_info, attachment_content = self.zip_importer.find_custom_stage_attachment(f"lib_{custom_stage_name}")
                #             _check_create_dir(Path(self.output_path).parent / "attachments")
                #             attachment_path = Path(self.output_path).parent / "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])

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

                # for f_info, f_content in self.zip_importer.subflows:
                #     subflow_json = json.loads(f_content)
                #     subflow_model = models.Flow(**subflow_json)
                #     subflow_dag_gen = DAGGenerator(subflow_model)
                #     dag = subflow_dag_gen.generate()._dag
                #     subflow = Subflow(dag=dag, name=subflow_model.name, is_local=False)
                #     subflow_code_gen = SubflowCodeGenerator(subflow=subflow, master_gen=master_gen, offline=self.offline)
                #     subflow_code = subflow_code_gen.generate_all()
                #     all_code.append(subflow_code)

                # for f_info, f_content in self.zip_importer.jobs:
                #     job_json = json.loads(f_content)["entity"]["job"]
                #     job_gen = JobGenerator(job_json)
                #     job_settings = job_gen.create_job_model()
                #     job_code_gen = JobSettingsCodeGenerator(job_settings, master_gen)
                #     job_code = job_code_gen.generate_code()
                #     job_var = master_gen.get_object_var(job_settings)
                #     job_vars.append(job_var)
                #     all_code.append(job_code)

                # 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")

                all_code.append(f"sdk.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="")')

            # for f_info, f_content in self.zip_importer.test_cases:
            #     json_data = json.loads(f_content)
            #     specification = json_data["entity"]["specification"]
            #     given = specification["given"]
            #     then = specification["then"]
            #     when = specification["when"]

            #     tcc = TestCaseComposer(
            #         name=json_data["metadata"]["name"],
            #         flow_name=flow_name,
            #     )

            #     for link in given:
            #         tcc.use_input_test_data((link["link"], link["path"]))
            #     for link in then:
            #         tcc.use_output_test_data((link["link"], link["path"]))
            #         if "checkRowCountOnly" in link and link["checkRowCountOnly"]:
            #             tcc.use_row_count_only(link["link"])
            #         if "ignore" in link:
            #             tcc.exclude_columns((link["link"], link["ignore"]))
            #         if "cluster" in link:
            #             tcc.use_cluster_key((link["link"], link["cluster"]))

            #     tcc.use_parameters(when["parameters"])
            #     test_case_gen = TestCaseCodeGenerator(tcc, master_gen)
            #     test_case_code = test_case_gen.generate_code()
            #     all_code.append(test_case_code)
            #     all_code.append(f'sdk.run_test_case(tcc={master_gen.get_object_var(tcc)}, test_case_name="dfsf", print_logs={self.print_logs})')

            formatted_all_code += _format_code("\n".join(all_code))

        except Exception:
            errors[Path(f_info.filename).stem] = traceback.format_exc()

        if self.write_to_output:
            with open(self.output_path, "w") as f:
                f.write(formatted_all_code)
            return {self.output_path: formatted_all_code}, errors
        return {flow_name + ".py": formatted_all_code}, errors
