import sys
from types import ModuleType
from typing import Optional
from modules import scripts
from lib_free_u import global_state


def patch():
    xyz_module = find_xyz_module()
    if xyz_module is None:
        print("[sd-webui-freeu]", "xyz_grid.py not found.", file=sys.stderr)
        return
    xyz_module.axis_options.extend([
        xyz_module.AxisOption("[FreeU] Enabled", str_to_bool, apply_global_state("enable"), choices=choices_bool),
        xyz_module.AxisOption("[FreeU] Version", str, apply_global_state("version", key_map=global_state.all_versions), choices=choices_version),
        xyz_module.AxisOption("[FreeU] Preset", str, apply_global_state("preset"), choices=choices_preset),
        xyz_module.AxisOption("[FreeU] Start At Step", int_or_float, apply_global_state("start_ratio")),
        xyz_module.AxisOption("[FreeU] Stop At Step", int_or_float, apply_global_state("stop_ratio")),
        xyz_module.AxisOption("[FreeU] Transition Smoothness", int_or_float, apply_global_state("transition_smoothness")),
        *[
            opt
            for index in range(global_state.STAGES_COUNT)
            for opt in [
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Backbone Scale", float, apply_global_state(f"b{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Backbone Offset", float, apply_global_state(f"o{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Backbone Width", float, apply_global_state(f"w{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Skip Scale", float, apply_global_state(f"s{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Skip Cutoff", float, apply_global_state(f"t{index}")),
                xyz_module.AxisOption(f"[FreeU] Stage {index+1} Skip High End Scale", float, apply_global_state(f"h{index}")),
            ]
        ]
    ])


def apply_global_state(k, key_map=None):
    def callback(_p, v, _vs):
        if key_map is not None:
            v = key_map[v]
        global_state.xyz_attrs[k] = v

    return callback


def str_to_bool(string):
    string = str(string)
    if string in ["None", ""]:
        return None
    elif string.lower() in ["true", "1"]:
        return True
    elif string.lower() in ["false", "0"]:
        return False
    else:
        raise ValueError(f"Could not convert string to boolean: {string}")


def int_or_float(string):
    try:
        return int(string)
    except ValueError:
        return float(string)


def choices_bool():
    return ["False", "True"]


def choices_version():
    return list(global_state.all_versions.keys())


def choices_preset():
    presets = list(global_state.all_presets.keys())
    presets.insert(0, "UI Settings")
    return presets
    

def find_xyz_module() -> Optional[ModuleType]:
    for data in scripts.scripts_data:
        if data.script_class.__module__ in {"xyz_grid.py", "xy_grid.py"} and hasattr(data, "module"):
            return data.module

    return None