from collections import defaultdict
import json
from pathlib import Path
import traceback
import os
import logging

from ibm_watsonx_data_integration.cpd_models.parameter_set_model import ParameterSet
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
from ibm_watsonx_data_integration.services.datastage.models.flow.subflow import Subflow
from ibm_watsonx_data_integration.services.datastage.models.flow.dag import SuperNodeRef, DAG, SuperNode

logger = logging.getLogger(__name__)


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>",
        base_auth_url: str = "https://cloud.ibm.com",
        base_api_url: str = "https://api.ca-tor.dai.cloud.ibm.com",
    ):
        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()
        self.base_auth_url = base_auth_url
        self.base_api_url = base_api_url

    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, base_auth_url=self.base_auth_url, base_api_url=self.base_api_url
        )
        errors = dict()
        formatted_all_code = _autogenerated_header()

        try:
            paramset_objs: list[ParameterSet] = []
            external_paramsets: dict[str, str] = dict()
            for f_info, f_content in self.zip_importer.paramsets:
                json_data = {"entity": json.loads(f_content)}
                param_set = ParameterSet(project=None, **json_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)
                all_code.append(f"{ps_code}\n\n")
                external_paramsets[param_set.name] = var_name

            for f_info, f_content in self.zip_importer.connections:
                try:
                    json_data = json.loads(f_content)
                    conn_gen = ConnectionGenerator(json_data)
                    conn_gen.create_connection_model()
                except Exception as exception:
                    logger.warning(f"Failed to generate a connection: {exception}")

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

                # key: DAG
                # value: dict of supernoderefs to replace
                replace_nodes: defaultdict[dict] = defaultdict(dict)

                def replace_super_node_refs(dtc: DAG) -> None:
                    for node in dtc.nodes():
                        if isinstance(node, SuperNodeRef):
                            subflow_name = node.name

                            try:
                                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:
                                    found = False
                                    for paramset_obj in paramset_objs:
                                        if paramset_obj.name == paramset.name:
                                            sfc_paramsets.append(paramset_obj)
                                            found = True
                                            break
                                    if not found:
                                        sfc_paramsets.append(paramset)
                                super_node = Subflow(
                                    parent_dag=dtc,
                                    dag=dag,
                                    name=subflow_name,
                                    label=node.label,
                                    parameter_sets=sfc_paramsets,
                                    local_parameters=sfc.local_parameters,
                                    local_parameter_values=node._local_parameter_values,
                                    is_local=False,
                                )
                            except KeyError as key_error:
                                # if there is a key error that means the subflow does not exist in the zip
                                # to not crash the entire code generation, we will just create an empty subflow
                                logger.warning(str(key_error))
                                super_node = Subflow(
                                    parent_dag=dtc, dag=DAG(), name=f"{subflow_name} MISSING", label=node.label, is_local=False
                                )

                            replace_nodes[dtc][node] = super_node
                            replace_super_node_refs(super_node._dag)
                        elif isinstance(node, SuperNode):
                            replace_super_node_refs(node._dag)

                replace_super_node_refs(fc._dag)

                # 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
                #         super_node = Subflow(parent_dag=fc._dag, dag=dag, name=subflow_name, label=node.label, is_local=False)
                #         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])
                for dag, node_dict in replace_nodes.items():
                    for node, replacement in node_dict.items():
                        dag.replace_node(node, replacement)

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

                # for each parameter set, try to use it. If it does not exist, we simply will add a placeholder
                if flow_model.external_paramsets:
                    for paramset in flow_model.external_paramsets:
                        if paramset.name in external_paramsets:
                            all_code.insert(0, "from ibm_watsonx_data_integration.cpd_models import ValueSet")
                            all_code.append(f"{code_gen.composer}.use_parameter_set({external_paramsets[paramset.name]})")
                        else:
                            all_code.extend(
                                [
                                    f"# NOTE: Parameter set '{paramset.name}' is missing from the zip file",
                                    f'{code_gen.composer}.use_parameter_set(project.create_parameter_set("PLACEHOLDER FOR {paramset.name}"))',
                                ]
                            )

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

            # 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:
            os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
            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
