import os
import comfy.utils
import json
import logging
import folder_paths
from aiohttp import web
from typing import Any
from . import config


def resolve_model_base_paths():
    folders = list(folder_paths.folder_names_and_paths.keys())
    config.model_base_paths = {}
    for folder in folders:
        if folder == "configs":
            continue
        if folder == "custom_nodes":
            continue
        config.model_base_paths[folder] = folder_paths.get_folder_paths(folder)


def get_full_path(model_type: str, path_index: int, filename: str):
    """
    Get the absolute path in the model type through string concatenation.
    """
    folders = config.model_base_paths.get(model_type, [])
    if not path_index < len(folders):
        raise RuntimeError(f"PathIndex {path_index} is not in {model_type}")
    base_path = folders[path_index]
    return os.path.join(base_path, filename)


def get_valid_full_path(model_type: str, path_index: int, filename: str):
    """
    Like get_full_path but it will check whether the file is valid.
    """
    folders = config.model_base_paths.get(model_type, [])
    if not path_index < len(folders):
        raise RuntimeError(f"PathIndex {path_index} is not in {model_type}")
    base_path = folders[path_index]
    full_path = os.path.join(base_path, filename)
    if os.path.isfile(full_path):
        return full_path
    elif os.path.islink(full_path):
        raise RuntimeError(
            f"WARNING path {full_path} exists but doesn't link anywhere, skipping."
        )


def get_download_path():
    download_path = os.path.join(config.extension_uri, "downloads")
    if not os.path.exists(download_path):
        os.makedirs(download_path)
    return download_path


def recursive_search_files(directory: str):
    files, folder_all = folder_paths.recursive_search(
        directory, excluded_dir_names=[".git"]
    )
    return files


def search_files(directory: str):
    entries = os.listdir(directory)
    files = [f for f in entries if os.path.isfile(os.path.join(directory, f))]
    return files


def file_list_to_name_dict(files: list[str]):
    file_dict: dict[str, str] = {}
    for file in files:
        filename = os.path.splitext(file)[0]
        file_dict[filename] = file
    return file_dict


def get_model_metadata(filename: str):
    if not filename.endswith(".safetensors"):
        return {}
    try:
        out = comfy.utils.safetensors_header(filename, max_size=1024 * 1024)
        if out is None:
            return {}
        dt = json.loads(out)
        if not "__metadata__" in dt:
            return {}
        return dt["__metadata__"]
    except:
        return {}


def get_model_all_images(model_path: str):
    base_dirname = os.path.dirname(model_path)
    files = search_files(base_dirname)
    files = folder_paths.filter_files_content_types(files, ["image"])

    basename = os.path.splitext(os.path.basename(model_path))[0]
    output: list[str] = []
    for file in files:
        file_basename = os.path.splitext(file)[0]
        if file_basename == basename:
            output.append(file)
        if file_basename == f"{basename}.preview":
            output.append(file)
    return output


def get_model_preview_name(model_path: str):
    images = get_model_all_images(model_path)
    return images[0] if len(images) > 0 else "no-preview.png"


def save_model_preview_image(model_path: str, image_file: Any):
    if not isinstance(image_file, web.FileField):
        raise RuntimeError("Invalid image file")

    content_type: str = image_file.content_type
    if not content_type.startswith("image/"):
        raise RuntimeError(f"FileTypeError: expected image, got {content_type}")

    base_dirname = os.path.dirname(model_path)

    # remove old preview images
    old_preview_images = get_model_all_images(model_path)
    a1111_civitai_helper_image = False
    for image in old_preview_images:
        if os.path.splitext(image)[1].endswith(".preview"):
            a1111_civitai_helper_image = True
        image_path = os.path.join(base_dirname, image)
        os.remove(image_path)

    # save new preview image
    basename = os.path.splitext(os.path.basename(model_path))[0]
    extension = f".{content_type.split('/')[1]}"
    new_preview_path = os.path.join(base_dirname, f"{basename}{extension}")

    with open(new_preview_path, "wb") as f:
        f.write(image_file.file.read())

    # TODO Is it possible to abandon the current rules and adopt the rules of a1111 civitai_helper?
    if a1111_civitai_helper_image:
        """
        Keep preview image of a1111_civitai_helper
        """
        new_preview_path = os.path.join(base_dirname, f"{basename}.preview{extension}")
        with open(new_preview_path, "wb") as f:
            f.write(image_file.file.read())


def get_model_all_descriptions(model_path: str):
    base_dirname = os.path.dirname(model_path)
    files = search_files(base_dirname)
    files = folder_paths.filter_files_extensions(files, [".txt", ".md"])

    basename = os.path.splitext(os.path.basename(model_path))[0]
    output: list[str] = []
    for file in files:
        file_basename = os.path.splitext(file)[0]
        if file_basename == basename:
            output.append(file)
    return output


def get_model_description_name(model_path: str):
    descriptions = get_model_all_descriptions(model_path)
    basename = os.path.splitext(os.path.basename(model_path))[0]
    return descriptions[0] if len(descriptions) > 0 else f"{basename}.md"


def save_model_description(model_path: str, content: Any):
    if not isinstance(content, str):
        raise RuntimeError("Invalid description")

    base_dirname = os.path.dirname(model_path)

    # remove old descriptions
    old_descriptions = get_model_all_descriptions(model_path)
    for desc in old_descriptions:
        description_path = os.path.join(base_dirname, desc)
        os.remove(description_path)

    # save new description
    basename = os.path.splitext(os.path.basename(model_path))[0]
    extension = ".md"
    new_desc_path = os.path.join(base_dirname, f"{basename}{extension}")

    with open(new_desc_path, "w", encoding="utf-8") as f:
        f.write(content)


def rename_model(model_path: str, new_model_path: str):
    if model_path == new_model_path:
        return

    if os.path.exists(new_model_path):
        raise RuntimeError(f"Model {new_model_path} already exists")

    model_name = os.path.splitext(os.path.basename(model_path))[0]
    new_model_name = os.path.splitext(os.path.basename(new_model_path))[0]

    model_dirname = os.path.dirname(model_path)
    new_model_dirname = os.path.dirname(new_model_path)

    if not os.path.exists(new_model_dirname):
        os.makedirs(new_model_dirname)

    # move model
    os.rename(model_path, new_model_path)

    # move preview
    previews = get_model_all_images(model_path)
    for preview in previews:
        preview_path = os.path.join(model_dirname, preview)
        preview_name = os.path.splitext(preview)[0]
        preview_ext = os.path.splitext(preview)[1]
        new_preview_path = (
            os.path.join(new_model_dirname, new_model_name + preview_ext)
            if preview_name == model_name
            else os.path.join(
                new_model_dirname, new_model_name + ".preview" + preview_ext
            )
        )
        os.rename(preview_path, new_preview_path)

    # move description
    description = get_model_description_name(model_path)
    description_path = os.path.join(model_dirname, description)
    if os.path.isfile(description_path):
        new_description_path = os.path.join(new_model_dirname, f"{new_model_name}.md")
        os.rename(description_path, new_description_path)


import pickle


def save_dict_pickle_file(filename: str, data: dict):
    with open(filename, "wb") as f:
        pickle.dump(data, f)


def load_dict_pickle_file(filename: str) -> dict:
    with open(filename, "rb") as f:
        data = pickle.load(f)
    return data


def resolve_setting_key(key: str) -> str:
    key_paths = key.split(".")
    setting_id = config.setting_key
    try:
        for key_path in key_paths:
            setting_id = setting_id[key_path]
    except:
        pass
    if not isinstance(setting_id, str):
        raise RuntimeError(f"Invalid key: {key}")

    return setting_id


def set_setting_value(key: str, value: Any):
    setting_id = resolve_setting_key(key)
    fake_request = config.FakeRequest()
    settings = config.serverInstance.user_manager.settings.get_settings(fake_request)
    settings[setting_id] = value
    config.serverInstance.user_manager.settings.save_settings(fake_request, settings)


def get_setting_value(key: str, default: Any = None) -> Any:
    setting_id = resolve_setting_key(key)
    fake_request = config.FakeRequest()
    settings = config.serverInstance.user_manager.settings.get_settings(fake_request)
    return settings.get(setting_id, default)


from dataclasses import asdict, is_dataclass


def unpack_dataclass(data: Any):
    if isinstance(data, dict):
        return {key: unpack_dataclass(value) for key, value in data.items()}
    elif isinstance(data, list):
        return [unpack_dataclass(x) for x in data]
    elif is_dataclass(data):
        return asdict(data)
    else:
        return data
