from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

from dynamicprompts.wildcards import WildcardManager
from modules import script_callbacks
from modules.generation_parameters_copypaste import parse_generation_parameters
from modules.script_callbacks import ImageSaveParams

from sd_dynamic_prompts.pnginfo_saver import strip_template_info
from sd_dynamic_prompts.prompt_writer import PromptWriter
from sd_dynamic_prompts.settings import on_ui_settings
from sd_dynamic_prompts.wildcards_tab import initialize as initialize_wildcards_tab

logger = logging.getLogger(__name__)


def register_prompt_writer(prompt_writer: PromptWriter) -> None:
    def on_save(image_save_params: ImageSaveParams) -> None:
        image_name = Path(image_save_params.filename)
        prompt_filename = image_name.with_suffix(".csv")
        prompt_writer.write_prompts(prompt_filename)

    script_callbacks.on_before_image_saved(on_save)


def register_on_infotext_pasted() -> None:
    def on_infotext_pasted(infotext: str, parameters: dict[str, Any]) -> None:
        new_parameters = {}
        if "Prompt" in parameters and "Template:" in parameters["Prompt"]:
            strip_template_info(parameters)
            new_parameters = parse_generation_parameters(parameters["Prompt"])
        elif (
            "Negative prompt" in parameters
            and "Template:" in parameters["Negative prompt"]
        ):
            strip_template_info(parameters)
            new_parameters = parse_generation_parameters(parameters["Negative prompt"])
            new_parameters["Negative prompt"] = new_parameters["Prompt"]
            new_parameters["Prompt"] = parameters["Prompt"]
        parameters.update(new_parameters)

    script_callbacks.on_infotext_pasted(on_infotext_pasted)


def register_settings():
    script_callbacks.on_ui_settings(on_ui_settings)


def register_wildcards_tab(wildcard_manager: WildcardManager):
    initialize_wildcards_tab(wildcard_manager)
