#!/usr/bin/env python3
"""Extract GBK-named ZIP archives safely and generate file indexes."""

from __future__ import annotations

import argparse
import csv
import hashlib
import os
import posixpath
import shutil
import sys
import zipfile
from pathlib import Path


def decode_zip_name(name: str) -> str:
    raw = name.encode("cp437", errors="replace")
    for encoding in ("gbk", "gb18030", "utf-8"):
        try:
            return raw.decode(encoding)
        except UnicodeDecodeError:
            continue
    return name


def safe_parts(name: str) -> list[str]:
    normalized = posixpath.normpath(name.replace("\\", "/"))
    parts = [part for part in normalized.split("/") if part and part != "."]
    if any(part == ".." for part in parts):
        raise ValueError(f"unsafe zip path: {name}")
    return parts


def copy_with_hash(src, dst: Path) -> str:
    h = hashlib.sha256()
    dst.parent.mkdir(parents=True, exist_ok=True)
    with dst.open("wb") as out:
        while True:
            chunk = src.read(1024 * 1024)
            if not chunk:
                break
            h.update(chunk)
            out.write(chunk)
    return h.hexdigest()


def extract_zip(zip_path: Path, dest_root: Path, index_rows: list[dict[str, str]], source_label: str) -> None:
    with zipfile.ZipFile(zip_path) as zf:
        for info in zf.infolist():
            decoded = decode_zip_name(info.filename)
            parts = safe_parts(decoded)
            if not parts:
                continue
            out_path = dest_root.joinpath(*parts)
            if info.is_dir() or decoded.endswith("/"):
                out_path.mkdir(parents=True, exist_ok=True)
                continue
            with zf.open(info) as src:
                sha256 = copy_with_hash(src, out_path)
            mtime = ""
            try:
                y, m, d, hh, mm, ss = info.date_time
                mtime = f"{y:04d}-{m:02d}-{d:02d} {hh:02d}:{mm:02d}:{ss:02d}"
            except Exception:
                pass
            index_rows.append(
                {
                    "source_label": source_label,
                    "source_zip": str(zip_path),
                    "zip_internal_path": decoded,
                    "extracted_path": str(out_path),
                    "relative_to_dest": str(out_path.relative_to(dest_root)),
                    "extension": out_path.suffix.lower(),
                    "file_size": str(info.file_size),
                    "compressed_size": str(info.compress_size),
                    "zip_mtime": mtime,
                    "sha256": sha256,
                }
            )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--zip", action="append", required=True, dest="zips")
    parser.add_argument("--dest-root", required=True)
    parser.add_argument("--index-csv", required=True)
    parser.add_argument("--replace", action="store_true")
    args = parser.parse_args()

    dest_root = Path(args.dest_root)
    if args.replace and dest_root.exists():
        shutil.rmtree(dest_root)
    dest_root.mkdir(parents=True, exist_ok=True)

    rows: list[dict[str, str]] = []
    for zip_arg in args.zips:
        zip_path = Path(zip_arg)
        if not zip_path.exists():
            raise FileNotFoundError(zip_path)
        source_label = zip_path.stem
        extract_zip(zip_path, dest_root, rows, source_label)

    fieldnames = [
        "source_label",
        "source_zip",
        "zip_internal_path",
        "extracted_path",
        "relative_to_dest",
        "extension",
        "file_size",
        "compressed_size",
        "zip_mtime",
        "sha256",
    ]
    index_path = Path(args.index_csv)
    index_path.parent.mkdir(parents=True, exist_ok=True)
    with index_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"extracted_files={len(rows)}")
    print(f"dest_root={dest_root}")
    print(f"index_csv={index_path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
