import os
import hashlib
import functools
from typing import List
from pathlib import Path

import git

import folder_paths


@functools.lru_cache(maxsize=128)
def addnet_hash_safetensors(filepath: str) -> str:
    """Cached version of kohya-ss hash for safetensors"""
    hash_sha256 = hashlib.sha256()
    blksize = 1024 * 1024

    with open(filepath, "rb") as b:
        header = b.read(8)
        n = int.from_bytes(header, "little")

        offset = n + 8
        b.seek(offset)
        for chunk in iter(lambda: b.read(blksize), b""):
            hash_sha256.update(chunk)

    return hash_sha256.hexdigest()

class Lora:
    def __init__(self, name: str, unet: float, te: float, path: str):
        self.name = name
        self.unet = unet
        self.te = te
        self.path = path
        self.full_path = folder_paths.get_full_path("loras", self.path)

        if self.full_path:
            self.hash = addnet_hash_safetensors(self.full_path)[:12]

    def __str__(self):
        if self.unet != self.te:
            return f"<lora:{self.name}:unet={self.unet}:te={self.te}>"
        return f"<lora:{self.name}:{self.unet}>"

def extract_loras(workflow: dict) -> List[Lora]:
    result = []
    for node in workflow["nodes"]:
        if node["type"] == "Power Lora Loader (rgthree)":
            for widget in node.get("widgets_values", []):
                if isinstance(widget, dict) and widget.get("on"):
                    lora_path = widget["lora"]
                    name = os.path.splitext(os.path.basename(lora_path))[0]
                    unet = widget["strength"]
                    te = widget["strengthTwo"]
                    result.append(Lora(name=name, unet=unet, te=te, path=lora_path))
    return result


class AddA1111LikeMetadata:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "positive": ("STRING", {"forceInput": True}),
                "negative": ("STRING", {"forceInput": True}),
                "seed": ("INT", {"forceInput": True}),
                "model": ("STRING", {"forceInput": True}),
                "include_hashes": ("BOOLEAN", {"default": True}),
            },
            "hidden": {"extra_pnginfo": "EXTRA_PNGINFO"},
        }

    DESCRIPTION = "Adds A1111-like metadata to the PNG info."

    RETURN_TYPES = ()
    OUTPUT_NODE = True

    FUNCTION = "add_a1111_metadata"

    _version = None

    @classmethod
    def get_version(cls):
        if cls._version is None:
            try:
                repo = git.Repo(os.getcwd())
                version_str = repo.git.describe(tags=True).strip()
                if repo.active_branch.name != "master":
                    version_str += f" ({repo.active_branch.name})"
                cls._version = version_str
            except Exception:
                cls._version = "(Unknown version)"
        return cls._version

    def build_metadata(self, positive: str, negative: str, seed: int, include_hashes: bool, model: str, loras: List[Lora]):
        metadata = {
            "Negative prompt": negative,
            "Seed": seed,
            "Model": model,
        }

        if include_hashes and loras:
            metadata.update({
                "Lora hashes": f'"{", ".join([f"{lora.name}: {lora.hash}" for lora in loras])}"'
            })

        parts = [(positive + '\n' + ' '.join(str(lora) for lora in loras)).strip()]
        parts.extend(f"{key}: {value}" for key, value in metadata.items())
        parts.append(f"Version: ComfyUI {self.get_version()}")

        final_string = "\n".join([parts[0], parts[1], ", ".join(parts[2:])]).strip() + " [Full workflow included in the file]"
        return final_string


    def add_a1111_metadata(self, positive, negative, seed, model, include_hashes, extra_pnginfo=None):
        if extra_pnginfo is not None:
            extra_pnginfo["parameters"] = self.build_metadata(
                positive,
                negative,
                seed,
                include_hashes,
                Path(model).stem,
                extract_loras(extra_pnginfo['workflow']),
            )

        return {}

NODE_CLASS_MAPPINGS = {
    "AddA1111LikeMetadata": AddA1111LikeMetadata,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "AddA1111LikeMetadata": "Add A1111-like Metadata",
}