"""批量解压 CDS 下载的 ZIP 格式 NetCDF 文件"""
import logging
import zipfile
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor

from src.utils.config import DATA_RAW, CITIES

logger = logging.getLogger(__name__)


def extract_one(nc_path: Path) -> bool:
    """解压单个 ZIP 伪装的 .nc 文件，原地替换为真实 NetCDF"""
    if nc_path.stat().st_size == 0:
        return False

    with open(nc_path, "rb") as f:
        header = f.read(4)

    if header[:2] == b"CDF":  # 已经是真实 NetCDF
        return True
    if header[:2] != b"PK":  # 不是 ZIP，跳过
        return False

    try:
        with zipfile.ZipFile(nc_path) as zf:
            nc_names = [n for n in zf.namelist() if n.endswith(".nc")]
            if not nc_names:
                return False
            data = zf.read(nc_names[0])
        nc_path.write_bytes(data)
        return True
    except Exception:
        return False


def extract_all():
    for city in CITIES:
        city_dir = Path(DATA_RAW) / "era5" / city
        files = sorted(city_dir.glob("*.nc"))
        done = 0
        for f in files:
            if extract_one(f):
                done += 1
        logger.info("%s: %d/%d 已解压", city, done, len(files))


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
    extract_all()
