prepare_data.py

"""Download spreadsheet data from Sheetpedia/Enron → upload to Google Drive → register in metadata.

Can be run as a CLI or called from the app's /api/ingest endpoint.

CLI usage:
    python prepare_data.py --source sheetpedia --limit 10
    python prepare_data.py --source enron
    python prepare_data.py setup --root-folder-id 1ABC...XYZ
"""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path

import config
import drive_client
import metadata_client

STAGING_DIR = config.TEMP_DIR / "staging"
STAGING_DIR.mkdir(parents=True, exist_ok=True)


def _get_known_filenames() -> set[str]:
    """Collect filenames already uploaded to Drive / registered in metadata."""
    known: set[str] = set()
    try:
        rows = metadata_client.get_all_rows()
        for r in rows:
            known.add(r["file_name"])
    except Exception:
        pass
    for p in (config.DATA_DIR / "raw" / "sheetpedia").glob("*.xlsx"):
        known.add(p.name)
    return known


def ingest_sheetpedia(limit: int = 10, offset: int = 0) -> int:
    """Stream xlsx files from HuggingFace Sheetpedia, upload to Drive, register metadata.

    Skips files that are already present in metadata or local raw data.
    Use ``offset`` to jump ahead in the stream (handy after prior ingestions).
    """
    try:
        from datasets import load_dataset
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "datasets"])
        from datasets import load_dataset

    try:
        from openpyxl import load_workbook
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "openpyxl"])
        from openpyxl import load_workbook

    print(f"Streaming Sheetpedia dataset (limit={limit}, offset={offset})...")
    ds = load_dataset("tianzl66/Sheetpedia_xlsx", split="train", streaming=True)

    known = _get_known_filenames()
    print(f"  Already known: {len(known)} files — will skip duplicates")

    metadata_client.ensure_header()
    entries = []
    saved = 0
    skipped = 0

    for i, row in enumerate(ds):
        if saved >= limit:
            break

        if i < offset:
            continue

        xlsx_bytes = row.get("xlsx")
        key = row.get("__key__", f"sheet_{i}")

        if xlsx_bytes is None:
            continue

        safe_name = Path(key).name.replace("/", "_").replace(" ", "_")
        if not safe_name.endswith(".xlsx"):
            safe_name += ".xlsx"

        if safe_name in known:
            skipped += 1
            continue

        if isinstance(xlsx_bytes, dict) and "bytes" in xlsx_bytes:
            raw = xlsx_bytes["bytes"]
        elif isinstance(xlsx_bytes, (bytes, bytearray)):
            raw = bytes(xlsx_bytes)
        elif isinstance(xlsx_bytes, list):
            raw = bytes(xlsx_bytes)
        else:
            continue

        local_path = STAGING_DIR / safe_name
        local_path.write_bytes(raw)

        try:
            wb = load_workbook(str(local_path), read_only=True)
            sheet_names = wb.sheetnames
            wb.close()
        except Exception as e:
            print(f"  Skipping {safe_name}: invalid xlsx ({e})")
            local_path.unlink(missing_ok=True)
            continue

        # Also save to raw dir for local queue
        raw_dir = config.DATA_DIR / "raw" / "sheetpedia"
        raw_dir.mkdir(parents=True, exist_ok=True)
        raw_copy = raw_dir / safe_name
        if not raw_copy.exists():
            raw_copy.write_bytes(local_path.read_bytes())

        # Upload to Drive
        drive_file_id = drive_client.upload_file(
            str(local_path),
            config.DRIVE_UNLABELLED_FOLDER_ID,
            file_name=safe_name,
        )

        for sn in sheet_names:
            entries.append({
                "file_name": safe_name,
                "drive_file_id": drive_file_id,
                "sheet_name": sn,
                "num_sheets_in_file": str(len(sheet_names)),
            })

        known.add(safe_name)
        local_path.unlink(missing_ok=True)
        saved += 1
        if saved % 5 == 0:
            print(f"  Uploaded {saved}/{limit} files (skipped {skipped} duplicates)...")

    added = metadata_client.add_rows(entries)
    print(f"Sheetpedia: uploaded {saved} files, skipped {skipped} duplicates, registered {added} sheets")
    return added


def ingest_enron() -> int:
    """Clone Enron XLS repo, convert to xlsx, upload to Drive, register metadata."""
    enron_dir = STAGING_DIR / "enron"
    repo_dir = enron_dir / "enron_xls"

    if repo_dir.exists():
        print("Enron repo already cloned, reusing...")
    else:
        enron_dir.mkdir(parents=True, exist_ok=True)
        print("Cloning Enron XLS repo...")
        subprocess.run(
            ["git", "clone", "--depth", "1",
             "https://github.com/SheetJS/enron_xls.git",
             str(repo_dir)],
            check=True,
        )

    try:
        from openpyxl import load_workbook
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "openpyxl"])
        from openpyxl import load_workbook

    xls_files = []
    for ext in ("*.xls", "*.xlsx", "*.XLS", "*.XLSX"):
        xls_files.extend(repo_dir.rglob(ext))

    print(f"Found {len(xls_files)} spreadsheet files in Enron repo")

    metadata_client.ensure_header()
    entries = []
    uploaded = 0

    for xls_path in xls_files:
        xlsx_path = xls_path
        if xls_path.suffix.lower() == ".xls":
            xlsx_path = xls_path.with_suffix(".xlsx")
            if not xlsx_path.exists():
                try:
                    import xlrd
                    from openpyxl import Workbook
                    xls_wb = xlrd.open_workbook(str(xls_path))
                    new_wb = Workbook()
                    for idx, sn in enumerate(xls_wb.sheet_names()):
                        xls_sheet = xls_wb.sheet_by_name(sn)
                        ws = new_wb.active if idx == 0 else new_wb.create_sheet(title=sn)
                        if idx == 0:
                            ws.title = sn
                        for r in range(xls_sheet.nrows):
                            for c in range(xls_sheet.ncols):
                                ws.cell(row=r + 1, column=c + 1, value=xls_sheet.cell_value(r, c))
                    new_wb.save(str(xlsx_path))
                    new_wb.close()
                    xls_wb.release_resources()
                except ImportError:
                    print("  xlrd not installed; skipping .xls files. pip install xlrd")
                    break
                except Exception as e:
                    print(f"  Could not convert {xls_path.name}: {e}")
                    continue

        try:
            wb = load_workbook(str(xlsx_path), read_only=True)
            sheet_names = wb.sheetnames
            wb.close()
        except Exception:
            continue

        drive_file_id = drive_client.upload_file(
            str(xlsx_path),
            config.DRIVE_UNLABELLED_FOLDER_ID,
            file_name=xlsx_path.name,
        )

        for sn in sheet_names:
            entries.append({
                "file_name": xlsx_path.name,
                "drive_file_id": drive_file_id,
                "sheet_name": sn,
                "num_sheets_in_file": str(len(sheet_names)),
            })

        uploaded += 1
        if uploaded % 10 == 0:
            print(f"  Uploaded {uploaded} files...")

    added = metadata_client.add_rows(entries)
    print(f"Enron: uploaded {uploaded} files, registered {added} sheets")
    return added


def setup_folders(root_folder_id: str) -> None:
    """Create Drive folder structure and print env vars to set."""
    folder_ids = drive_client.setup_folder_structure(root_folder_id)
    metadata_client.ensure_header()
    print("\nDrive folder structure created!")
    print("\nSet these environment variables:\n")
    for key, value in folder_ids.items():
        print(f"  {key}={value}")
    print(f"  METADATA_SHEET_ID={config.METADATA_SHEET_ID}")
    print()


def main():
    parser = argparse.ArgumentParser(description="Prepare spreadsheet data for labeling")
    sub = parser.add_subparsers(dest="command")

    ingest_parser = sub.add_parser("ingest", help="Download and upload data")
    ingest_parser.add_argument("--source", choices=["sheetpedia", "enron", "both"], default="sheetpedia")
    ingest_parser.add_argument("--limit", type=int, default=10, help="Max new Sheetpedia files to download")
    ingest_parser.add_argument("--offset", type=int, default=0, help="Skip first N entries in the stream")

    setup_parser = sub.add_parser("setup", help="Create Drive folder structure")
    setup_parser.add_argument("--root-folder-id", required=True, help="ID of the shared Drive folder")

    sync_parser = sub.add_parser("sync", help="Sync metadata with Drive folder contents")

    args = parser.parse_args()

    if args.command == "setup":
        setup_folders(args.root_folder_id)
    elif args.command == "sync":
        from app import sync_metadata
        result = sync_metadata()
        print(f"Scanned: {result['scanned']}, New sheets added: {result['new_sheets_added']}")
    elif args.command == "ingest" or args.command is None:
        if not hasattr(args, "source"):
            args.source = "sheetpedia"
            args.limit = 10
            args.offset = 0

        if args.source in ("sheetpedia", "both"):
            ingest_sheetpedia(limit=args.limit, offset=args.offset)
        if args.source in ("enron", "both"):
            ingest_enron()
        print("Done!")
    else:
        parser.print_help()


if __name__ == "__main__":
    main()