import json
import os
from pathlib import Path
from typing import Literal
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,
    MasterCodeGenerator,
    # TestCaseCodeGenerator,
)
from ibm_watsonx_data_integration.services.datastage.codegen.dag_generator import DAGGenerator
from ibm_watsonx_data_integration.services.datastage.codegen.exporters.flow_file_exporter import FlowFileExporter
from ibm_watsonx_data_integration.services.datastage.codegen.exporters.single_file_exporter import SingleFileExporter
from ibm_watsonx_data_integration.services.datastage.codegen.exporters.util import (
    _autogenerated_header,
)
from ibm_watsonx_data_integration.services.datastage.codegen.importers import ZipImporter

# from ibm_watsonx_data_integration.services.datastage.models.test_case_composer import TestCaseComposer


def _delete_output_path(output_path: str):
    """
    Deletes the output path if it exists.
    """
    if os.path.exists(output_path):
        if os.path.isfile(output_path):
            os.remove(output_path)
        elif os.path.isdir(output_path):
            for root, dirs, files in os.walk(output_path, topdown=False):
                for name in files:
                    os.remove(os.path.join(root, name))
                for name in dirs:
                    os.rmdir(os.path.join(root, name))
            os.rmdir(output_path)


def _get_extension(file_path: str) -> str:
    """
    Returns the file extension of the given file path.
    """
    return Path(file_path).suffix.lstrip(".").lower()


ExportMode = Literal["single_file", "file_per_asset", "file_per_flow"]


# need to have options for:
# converting everything locally
# using api for conversion
# preserve positioning
# optimize
class PythonGeneratorConfig:
    """
    The underlying settngs for the python generator.

    Args:
        mode: The export mode, or how the generated code will be structured. Options are:
            - ``single_file``: All code in one file.
            - ``file_per_flow``: (default) Each flow in its own file, with relevant assets (connections, subflows, parameter sets, components, etc) included as needed.
        create_job: Whether to create a job for the flow(s) (default is True).
        run_job: Whether to run the created job
        overwrite: Whether to overwrite the directory being written to (default is False).
        use_flow_name: Whether the run/create flow statement should use the flow's original name (applies for zip structure only, default is True).
        api_key: TODO
        project_id: TODO
    """

    def __init__(
        self,
        *,
        mode: ExportMode = "file_per_flow",
        create_job: bool = True,
        run_job: bool = True,
        overwrite: bool = False,
        use_flow_name: bool = True,
        api_key: str = "<TODO: insert your api_key>",
        project_id: str = "<TODO: insert your project_id>",
    ):
        self.create_job = create_job
        self.run_job = run_job
        self.mode = mode
        self.overwrite = overwrite
        self.use_flow_name = use_flow_name
        self.api_key = api_key
        self.project_id = project_id


class PythonGenerator:
    """
    Converts DataStage flows exported as a flow JSON or ZIP file to Python SDK code.
    """

    def __init__(self):
        self.configuration = PythonGeneratorConfig()

    def get_code(self, input_path):
        """
        Returns the generated code as a string. Applicable only for JSON inputs.
        """
        if _get_extension(input_path) != "json":
            raise ValueError(f"Input file must be a JSON file, got '{input_path}'")

        with open(input_path, "r") as f:
            flow_json = json.load(f)

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

        dag_gen = DAGGenerator(flow_model)
        fc = dag_gen.generate()
        code_gen = FlowCodeGenerator(fc=fc, master_gen=MasterCodeGenerator())
        header = _autogenerated_header()
        setup = "\n".join(code_gen.generate_setup(self.configuration.api_key, self.configuration.project_id))
        code = code_gen.generate_all()
        return header + setup + code, code_gen

    # def get_test_case_code(self, input_path):
    #     """
    #     Returns the generated code as a string. Applicable only for JSON inputs.
    #     """
    #     with open(input_path, "r") as f:
    #         flow_json = json.load(f)

    #     tcc = TestCaseComposer.from_dict(flow_json)
    #     code_gen = TestCaseCodeGenerator(test_case_composer=tcc, master_gen=MasterCodeGenerator())
    #     code = "from ibm.datastage._framework.test_case_composer import TestCaseComposer\n"
    #     code += "from ibm.datastage._framework.sdk import DataStageSDK\n"
    #     code += "from ibm.datastage.config.config import AutoDetectConfig\n\n"
    #     code += "config = AutoDetectConfig()\n"
    #     code += "sdk = DataStageSDK(config)\n"
    #     code += code_gen.generate_code()
    #     return code

    def _generate_from_json(self, input_path, output_path=None):
        generated_code = ""
        errors = {}

        try:
            generated_code, code_gen = self.get_code(input_path)
            generated_code += f"\n\nproject.update_flow({code_gen.composer})"
            if self.configuration.create_job:
                job_name = code_gen.flow_name + "_job"
                job_var = "job_1"
                generated_code += f'\n\n{job_var} = project.create_job(name="{job_name}", flow=flow)'
                if self.configuration.run_job:
                    job_run_var = "job_run_1"
                    generated_code += f'\n\n{job_run_var} = {job_var}.start(name="{job_name}_job", description="")'
            if output_path:
                with open(output_path, "w") as f:
                    f.write(generated_code)
        except Exception:
            errors[Path(input_path).stem] = traceback.format_exc()
        return {output_path: generated_code}, errors

    def _generate_from_zip(self, input_path, output_path=None):
        zip_importer = ZipImporter(input_path)
        zip_importer.run()

        match self.configuration.mode:
            case "single_file":
                return SingleFileExporter(
                    zip_importer=zip_importer,
                    output_path=output_path,
                    create_job=self.configuration.create_job,
                    run_job=self.configuration.run_job,
                    use_flow_name=self.configuration.use_flow_name,
                    api_key=self.configuration.api_key,
                    project_id=self.configuration.project_id,
                ).run()
            case "file_per_flow":
                return FlowFileExporter(
                    zip_importer=zip_importer,
                    output_path=output_path,
                    create_job=self.configuration.create_job,
                    run_job=self.configuration.run_job,
                    use_flow_name=self.configuration.use_flow_name,
                    api_key=self.configuration.api_key,
                    project_id=self.configuration.project_id,
                ).run()
            # case "file_per_asset": # Removing this feature for now
            #     return MultiFileExporter(
            #         zip_importer=zip_importer,
            #         output_path=output_path,
            #         create_job=self.configuration.create_job,
            #         use_flow_name=self.configuration.use_flow_name,
            #     ).run()

    def generate(self, input_path, output_path=None):
        """
        Runs the Python converter, generating Python code and writing it to the specified output path.

        Args:
            input_path: The path to the input file. Can be a flow JSON file or a ZIP file containing the flow and its dependencies.
            output_path (optional): The output path where the generated Python code will be written.

        Returns:
            A dictionary containing the generated code strings.
        """

        if not input_path:
            raise ValueError("Input path must be specified.")

        if self.configuration.overwrite:
            _delete_output_path(output_path)

        match _get_extension(input_path):
            case "json":
                output_path = output_path or "code"
                generated_code, errors = self._generate_from_json(input_path, output_path)
                for filename, err in errors.items():
                    print(err)
                return generated_code, errors
            case "zip":
                generated_code, errors = self._generate_from_zip(input_path, output_path)
                for filename, err in errors.items():
                    print(err)
                return generated_code, errors
            case _:
                raise ValueError(f"Unsupported input file type: '{input_path}'. Only .json and .zip files are supported.")
