# -*- coding: utf-8 -*-
"""
디스크 기반 노드 출력 캐시 (ComfyUI cache provider API).

ComfyUI는 노드 출력을 "입력 서명"(위젯 값 + 모든 조상 노드의 서명, 세션 간
결정적 SHA256)으로 캐싱하는데, 그 캐시는 메모리에만 있어서 (1) 재시작하면
사라지고 (2) 기본 RAM_PRESSURE 모드에서는 RAM이 부족하면 수시로 증발한다.

이 프로바이더는 캐시 저장 시 디스크에도 쓰고, 메모리 캐시 미스 시 디스크에서
복원한다. 효과:
  - 재시작해도 시드/프롬프트/이미지가 같은 클립은 JoyCaption부터 샘플러까지
    전부 스킵 (실행되지 않고 디스크에서 결과만 복원)
  - 시드 하나만 바꾸면 정확히 그 클립 체인만 다시 생성 (서명에 모든 입력이
    포함되므로 무효화 판정은 ComfyUI가 알아서 정확하게 해줌)
  - MODEL/CLIP/VAE/NOISE 같은 객체 출력은 직렬화 불가로 자동 제외되어 저장 안 됨

캐시 위치/용량은 환경변수로 조정:
  DOLPHIN_DISK_CACHE_DIR      (기본: <output>/_node_disk_cache)
  DOLPHIN_DISK_CACHE_MAX_GB   (기본: 40  — 초과 시 오래된 것부터 삭제)
  DOLPHIN_DISK_CACHE_MAX_ENTRY_GB (기본: 2 — 이보다 큰 단일 출력은 저장 안 함)
"""
import os
import time
import asyncio
import logging

import torch

import folder_paths
from comfy_execution.cache_provider import register_cache_provider
from comfy_api.latest._caching import CacheProvider, CacheValue

log = logging.getLogger("dolphin.diskcache")

_DEFAULT_DIR = os.path.join(folder_paths.get_output_directory(), "_node_disk_cache")
CACHE_DIR = os.environ.get("DOLPHIN_DISK_CACHE_DIR", _DEFAULT_DIR)
MAX_TOTAL_BYTES = int(float(os.environ.get("DOLPHIN_DISK_CACHE_MAX_GB", "15")) * (1024 ** 3))
# VAEDecode/이미지 배치 결과(수백MB~1GB, 디코드된 프레임)는 일부러 상한 아래로 뒀다.
# 비싼 건 디퓨전 샘플링(latent, 수십MB)이고 VAE 디코드는 latent만 있으면 몇 초면
# 다시 되므로, 큰 디코드 결과까지 디스크에 쌓아둘 실익이 없다 - 용량만 먹는다.
MAX_ENTRY_BYTES = int(float(os.environ.get("DOLPHIN_DISK_CACHE_MAX_ENTRY_GB", "0.2")) * (1024 ** 3))

_ALLOWED_SCALARS = (str, int, float, bool, bytes, type(None))


def _serializable(obj):
    """텐서/스칼라/컨테이너만 허용. MODEL, NOISE, SAMPLER 등 객체가 섞이면 False."""
    if isinstance(obj, torch.Tensor):
        return True
    if isinstance(obj, _ALLOWED_SCALARS):
        return True
    if isinstance(obj, (list, tuple)):
        return all(_serializable(x) for x in obj)
    if isinstance(obj, dict):
        return all(isinstance(k, _ALLOWED_SCALARS) and _serializable(v) for k, v in obj.items())
    return False


def _to_cpu(obj):
    if isinstance(obj, torch.Tensor):
        return obj.detach().to("cpu")
    if isinstance(obj, list):
        return [_to_cpu(x) for x in obj]
    if isinstance(obj, tuple):
        return tuple(_to_cpu(x) for x in obj)
    if isinstance(obj, dict):
        return {k: _to_cpu(v) for k, v in obj.items()}
    return obj


def _tensor_bytes(obj):
    if isinstance(obj, torch.Tensor):
        return obj.numel() * obj.element_size()
    if isinstance(obj, dict):
        return sum(_tensor_bytes(v) for v in obj.values())
    if isinstance(obj, (list, tuple)):
        return sum(_tensor_bytes(v) for v in obj)
    return 0


class DolphinDiskCache(CacheProvider):
    def __init__(self, directory=CACHE_DIR):
        self.dir = directory
        os.makedirs(self.dir, exist_ok=True)
        self.log_path = os.path.join(self.dir, "_diskcache.log")

    def _flog(self, msg):
        # 콘솔이 안 보이는 환경에서도 저장/스킵 경로를 추적할 수 있게 파일로 남긴다.
        # 주의: should_cache가 아예 호출 안 된 노드는 ComfyUI 쪽 NaN 키 게이트
        # (_contains_self_unequal)에서 걸러진 것 - 로그에 안 찍히는 것 자체가 단서.
        try:
            with open(self.log_path, "a", encoding="utf-8") as f:
                f.write(f"{time.strftime('%H:%M:%S')} {msg}\n")
        except OSError:
            pass

    def _path(self, context):
        return os.path.join(self.dir, f"{context.cache_key_hash}.pt")

    async def on_lookup(self, context):
        path = self._path(context)
        if not os.path.isfile(path):
            return None

        def load():
            return torch.load(path, map_location="cpu", weights_only=False)

        try:
            data = await asyncio.to_thread(load)
        except Exception as e:
            log.warning(f"[DiskCache] 손상된 캐시 삭제 ({context.class_type}): {e}")
            try:
                os.remove(path)
            except OSError:
                pass
            return None
        try:
            os.utime(path, None)  # LRU: 접근 시각 갱신 (prune 시 오래된 것부터 삭제)
        except OSError:
            pass
        print(f"💾 [DiskCache] {context.class_type} 디스크에서 복원 (node {context.node_id})")
        self._flog(f"RESTORE {context.class_type} (node {context.node_id})")
        return CacheValue(outputs=data["outputs"], ui=data.get("ui"))

    def should_cache(self, context, value=None):
        if value is None:  # lookup 시점 - 파일 존재 여부로 판단하므로 항상 시도
            return True
        if not _serializable(value.outputs):
            self._flog(f"SKIP not-serializable {context.class_type} (node {context.node_id})")
            return False
        size = _tensor_bytes(value.outputs)
        if size > MAX_ENTRY_BYTES:
            self._flog(f"SKIP too-big {context.class_type} {size/1024/1024:.0f}MB (node {context.node_id})")
            return False
        self._flog(f"STORE-OK {context.class_type} {size/1024/1024:.1f}MB (node {context.node_id})")
        return True

    async def on_store(self, context, value):
        path = self._path(context)
        if os.path.exists(path):
            return
        try:
            payload = {
                "outputs": _to_cpu(value.outputs),
                "ui": value.ui,
                "class_type": context.class_type,
                "saved_at": time.time(),
            }

            def save():
                tmp = path + ".tmp"
                torch.save(payload, tmp)
                os.replace(tmp, path)

            await asyncio.to_thread(save)
            self._flog(f"SAVED {context.class_type} (node {context.node_id})")
        except Exception as e:
            log.warning(f"[DiskCache] 저장 실패 ({context.class_type}): {e}")
            self._flog(f"SAVE-FAIL {context.class_type} (node {context.node_id}): {type(e).__name__}: {e}")
            try:
                os.remove(path + ".tmp")
            except OSError:
                pass

    def on_prompt_end(self, prompt_id):
        try:
            self._prune()
        except Exception as e:
            log.warning(f"[DiskCache] prune 실패: {e}")

    def _prune(self):
        entries = []
        total = 0
        with os.scandir(self.dir) as it:
            for e in it:
                if e.name.endswith(".pt") and e.is_file():
                    st = e.stat()
                    entries.append((st.st_mtime, st.st_size, e.path))
                    total += st.st_size
        if total <= MAX_TOTAL_BYTES:
            return
        entries.sort()  # mtime 오래된 순
        for _, size, path in entries:
            try:
                os.remove(path)
                total -= size
            except OSError:
                pass
            if total <= MAX_TOTAL_BYTES:
                break


_provider = None


def register():
    global _provider
    if _provider is None:
        _provider = DolphinDiskCache()
        register_cache_provider(_provider)
        print(f"💾 [Dolphin] 디스크 노드 캐시 활성화: {_provider.dir}")