from __future__ import annotations

import json
import logging
import random
import shutil
import traceback
from pathlib import Path

import gradio as gr
from dynamicprompts.wildcards import WildcardManager
from dynamicprompts.wildcards.collection import WildcardTextFile
from dynamicprompts.wildcards.tree import WildcardTreeNode
from modules import script_callbacks
from send2trash import send2trash

from sd_dynamic_prompts.element_ids import make_element_id

COPY_COLLECTION_ACTION = "copy collection"
LOAD_FILE_ACTION = "load file"
LOAD_TREE_ACTION = "load tree"
MESSAGE_PROCESSING_ACTION = "message processing"

logger = logging.getLogger(__name__)

wildcard_manager: WildcardManager


def get_collection_dirs() -> dict[str, Path]:
    """
    Get a mapping of name -> subdirectory path for the extension's collections/ directory.
    """
    from sd_dynamic_prompts.paths import get_extension_base_path

    collections_path = get_extension_base_path() / "collections"

    return {
        str(pth.relative_to(collections_path)): pth
        for pth in collections_path.iterdir()
        if pth.is_dir()
    }


def initialize(manager: WildcardManager):
    global wildcard_manager
    wildcard_manager = manager
    script_callbacks.on_ui_tabs(on_ui_tabs)


def _format_node_for_json(
    wildcard_manager: WildcardManager,
    node: WildcardTreeNode,
) -> list[dict]:
    collections = [
        {
            "name": node.qualify_name(coll),
            "wrappedName": wildcard_manager.to_wildcard(node.qualify_name(coll)),
            "children": [],
        }
        for coll in sorted(node.collections)
    ]
    child_items = [
        {"name": name, "children": _format_node_for_json(wildcard_manager, child_node)}
        for name, child_node in sorted(node.child_nodes.items())
    ]
    return [*collections, *child_items]


def on_ui_tabs():
    help_html = f"""
    <ol>
        <li>Create your wildcard library by copying a collection using the dropdown in the Collection actions.</li>
        <li>Click on the files that appear in the tree to edit them.</li>
        <li>Use the wildcard in your script by typing the name of the file or copying the text from the Wildcards file text box.</li>
        <li>You can also add your own wildcard files to the {wildcard_manager.path} folder.</li>
    </ol>
    """

    with gr.Blocks() as wildcards_tab:
        with gr.Row():
            with gr.Column():
                gr.Textbox(
                    placeholder="Search in wildcard names...",
                    elem_id=make_element_id("wildcard-search"),
                    label="",
                )
                gr.HTML("Loading...", elem_id=make_element_id("wildcard-tree"))
                with gr.Accordion("Help", open=False):
                    gr.HTML(help_html)
                with gr.Accordion("Collection actions", open=False):
                    collection_dropdown = gr.Dropdown(
                        choices=sorted(get_collection_dirs()),
                        type="value",
                        label="Select a collection",
                    )
                    with gr.Row():
                        collection_copy_button = gr.Button(
                            "Copy collection",
                            scale=1,
                        )
                        overwrite_checkbox = gr.Checkbox(
                            label="Overwrite existing",
                            value=False,
                        )
                    with gr.Row():
                        refresh_wildcards_button = gr.Button(
                            "Refresh wildcards",
                            elem_id=make_element_id("wildcard-load-tree-button"),
                        )
                        delete_tree_button = gr.Button(
                            "Delete all wildcards",
                            elem_id=make_element_id("wildcard-delete-tree-button"),
                        )
            with gr.Column():
                gr.Textbox(
                    "",
                    elem_id=make_element_id("wildcard-file-name"),
                    interactive=False,
                    label="Wildcards file",
                )
                gr.Textbox(
                    "",
                    elem_id=make_element_id("wildcard-file-editor"),
                    lines=10,
                    interactive=True,
                    label="File editor",
                )
                save_button = gr.Button(
                    "Save wildcards",
                    scale=1,
                    elem_id=make_element_id("wildcard-save-button"),
                )

        # Hidden scratch textboxes and button for communication with JS bits.
        client_to_server_message_textbox = gr.Textbox(
            "",
            elem_id=make_element_id("wildcard-c2s-message-textbox"),
            visible=False,
        )
        server_to_client_message_textbox = gr.Textbox(
            "",
            elem_id=make_element_id("wildcard-s2c-message-textbox"),
            visible=False,
        )
        client_to_server_message_action_button = gr.Button(
            "Action",
            elem_id=make_element_id("wildcard-c2s-action-button"),
            visible=False,
        )

        # Handle the frontend sending a message
        client_to_server_message_action_button.click(
            handle_message,
            inputs=[client_to_server_message_textbox],
            outputs=[server_to_client_message_textbox],
        )

        refresh_wildcards_button.click(
            refresh_wildcards_callback,
            inputs=[],
            outputs=[server_to_client_message_textbox],
        )

        delete_tree_button.click(
            delete_tree_callback,
            _js="SDDP.onDeleteTreeClick",
            inputs=[client_to_server_message_textbox],
            outputs=[server_to_client_message_textbox],
        )

        save_button.click(
            save_file_callback,
            _js="SDDP.onSaveFileClick",
            inputs=[client_to_server_message_textbox],
            outputs=[server_to_client_message_textbox],
        )

        collection_copy_button.click(
            copy_collection_callback,
            inputs=[overwrite_checkbox, collection_dropdown],
            outputs=[server_to_client_message_textbox],
        )

    return ((wildcards_tab, "Wildcards Manager", "sddp-wildcard-manager"),)


def create_payload(*, action: str, success: bool, **rest) -> str:
    return json.dumps(
        {
            "id": random.randint(0, 1000000),
            "action": action,
            "success": success,
            **rest,
        },
    )


def copy_collection_callback(overwrite_checkbox, collection):
    collection_paths = get_collection_dirs()
    if collection in collection_paths:
        collection_path = collection_paths[collection]
        for file in collection_path.glob("**/*"):
            if file.is_file():
                target_path = (
                    wildcard_manager.path
                    / collection
                    / file.relative_to(collection_path)
                )
                if not target_path.exists() or overwrite_checkbox:
                    target_path.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy(file, target_path)

        return refresh_wildcards_callback()

    return create_payload(
        action=COPY_COLLECTION_ACTION,
        success=False,
    )


def refresh_wildcards_callback():
    wildcard_manager.clear_cache()
    root = wildcard_manager.tree.root
    tree = _format_node_for_json(wildcard_manager, root)
    collection_count = len(list(root.walk_full_names()))
    return create_payload(
        action=LOAD_TREE_ACTION,
        success=True,
        tree=tree,
        collection_count=collection_count,
    )


def delete_tree_callback(confirm_delete):
    if confirm_delete == "True":
        send2trash(wildcard_manager.path)
        wildcard_manager.path.mkdir(parents=True, exist_ok=True)
        return refresh_wildcards_callback()
    return create_payload(action="delete tree", success=False)


def handle_message(event_str: str) -> str:
    try:
        event = json.loads(event_str)
        if event["action"] == LOAD_FILE_ACTION:
            return handle_load_wildcard(event)
        raise ValueError(f"Unknown event: {event}")
    except Exception as e:
        traceback.print_exc()
        return create_payload(
            action=MESSAGE_PROCESSING_ACTION,
            success=False,
            message=f"Error processing message: {e}",
        )


def handle_load_wildcard(event: dict) -> str:
    name = event["name"]
    wf = wildcard_manager.get_file(name)
    if isinstance(wf, WildcardTextFile):
        # For text files, just return the raw text.
        contents = wf.read_text()
        can_edit = True
    else:
        # Otherwise, return a preview of the values,
        # with a header to indicate that the file can't be edited.
        values = "\n".join(str(val) for val in wf.get_values())
        contents = f"# File can't be edited\n{values}"
        can_edit = False

    return create_payload(
        action=LOAD_FILE_ACTION,
        success=True,
        contents=contents,
        can_edit=can_edit,
        name=name,
        wrapped_name=wildcard_manager.to_wildcard(name),
    )


def save_file_callback(event_str: str):
    try:
        event = json.loads(event_str)
        wf = wildcard_manager.get_file(event["wildcard"]["name"])
        if isinstance(wf, WildcardTextFile):
            wf.write_text(event["contents"].strip())
        else:
            raise Exception("Can't save non-text files")
        wildcard_manager.clear_cache()
        return handle_load_wildcard({"name": event["wildcard"]["name"]})
    except Exception as e:
        logger.exception(e)
