import hashlib
import re
import secrets
import threading
import time

from aiohttp import web

import folder_paths
from server import PromptServer

_RNG = secrets.SystemRandom()
_STATES = {}
_STATE_LOCK = threading.Lock()


def _get_checkpoints():
    """
    Use ComfyUI's native checkpoint list order.

    Important:
      CheckpointLoaderSimple validates combo/list types against this list.
      If this node sorts the list independently, the list type can differ from
      CheckpointLoaderSimple on first run after loading a workflow, causing:
        Return type mismatch between linked nodes
    """
    try:
        return list(folder_paths.get_filename_list("checkpoints"))
    except Exception:
        return []


def _checkpoint_type():
    checkpoints = _get_checkpoints()
    return checkpoints if checkpoints else [""]


def _checkpoint_hash(checkpoints):
    joined = "\n".join(str(item) for item in checkpoints)
    return hashlib.sha1(joined.encode("utf-8")).hexdigest()


def _clamp_index(index, count):
    if count <= 0:
        return 0
    return max(0, min(int(index), count - 1))


def _index_of_checkpoint(checkpoints, checkpoint_name):
    if not checkpoints:
        return 0
    try:
        return checkpoints.index(checkpoint_name)
    except ValueError:
        return 0


def _make_shuffle_order(count, first_index=None, previous_index=None):
    if count <= 0:
        return []

    indices = list(range(count))

    if first_index is not None:
        first_index = _clamp_index(first_index, count)
        rest = [index for index in indices if index != first_index]
        _RNG.shuffle(rest)
        return [first_index] + rest

    _RNG.shuffle(indices)

    if count > 1 and previous_index is not None and indices and indices[0] == previous_index:
        swap_index = _RNG.randrange(1, count)
        indices[0], indices[swap_index] = indices[swap_index], indices[0]

    return indices


def _safe_checkpoint_name(name):
    value = str(name).replace("\\", "/")
    value = value.replace("/", "_")

    for ext in [".safetensors", ".ckpt", ".pt", ".pth", ".bin"]:
        if value.lower().endswith(ext):
            value = value[:-len(ext)]
            break

    value = re.sub(r'[<>:"|?*\x00-\x1f]', "_", value)
    value = value.strip().strip(".")
    return value or "checkpoint"


def _new_state(checkpoints, start_checkpoint, mode, change_every):
    count = len(checkpoints)
    start_index = _index_of_checkpoint(checkpoints, start_checkpoint)

    state = {
        "start_checkpoint": start_checkpoint,
        "start_index": start_index,
        "mode": mode,
        "change_every": max(1, int(change_every)),
        "repeat_count": 0,
        "current_index": start_index,
        "shuffle_order": [],
        "shuffle_position": 0,
        "cycle_count": 0,
        "checkpoint_hash": _checkpoint_hash(checkpoints),
    }

    if count > 0 and mode == "shuffle_once":
        state["shuffle_order"] = _make_shuffle_order(count, first_index=start_index)
        state["shuffle_position"] = 0
        state["current_index"] = state["shuffle_order"][0]

    return state


def _state_needs_reset(state, checkpoints, start_checkpoint, mode, change_every):
    if state is None:
        return True

    if state.get("checkpoint_hash") != _checkpoint_hash(checkpoints):
        return True

    if state.get("start_checkpoint") != start_checkpoint:
        return True

    if state.get("mode") != mode:
        return True

    if state.get("change_every") != max(1, int(change_every)):
        return True

    return False


def _select_index(state, checkpoints, mode):
    count = len(checkpoints)
    if count <= 0:
        return 0

    if mode == "shuffle_once":
        order = state.get("shuffle_order") or list(range(count))
        position = state.get("shuffle_position", 0)

        if position >= len(order):
            position = 0
            state["shuffle_position"] = 0

        return _clamp_index(order[position], count)

    return _clamp_index(state.get("current_index", 0), count)


def _advance_state(state, checkpoints, mode):
    count = len(checkpoints)
    if count <= 0:
        return

    state["repeat_count"] = int(state.get("repeat_count", 0)) + 1

    if state["repeat_count"] < state["change_every"]:
        return

    state["repeat_count"] = 0

    if mode == "fixed":
        return

    if mode == "increment":
        current = _clamp_index(state.get("current_index", 0), count)
        next_index = current + 1
        if next_index >= count:
            next_index = 0
            state["cycle_count"] = int(state.get("cycle_count", 0)) + 1
        state["current_index"] = next_index
        return

    if mode == "randomize":
        state["current_index"] = _RNG.randrange(count)
        return

    if mode == "shuffle_once":
        current_index = _select_index(state, checkpoints, mode)
        position = int(state.get("shuffle_position", 0)) + 1

        if position >= count:
            state["cycle_count"] = int(state.get("cycle_count", 0)) + 1
            state["shuffle_order"] = _make_shuffle_order(count, previous_index=current_index)
            state["shuffle_position"] = 0
        else:
            state["shuffle_position"] = position

        state["current_index"] = _select_index(state, checkpoints, mode)


routes = PromptServer.instance.routes


@routes.post("/checkpoint_name_cycler/reset")
async def checkpoint_name_cycler_reset(_request):
    with _STATE_LOCK:
        count = len(_STATES)
        _STATES.clear()

    return web.json_response({
        "ok": True,
        "reset_count": count,
    })


class CheckpointNameCycler:
    @classmethod
    def INPUT_TYPES(cls):
        checkpoints = _checkpoint_type()

        return {
            "required": {
                "start_checkpoint": (checkpoints,),
                "mode": (["fixed", "increment", "randomize", "shuffle_once"], {"default": "increment"}),
                "change_every": ("INT", {"default": 1, "min": 1, "max": 999999}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = (_checkpoint_type(), "STRING", "STRING", "INT", "INT", "INT")
    RETURN_NAMES = ("ckpt_name", "ckpt_name_str", "ckpt_name_safe", "index", "count", "cycle")
    FUNCTION = "cycle"
    CATEGORY = "utils/model"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return time.time_ns()

    def cycle(self, start_checkpoint, mode, change_every, unique_id=None):
        checkpoints = _get_checkpoints()
        count = len(checkpoints)
        node_key = str(unique_id) if unique_id is not None else "__default__"

        with _STATE_LOCK:
            state = _STATES.get(node_key)

            if _state_needs_reset(state, checkpoints, start_checkpoint, mode, change_every):
                state = _new_state(checkpoints, start_checkpoint, mode, change_every)
                _STATES[node_key] = state

            change_every = max(1, int(change_every))

            if count <= 0:
                return {
                    "ui": {
                        "ckpt_name": [""],
                        "repeat_index": [0],
                        "change_every": [change_every],
                    },
                    "result": ("", "", "checkpoint", 0, 0, int(state.get("cycle_count", 0))),
                }

            index = _select_index(state, checkpoints, mode)
            ckpt_name = checkpoints[index]
            ckpt_name_str = str(ckpt_name)
            ckpt_name_safe = _safe_checkpoint_name(ckpt_name_str)
            cycle = int(state.get("cycle_count", 0))
            repeat_index = int(state.get("repeat_count", 0)) + 1

            _advance_state(state, checkpoints, mode)

            return {
                "ui": {
                    "ckpt_name": [ckpt_name_safe],
                    "repeat_index": [repeat_index],
                    "change_every": [change_every],
                },
                "result": (ckpt_name, ckpt_name_str, ckpt_name_safe, int(index), int(count), cycle),
            }


NODE_CLASS_MAPPINGS = {
    "CheckpointNameCycler": CheckpointNameCycler,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CheckpointNameCycler": "Checkpoint Name Cycler",
}