#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import zipfile
import logging
from pathlib import Path
from tqdm import tqdm

# ─── 1. 설정 ───────────────────────────────────────────────
ROOT = Path("/data/ai-hub/ppe-dataset/105.공사현장_안전장비_인식_데이터/01.데이터")
LOG_FILE = ROOT / "zip_process.log"

# ─── 2. 로거 세팅 ───────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[
        logging.FileHandler(LOG_FILE, encoding="utf-8"),
        logging.StreamHandler()
    ]
)

def process_zip(zippath: Path):
    """ 1) 무결성 테스트 2) 해제 디렉토리 존재/내용 검사 3) 필요 시 압축 해제 """
    target_dir = zippath.with_suffix("")  # .zip 제거한 폴더
    # 1) integrity check
    try:
        with zipfile.ZipFile(zippath, 'r') as zf:
            bad = zf.testzip()
            if bad:
                logging.error(f"✖ 손상: {zippath} → 첫 오류 엔트리: {bad}")
                return
    except zipfile.BadZipFile:
        logging.error(f"✖ BadZipFile: {zippath}")
        return
    except Exception as e:
        logging.error(f"✖ 무결성 검사 실패: {zippath} → {e}")
        return

    # 2) 이미 해제된 상태인지 검사
    if target_dir.exists() and any(target_dir.iterdir()):
        logging.info(f"✔ 이미 해제됨, 스킵: {target_dir}")
        return

    # 3) 압축 해제
    logging.info(f"→ 해제 시작: {zippath}")
    try:
        target_dir.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(zippath, 'r') as zf:
            for member in tqdm(zf.infolist(),
                               desc=f"Extracting {zippath.name}",
                               unit="file"):
                try:
                    zf.extract(member, path=target_dir)
                except Exception:
                    # 유니코드 경로 오류 등 무시
                    pass
        logging.info(f"✔ 해제 완료: {target_dir}")
    except Exception as e:
        logging.error(f"✖ 해제 실패: {zippath} → {e}")
        # 실패 시, 디렉토리 풀던 도중 생긴 쓰레기 제거
        # (필요하면 아래 주석 해제)
        # import shutil; shutil.rmtree(target_dir, ignore_errors=True)

def main():
    zip_files = list(ROOT.rglob("*.zip"))
    logging.info(f"총 {len(zip_files)}개의 ZIP 파일을 찾았습니다.")
    for z in zip_files:
        process_zip(z)
    logging.info("=== 모든 ZIP 처리 완료 ===")

if __name__ == "__main__":
    main()