import concurrent.futures
import json
import multiprocessing
import os
import re
import shutil
import subprocess
import time
from ibm_watsonx_data_integration.services.datastage.codegen.python_generator import *

global lock
lock = multiprocessing.Lock()


def _sort_keys(d):
    if not isinstance(d, dict):
        return d
    return {k: _sort_keys(d[k]) for k in sorted(d.keys(), key=lambda x: str(x))}


###############################################################################


def _filter_ignored(errors):
    errors_filtered = dict()

    for p, err in errors.items():
        if "SuperNodes are not supported" in err:
            continue  # Ignore subflows
        elif bool(re.search(r"ValueError: .*? is not a valid .*?", err)):
            continue  # Ignore parameter/value sets
        else:
            errors_filtered[p] = err

    return errors_filtered


# Internal helper that runs PythonGenerator on one zip
def _generate_for_one_zip(i, filename, input_dir, output_dir, generator: PythonGenerator, errors):
    if filename.endswith(".zip"):
        folder = filename.replace(".zip", "")
        print(f"({i}) Regenerating folder {folder}")
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, folder)
        _, errors_for_zip = generator.generate(input_path=input_path, output_path=output_path)

        # Filter out errors that aren't important
        errors_filtered = _filter_ignored(errors_for_zip)
        if errors_filtered:
            errors[folder] = errors_filtered

    return errors


# Internal helper that runs PythonGenerator on one json
def _generate_for_one_json(i, filename, input_dir, output_dir, generator: PythonGenerator, errors):
    if filename.endswith(".json"):
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, filename.replace(".json", ".py"))
        _, errors_for_json = generator.generate(input_path=input_path, output_path=output_path)

        # Filter out errors that aren't important
        errors_filtered = _filter_ignored(errors_for_json)
        if errors_filtered:
            errors[filename.replace(".json", ".py")] = errors_filtered

    return errors


# Run PythonGenerator on all files in the input_dir
def generate(file_type, input_dir, output_dir, results_dir, generator):
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(results_dir, exist_ok=True)
    helper = _generate_for_one_zip if file_type == "zip" else _generate_for_one_json

    start_time = time.time()

    manager = multiprocessing.Manager()
    errors = manager.dict()
    with concurrent.futures.ProcessPoolExecutor(max_workers=12) as executor:
        filenames = [filename for filename in os.listdir(input_dir)]
        futures = [
            executor.submit(helper, i, filename, input_dir, output_dir, generator, errors)
            for i, filename in enumerate(filenames)
        ]
        for future in futures:
            future.result()

    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Elapsed time: {elapsed_time} seconds")

    results_path = os.path.join(results_dir, f"generate_errors_{file_type}.json")
    with open(results_path, "w") as f:
        json.dump(_sort_keys(dict(errors)), f, indent=4)
    return _sort_keys(dict(errors))


###############################################################################


# Helper to execute one code file at 'code_path' and return the prints/errors
def _execute_code(id, code_path, temp_dir, prints, errors, successes):
    global lock
    os.makedirs(temp_dir, exist_ok=True)

    # Modifications: replace placeholder api key and project id, then save temporarily
    with open(code_path) as f:
        original = f.read()
    modified = original.replace(
        "# This code was autogenerated by PythonGenerator",
        "# This code was autogenerated by PythonGenerator\nimport os",
    )
    modified = modified.replace('api_key = "<your_api_key>"', 'api_key = os.environ["WXDI_API_KEY"]')
    modified = modified.replace('guid="<your_project_id>"', 'guid=os.environ["WXDI_PROJECT_ID"]')
    file_name = f"{id}_{Path(code_path).stem}"
    temp_code_path = os.path.join(temp_dir, f"temp_{file_name}.py")
    os.makedirs(os.path.dirname(temp_code_path), exist_ok=True)
    with open(temp_code_path, "w") as temp_file:
        temp_file.write(modified)

    # Run it and return any print outputs and errors
    timeout = 50
    result_print, result_error = [], []
    try:
        result = subprocess.run(["python", temp_code_path], capture_output=True, text=True, timeout=timeout)
        os.remove(temp_code_path)
        result_print = result.stdout
        result_error = result.stderr
    except subprocess.TimeoutExpired:
        result_print = ""
        result_error = "timed out"

    # Process results
    with lock:
        if not result_print and not result_error:
            successes.append(file_name)
            print(id, "success")
        if result_print:
            prints[file_name] = result_print
            print(id, "print")
        if result_error:
            errors[file_name] = result_error
            print(id, "error")


# Saves a checkpoint of prints, errors, and successes
def save_checkpoint(i, prints, errors, successes, result_dir):
    print(f"Saving checkpoint {i}: {len(prints)} Print | {len(errors)} Error | {len(successes)} Success")

    with open(os.path.join(result_dir, "run_print.json"), "w") as f:
        json.dump(_sort_keys(dict(prints)), f, indent=4)
    with open(os.path.join(result_dir, "run_error.json"), "w") as f:
        json.dump(_sort_keys(dict(errors)), f, indent=4)
    with open(os.path.join(result_dir, "run_success.json"), "w") as f:
        json.dump(sorted(successes), f, indent=4)


# Run generated code
def execute_code(code_dir, results_dir, temp_dir):
    manager = multiprocessing.Manager()
    errors = manager.dict()
    prints = manager.dict()
    successes = manager.list()

    # Collect valid code files to run
    valid_files = []
    for root, _, files in os.walk(code_dir):
        for file in files:
            full_path = os.path.join(root, file)
            if full_path.endswith(".py"):
                with open(full_path, encoding="utf-8") as f:
                    first_line = f.readline().strip()
                    if first_line == "# This code was autogenerated by PythonGenerator":
                        valid_files.append(full_path)

    # Run code files in parallel
    start_time = time.time()
    with concurrent.futures.ProcessPoolExecutor(max_workers=12) as executor:
        futures = []
        for i, code_path in enumerate(valid_files):
            futures.append(executor.submit(_execute_code, i, code_path, temp_dir, prints, errors, successes))
            # Checkpoint
            if (i + 1) % 50 == 0:
                concurrent.futures.wait(futures)
                save_checkpoint(i, prints, errors, successes, results_dir)
                futures = []

        # Final Checkpoint
        concurrent.futures.wait(futures)
        save_checkpoint(i, prints, errors, successes, results_dir)

    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"Elapsed time: {elapsed_time} seconds")


###############################################################################

if __name__ == "__main__":
    test_dir = os.path.join(os.getcwd(), "ibm_watsonx_data_integration/services/datastage/tests/pythongenerator")
    results_dir = os.path.join(test_dir, "results")
    temp_dir = os.path.join(test_dir, "temp")
    os.makedirs(temp_dir, exist_ok=True)

    # Choose PythonGenerator settings here - TODO: need to test both modes next
    generator = PythonGenerator()
    generator.configuration.mode = "file_per_flow"  # or single_file
    generator.configuration.api_key = 'os.environ["WXDI_API_KEY"]'
    generator.configuration.project_id = 'os.environ["WXDI_PROJECT_ID"]'
    generator.configuration.create_job = False
    generator.configuration.run_job = False

    # Run PythonGenerator on zip format
    # zip_dir = os.path.join(test_dir, "flows/zip")
    # code_dir = os.path.join(test_dir, "code/zip")
    # generate("zip", zip_dir, code_dir, results_dir, generator)

    # Run PythonGenerator on one zip - debug only
    # zip_name = "001"
    # shutil.copy2(os.path.join(test_dir, f"flows/zip/{zip_name}.zip"), temp_dir)
    # generate("zip", temp_dir, os.path.join(test_dir, "code/zip"), results_dir, generator)
    # os.remove(os.path.join(temp_dir, zip_name + ".zip"))

    # Run conversion for connector tests - zip
    shutil.copy2(os.path.join(test_dir, "flows/zip/connector_tests.zip"), temp_dir)
    generate("zip", temp_dir, os.path.join(test_dir, "code/zip"), results_dir, generator)
    os.remove(os.path.join(temp_dir, "connector_tests.zip"))

    # Run conversion for connector tests - json
    # input_dir = os.path.join(test_dir, "flows/json/connector")
    # output_dir = os.path.join(test_dir, "code/json/connector")
    # generate("json", input_dir, output_dir, results_dir, generator)

    # Run PythonGnenerator on json format of customer zips
    # all_errors = dict()
    # for i in range(1, 562):
    #     json_dir = os.path.join(test_dir, "flows/json", f"{i:03}")
    #     code_dir = os.path.join(test_dir, "code/json", f"{i:03}")
    #     curr_errors = generate("json", json_dir, code_dir, results_dir, generator)
    #     if curr_errors: all_errors[f"{i:03}"] = curr_errors
    #     if i > 0 and i % 20 == 0:
    #         print(f"Saving checkpoint {i} with {len(curr_errors.keys())} folders with errors")
    #         with open(os.path.join(results_dir, "generate_errors_json.json"), "w") as f:
    #             json.dump(all_errors, f, indent=4)

    # Run generated code - make sure to set WXDI_API_KEY and WXDI_PROJECT_ID
    # code_dir = os.path.join(test_dir, "code/custom")
    # execute_code(code_dir, results_dir, temp_dir)
