import json
import random
from pathlib import Path

from loguru import logger

import config


# ── OCR 캐시 (파일별 json 분리) ───────────────────

def _cache_path(image_path: str) -> Path:
    """
    이미지 경로 → json 저장 경로
    files/data/TSMC_A/1000114581_INV/page_001.png
    → files/json/TSMC_A/1000114581_INV/page_001.json
    """
    img = Path(image_path)
    try:
        rel = img.relative_to(config.DATA_DIR)
    except ValueError:
        rel = Path(img.parent.name) / img.name
    return config.OCR_CACHE_DIR / rel.with_suffix(".json")


def save_cache_by_file(ocr_result: dict):
    """이미지 1장의 OCR 결과를 개별 json으로 저장"""
    path = _cache_path(ocr_result["image_path"])
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(ocr_result, f, ensure_ascii=False, indent=2)
    logger.debug(f"캐시 저장: {path}")


def load_cache() -> list:
    """전체 캐시 통합 로드 (학습/평가 시 사용)"""
    samples = []
    for p in config.OCR_CACHE_DIR.rglob("*.json"):
        try:
            with open(p, encoding="utf-8") as f:
                samples.append(json.load(f))
        except (json.JSONDecodeError, UnicodeDecodeError) as e:
            logger.warning(f"캐시 로드 실패 (건너뜀): {p} | {e}")
    return samples


def get_cached_paths() -> set:
    """캐시된 이미지 경로 전체 반환 (중복 방지용)"""
    cached = []
    for p in config.OCR_CACHE_DIR.rglob("*.json"):
        try:
            with open(p, encoding="utf-8") as f:
                data = json.load(f)
            cached.append(data.get("image_path", ""))
            # image_path = data.get("image_path", "")
            # if image_path:
            #     cached.append(str(Path(image_path)))  # 정규화
            #     print(str(Path(image_path)))
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass
    return set(cached)


def cache_status() -> dict:
    """양식별 캐시 현황 (폴더별 json 파일 수 집계)"""
    if not config.OCR_CACHE_DIR.exists():
        return {}
    status = {}
    for cls_dir in sorted(config.OCR_CACHE_DIR.iterdir()):
        if cls_dir.is_dir():
            status[cls_dir.name] = len(list(cls_dir.rglob("*.json")))
    return status


# ── label2id ──────────────────────────────────────

def load_label2id() -> dict:
    path = config.MODEL_SAVE_PATH / "label2id.json"
    if path.exists():
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    if config.DATA_DIR.exists():
        classes = sorted([d.name for d in config.DATA_DIR.iterdir() if d.is_dir()])
        return {c: i for i, c in enumerate(classes)}
    return {}


def save_label2id(label2id: dict):
    config.MODEL_SAVE_PATH.mkdir(exist_ok=True)
    with open(config.MODEL_SAVE_PATH / "label2id.json", "w", encoding="utf-8") as f:
        json.dump(label2id, f, ensure_ascii=False, indent=2)
    logger.info(f"label2id 저장: {label2id}")


# ── 리플레이 버퍼 ──────────────────────────────────

def load_replay_buffer() -> dict:
    if not config.REPLAY_BUFFER_FILE.exists():
        return {}
    try:
        with open(config.REPLAY_BUFFER_FILE, encoding="utf-8") as f:
            return json.load(f)
    except (json.JSONDecodeError, UnicodeDecodeError) as e:
        backup = config.REPLAY_BUFFER_FILE.with_suffix(".json.bak")
        config.REPLAY_BUFFER_FILE.rename(backup)
        logger.warning(f"replay_buffer.json 손상 → 백업 후 초기화: {e}")
        return {}


def update_replay_buffer(new_samples: list, id2label: dict, per_class: int = 10):
    buf      = load_replay_buffer()
    by_label = {}
    for s in new_samples:
        name = id2label[s["label"]]
        by_label.setdefault(name, []).append(s)

    for name, samples in by_label.items():
        combined  = buf.get(name, []) + samples
        buf[name] = random.sample(combined, min(per_class, len(combined)))

    with open(config.REPLAY_BUFFER_FILE, "w", encoding="utf-8") as f:
        json.dump(buf, f, ensure_ascii=False, indent=2)
    logger.info(f"리플레이 버퍼 갱신: { {k: len(v) for k, v in buf.items()} }")