gongxiangp2/scripts/scan_sku_folders_to_sqlite.py

#!/usr/bin/env python3
import argparse
import json
import os
import re
import sqlite3
import sys
import uuid
from datetime import datetime, timezone
from pathlib import Path


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".bmp", ".gif"}
NON_DIRECTORY_EXTENSIONS = IMAGE_EXTENSIONS | {
    ".xlsx", ".xls", ".psd", ".psb", ".zip", ".rar", ".7z", ".pdf", ".txt", ".doc", ".docx"
}
SKU_DIR_PATTERN = re.compile(r"sku", re.IGNORECASE)
SPU_PATTERN = re.compile(r"^[A-Za-z]{1,5}\d{3,5}$")
SPU_TOKEN_PATTERN = re.compile(r"([A-Za-z]{1,5}\d{3,5})")
BRANDED_SPU_TOKEN_PATTERN = re.compile(r"([A-Za-z]{2,5}-[A-Za-z]{1,5}\d{3,5})")
BRANDED_SKU_FILE_PATTERN = re.compile(r"([A-Za-z]{2,5})-([A-Za-z]{1,5})(\d{4,5})(\d{2})(?:\D|$)")
KNOWN_SHOP_PREFIXES = {"GBPM", "GBSFY", "GBGJ", "GBJP", "XYL"}
IGNORED_SPU_PREFIXES = {"TASK", "SKU", "IMG"}
KNOWN_CATEGORY_CODES = {
    "SN", "DG", "XG", "BG", "SG", "CBG", "CZ", "SZ", "DNZ", "SJZ", "ZDZ",
    "CBZ", "CY", "AD", "SZD", "HXD", "SJ", "ZSJ", "ZWJ", "CFZW", "CFZWJ",
    "YMJ", "BJ", "CJ", "DSG", "CTG", "CYJ", "SZT", "XJ", "LRSF", "SMC",
}
IGNORED_PARTS = {"@eadir", "dingtalkappdata", ".recycle", "$nsec.recycle.bin"}
DETAIL_DIR_KEYWORDS = ("images", "详情页", "详情")
MAIN_DIR_KEYWORD = "主图"
MAIN_DIMENSION_DIR_NAMES = {"750", "800"}


def utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def normalize_spu(value: str) -> str:
    return value.strip().upper()


def looks_like_spu_token(value: str) -> bool:
    normalized = normalize_spu(value)
    if "-" in normalized:
        prefix, _, sku_part = normalized.partition("-")
        return prefix in KNOWN_SHOP_PREFIXES and bool(SPU_PATTERN.match(sku_part))

    return bool(SPU_PATTERN.match(normalized)) and not any(
        normalized.startswith(prefix) for prefix in IGNORED_SPU_PREFIXES
    )


def is_standalone_spu_match(text: str, start: int, end: int) -> bool:
    prev_char = text[start - 1] if start > 0 else ""
    next_char = text[end] if end < len(text) else ""
    if re.match(r"[A-Za-z0-9]", prev_char or ""):
        return False
    if not re.match(r"[A-Za-z0-9]", next_char or ""):
        return True
    return bool(re.match(r"K\d{1,2}(?![A-Za-z0-9])", text[end:], re.IGNORECASE))


def get_sku_part(spu: str) -> str:
    normalized = normalize_spu(spu)
    return normalized.partition("-")[2] if "-" in normalized else normalized


def has_known_category_code(spu: str) -> bool:
    sku_part = get_sku_part(spu)
    for code in sorted(KNOWN_CATEGORY_CODES, key=len, reverse=True):
        if sku_part.startswith(code) and len(sku_part) > len(code) and sku_part[len(code)].isdigit():
            return True
    return False


def collect_spu_tokens(path_obj: Path):
    tokens = []
    seen = set()
    for part in reversed(path_obj.parts):
        matches = list(BRANDED_SPU_TOKEN_PATTERN.finditer(part))
        for match in reversed(matches):
            token = normalize_spu(match.group(1))
            if (
                token not in seen
                and looks_like_spu_token(token)
                and is_standalone_spu_match(part, match.start(), match.end())
            ):
                tokens.append(token)
                seen.add(token)

    for part in reversed(path_obj.parts):
        normalized = normalize_spu(part)
        if normalized not in seen and looks_like_spu_token(normalized):
            tokens.append(normalized)
            seen.add(normalized)

        matches = list(SPU_TOKEN_PATTERN.finditer(part))
        for match in reversed(matches):
            token = normalize_spu(match.group(1))
            if (
                token not in seen
                and looks_like_spu_token(token)
                and is_standalone_spu_match(part, match.start(), match.end())
            ):
                tokens.append(token)
                seen.add(token)

    return tokens


def guess_spu(path_obj: Path) -> str:
    tokens = collect_spu_tokens(path_obj)
    for token in tokens:
        if has_known_category_code(token):
            return token
    if tokens:
        return tokens[0]
    return ""


def guess_spu_from_file_names(file_names) -> str:
    for file_name in file_names:
        match = BRANDED_SKU_FILE_PATTERN.search(file_name)
        if not match:
            continue
        token = normalize_spu(f"{match.group(1)}-{match.group(2)}{match.group(3)}")
        if looks_like_spu_token(token):
            return token
    return ""


def matches_sku_dir(name: str) -> bool:
    return bool(SKU_DIR_PATTERN.search(name or ""))


def should_ignore_path(path_obj: Path) -> bool:
    for part in path_obj.parts:
        if part.strip().lower() in IGNORED_PARTS:
            return True
    return False


def looks_like_file_name(name: str) -> bool:
    return Path(name).suffix.lower() in NON_DIRECTORY_EXTENSIONS


def normalize_dir_name(name: str) -> str:
    return re.sub(r"[\s_\-]+", "", (name or "")).strip().lower()


def is_detail_dir_name(name: str) -> bool:
    normalized = normalize_dir_name(name)
    return any(keyword in normalized for keyword in DETAIL_DIR_KEYWORDS)


def is_main_dir_name(name: str) -> bool:
    normalized = normalize_dir_name(name)
    return MAIN_DIR_KEYWORD in normalized or normalized in MAIN_DIMENSION_DIR_NAMES


def list_sibling_dirs(parent_path: Path, current_dir_name: str):
    try:
        entries = list(parent_path.iterdir())
    except OSError:
        return [], [], []

    sibling_names = []
    detail_names = []
    main_names = []

    for entry in entries:
        if not entry.is_dir():
            continue
        if entry.name == current_dir_name:
            continue
        if should_ignore_path(entry) or looks_like_file_name(entry.name):
            continue

        sibling_names.append(entry.name)
        if is_detail_dir_name(entry.name):
            detail_names.append(entry.name)
        if is_main_dir_name(entry.name):
            main_names.append(entry.name)

    return sibling_names, detail_names, main_names


def ensure_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS sku_folder_scans (
          scan_id TEXT PRIMARY KEY,
          root_path TEXT NOT NULL,
          started_at TEXT NOT NULL,
          finished_at TEXT,
          total_rows INTEGER NOT NULL DEFAULT 0
        );

        CREATE TABLE IF NOT EXISTS sku_folder_scan_items (
          id INTEGER PRIMARY KEY AUTOINCREMENT,
          scan_id TEXT NOT NULL,
          root_path TEXT NOT NULL,
          sku_dir_path TEXT NOT NULL,
          relative_path TEXT NOT NULL,
          directory_name TEXT NOT NULL,
          parent_path TEXT,
          grandparent_path TEXT,
          spu TEXT,
          depth INTEGER NOT NULL DEFAULT 0,
          file_count INTEGER NOT NULL DEFAULT 0,
          image_count INTEGER NOT NULL DEFAULT 0,
          child_dir_count INTEGER NOT NULL DEFAULT 0,
          sibling_dir_names_json TEXT NOT NULL DEFAULT '[]',
          sibling_detail_dir_names_json TEXT NOT NULL DEFAULT '[]',
          sibling_main_dir_names_json TEXT NOT NULL DEFAULT '[]',
          has_detail_dir_sibling INTEGER NOT NULL DEFAULT 0,
          has_main_dir_sibling INTEGER NOT NULL DEFAULT 0,
          has_detail_and_main_siblings INTEGER NOT NULL DEFAULT 0,
          sample_files_json TEXT NOT NULL DEFAULT '[]',
          created_at TEXT NOT NULL,
          FOREIGN KEY (scan_id) REFERENCES sku_folder_scans(scan_id)
        );

        CREATE INDEX IF NOT EXISTS idx_sku_folder_scan_items_scan_id
        ON sku_folder_scan_items(scan_id);

        CREATE INDEX IF NOT EXISTS idx_sku_folder_scan_items_spu
        ON sku_folder_scan_items(spu);

        CREATE INDEX IF NOT EXISTS idx_sku_folder_scan_items_relative_path
        ON sku_folder_scan_items(relative_path);
        """
    )

    expected_columns = {
        "sibling_dir_names_json": "TEXT NOT NULL DEFAULT '[]'",
        "sibling_detail_dir_names_json": "TEXT NOT NULL DEFAULT '[]'",
        "sibling_main_dir_names_json": "TEXT NOT NULL DEFAULT '[]'",
        "has_detail_dir_sibling": "INTEGER NOT NULL DEFAULT 0",
        "has_main_dir_sibling": "INTEGER NOT NULL DEFAULT 0",
        "has_detail_and_main_siblings": "INTEGER NOT NULL DEFAULT 0",
    }
    existing_columns = {
        row[1]
        for row in conn.execute("PRAGMA table_info(sku_folder_scan_items)").fetchall()
    }
    for column_name, column_definition in expected_columns.items():
        if column_name in existing_columns:
            continue
        conn.execute(
            f"ALTER TABLE sku_folder_scan_items ADD COLUMN {column_name} {column_definition}"
        )
    conn.commit()


def scan_sku_dirs(root_path: Path, skip_repeated_root_name: bool = False):
    for current_root, dir_names, file_names in os.walk(root_path):
        current_path = Path(current_root)
        if should_ignore_path(current_path):
            dir_names[:] = []
            continue
        if skip_repeated_root_name and current_path == root_path:
            root_name = root_path.name.casefold()
            dir_names[:] = [name for name in dir_names if name.casefold() != root_name]

        for dir_name in list(dir_names):
            if not matches_sku_dir(dir_name):
                continue

            sku_dir_path = current_path / dir_name
            if should_ignore_path(sku_dir_path) or looks_like_file_name(dir_name) or not sku_dir_path.is_dir():
                continue
            try:
                child_entries = list(sku_dir_path.iterdir())
            except OSError:
                child_entries = []

            files = [entry for entry in child_entries if entry.is_file()]
            child_dirs = [entry for entry in child_entries if entry.is_dir()]
            image_files = [entry.name for entry in files if entry.suffix.lower() in IMAGE_EXTENSIONS]
            sibling_names, detail_sibling_names, main_sibling_names = list_sibling_dirs(
                current_path, sku_dir_path.name
            )

            relative_path = sku_dir_path.relative_to(root_path).as_posix()
            yield {
                "root_path": str(root_path),
                "sku_dir_path": str(sku_dir_path),
                "relative_path": relative_path,
                "directory_name": sku_dir_path.name,
                "parent_path": str(sku_dir_path.parent),
                "grandparent_path": str(sku_dir_path.parent.parent) if sku_dir_path.parent != root_path else "",
                "spu": guess_spu(sku_dir_path) or guess_spu_from_file_names(image_files),
                "depth": len(sku_dir_path.relative_to(root_path).parts),
                "file_count": len(files),
                "image_count": len(image_files),
                "child_dir_count": len(child_dirs),
                "sibling_dir_names_json": json.dumps(sibling_names, ensure_ascii=False),
                "sibling_detail_dir_names_json": json.dumps(detail_sibling_names, ensure_ascii=False),
                "sibling_main_dir_names_json": json.dumps(main_sibling_names, ensure_ascii=False),
                "has_detail_dir_sibling": 1 if detail_sibling_names else 0,
                "has_main_dir_sibling": 1 if main_sibling_names else 0,
                "has_detail_and_main_siblings": 1 if detail_sibling_names and main_sibling_names else 0,
                "sample_files_json": json.dumps(image_files[:12], ensure_ascii=False),
            }


def insert_rows(conn: sqlite3.Connection, scan_id: str, rows) -> int:
    created_at = utc_now()
    total = 0
    for row in rows:
        conn.execute(
            """
            INSERT INTO sku_folder_scan_items (
              scan_id, root_path, sku_dir_path, relative_path, directory_name,
              parent_path, grandparent_path, spu, depth, file_count, image_count,
              child_dir_count, sibling_dir_names_json, sibling_detail_dir_names_json,
              sibling_main_dir_names_json, has_detail_dir_sibling, has_main_dir_sibling,
              has_detail_and_main_siblings, sample_files_json, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                scan_id,
                row["root_path"],
                row["sku_dir_path"],
                row["relative_path"],
                row["directory_name"],
                row["parent_path"],
                row["grandparent_path"],
                row["spu"],
                row["depth"],
                row["file_count"],
                row["image_count"],
                row["child_dir_count"],
                row["sibling_dir_names_json"],
                row["sibling_detail_dir_names_json"],
                row["sibling_main_dir_names_json"],
                row["has_detail_dir_sibling"],
                row["has_main_dir_sibling"],
                row["has_detail_and_main_siblings"],
                row["sample_files_json"],
                created_at,
            ),
        )
        total += 1
    conn.commit()
    return total


def parse_args():
    parser = argparse.ArgumentParser(description="Scan all sku folders and save results to sqlite.")
    parser.add_argument(
        "--root",
        default="/mnt/history-sc/大货营销策划图片",
        help="Root directory to scan.",
    )
    parser.add_argument(
        "--roots",
        default="",
        help="Multiple root directories separated by |. When provided, --root is ignored.",
    )
    parser.add_argument(
        "--db-path",
        default="/root/gongxiangp2/output/history-sku-folders.db",
        help="SQLite database file path.",
    )
    parser.add_argument(
        "--skip-repeated-root-name",
        action="store_true",
        help="Skip an immediate child directory whose name equals the scan root name.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    root_paths = [Path(item.strip()) for item in args.roots.split("|") if item.strip()] if args.roots else [Path(args.root)]
    db_path = Path(args.db_path)

    missing_roots = [str(root_path) for root_path in root_paths if not root_path.exists()]
    if missing_roots:
        print(f"Root path not found: {missing_roots[0]}", file=sys.stderr)
        return 1

    db_path.parent.mkdir(parents=True, exist_ok=True)

    conn = sqlite3.connect(str(db_path))
    try:
        ensure_schema(conn)
        all_previews = []
        grand_total = 0
        scan_ids = []
        for root_path in root_paths:
            scan_id = str(uuid.uuid4())
            scan_ids.append(scan_id)
            started_at = utc_now()
            conn.execute(
                "INSERT INTO sku_folder_scans (scan_id, root_path, started_at) VALUES (?, ?, ?)",
                (scan_id, str(root_path), started_at),
            )
            conn.commit()

            rows = list(scan_sku_dirs(root_path, args.skip_repeated_root_name))
            total_rows = insert_rows(conn, scan_id, rows)
            grand_total += total_rows
            finished_at = utc_now()
            conn.execute(
                "UPDATE sku_folder_scans SET finished_at = ?, total_rows = ? WHERE scan_id = ?",
                (finished_at, total_rows, scan_id),
            )
            conn.commit()
            all_previews.extend(rows[:3])

        print("SKU folder scan completed")
        print(f"scan_ids={','.join(scan_ids)}")
        print(f"db_path={db_path}")
        print(f"total_rows={grand_total}")
        print(json.dumps(all_previews[:12], ensure_ascii=False, indent=2))
        return 0
    finally:
        conn.close()


if __name__ == "__main__":
    raise SystemExit(main())