script_torrent.py

import os
import sys
import json
import logging
import tarfile
import argparse
import gc
from pathlib import Path
from typing import Optional, Dict, Any, Tuple, List
from urllib.parse import quote
from concurrent.futures import ProcessPoolExecutor, as_completed, wait, FIRST_COMPLETED

import io
import zstandard as zstd
import pypdfium2 as pdfium
import pyarrow as pa
import pyarrow.parquet as pq
import pandas as pd

# --- Configuration ---
DEFAULT_BATCH_SIZE = 100
MAX_BATCH_BYTES = 512 * 1024 * 1024
MAX_SINGLE_FILE_BYTES = 2 * 1024 * 1024 * 1024 - 1024
MAX_PENDING_FUTURES = os.cpu_count() * 2 if os.cpu_count() else 8

# --- Logging Setup ---
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("EpsteinProcessor")

# --- Schema ---
media_struct = pa.struct([("bytes", pa.binary()), ("path", pa.string())])

SCHEMA = pa.schema(
    [
        ("dataset_id", pa.int32()),
        ("doc_id", pa.string()),
        ("file_name", pa.string()),
        ("file_type", pa.string()),
        ("online_url", pa.string()),
        ("text_content", pa.string()),
        ("audio", media_struct),
        ("image", media_struct),
        ("video", media_struct),
        ("metadata", pa.string()),
        ("error", pa.string()),
    ]
)

HF_METADATA = {
    "huggingface": json.dumps(
        {
            "dataset_info": {
                "features": {
                    "audio": {"_type": "Audio"},
                    "image": {"_type": "Image"},
                    "video": {"_type": "Video"},
                }
            }
        }
    )
}


# --- Classes ---


class ChunkedParquetWriter:
    def __init__(
        self,
        output_dir: Path,
        schema: pa.Schema,
        target_file_size: int = 500 * 1024 * 1024,
    ):
        self.output_dir = output_dir
        self.schema = schema
        self.target_file_size = target_file_size
        self.current_file_index = 0
        self.current_writer = None
        self.current_file_size = 0
        self._open_new_writer()

    def _open_new_writer(self):
        if self.current_writer:
            self.current_writer.close()

        filename = (
            self.output_dir / f"epstein_files-{self.current_file_index:04d}.parquet"
        )
        self.current_writer = pq.ParquetWriter(
            str(filename),
            schema=self.schema,
            compression="zstd",
            use_content_defined_chunking=True,
        )
        self.current_file_index += 1
        self.current_file_size = 0
        logger.info(f"Started new parquet chunk: {filename.name}")

    def write_table(self, table: pa.Table):
        self.current_writer.write_table(table)
        # Estimate size based on uncompressed nbytes (conservative)
        self.current_file_size += table.nbytes

        if self.current_file_size >= self.target_file_size:
            self._open_new_writer()

    def close(self):
        if self.current_writer:
            self.current_writer.close()
            self.current_writer = None


# --- Helper Functions (Shared) ---


def get_online_url(dataset_num: int, filename: str) -> str:
    return f"https://www.justice.gov/epstein/files/{quote(f'DataSet {dataset_num}')}/{filename}"


def determine_type(filename: str) -> str:
    ext = os.path.splitext(filename)[1].lower()
    if ext == ".pdf":
        return "pdf"
    if ext in [".jpg", ".jpeg", ".png", ".tif", ".tiff", ".bmp", ".gif"]:
        return "image"
    if ext in [".mp4", ".mov", ".avi", ".wmv", ".mkv", ".m4v", ".3gp"]:
        return "video"
    if ext in [".mp3", ".wav", ".aac", ".wma", ".flac", ".opus", ".amr", ".m4a"]:
        return "audio"
    if ext in [".csv", ".xlsx", ".xls", ".txt", ".msg"]:
        return "data"
    return "other"


def parse_dataset_number(path_parts: list) -> int:
    for part in path_parts:
        if part.startswith("dataset_") and part[8:].isdigit():
            return int(part[8:])
    return -1


# --- Worker Logic (CPU Bound) ---


def extract_pdf_content(
    file_bytes: bytes, filename: str
) -> Tuple[List[Dict[str, Any]], Optional[str]]:
    results = []
    error_msg = None
    try:
        pdf = pdfium.PdfDocument(file_bytes)
        try:
            # 1. Extract text from all pages
            text_parts = []
            for i in range(len(pdf)):
                try:
                    text_parts.append(pdf[i].get_textpage().get_text_range())
                except Exception:
                    continue
            full_text = "\n".join(text_parts)
            results.append(
                {
                    "file_type": "pdf",
                    "text_content": full_text,
                    "image": None,
                    "file_name": filename,
                }
            )

            # 2. Extract images
            img_idx = 0
            for i in range(len(pdf)):
                page = pdf[i]
                for obj in page.get_objects():
                    if isinstance(obj, pdfium.PdfImage):
                        img_idx += 1
                        try:
                            # Try to get the image directly as a PIL object
                            bitmap = obj.get_bitmap()
                            pil_img = bitmap.to_pil()

                            # Encode to PNG in memory
                            img_byte_arr = io.BytesIO()
                            pil_img.save(img_byte_arr, format="PNG")
                            img_bytes = img_byte_arr.getvalue()

                            img_filename = (
                                f"{os.path.splitext(filename)[0]}_img_{img_idx}.png"
                            )
                            results.append(
                                {
                                    "file_type": "image",
                                    "text_content": None,
                                    "image": {"bytes": img_bytes, "path": img_filename},
                                    "file_name": img_filename,
                                }
                            )
                        except Exception as img_err:
                            logger.debug(
                                f"Failed to extract image {img_idx} from {filename}: {img_err}"
                            )
                            continue

            return results, None
        finally:
            pdf.close()
    except Exception as e:
        return [], f"PDF extraction error: {str(e)}"


def extract_text_from_data(
    content: bytes, filename: str
) -> Tuple[Optional[str], Optional[str]]:
    if not (filename.endswith(".txt") or filename.endswith(".csv")):
        return None, None
    try:
        return content.decode("utf-8", errors="ignore"), None
    except Exception as e:
        return None, f"Text decode error: {str(e)}"


def worker_process_content(
    content: bytes,
    filename: str,
    path_parts: list,
    f_type: str,
    member_size: int,
    ds_num: int,
) -> List[Dict[str, Any]]:
    """
    Worker function to handle CPU-intensive processing (PDF/Text).
    Returns a list of records to be added to the Parquet file.
    """
    results = []

    try:
        if f_type == "pdf":
            pdf_results, error_msg = extract_pdf_content(content, filename)
            for res in pdf_results:
                results.append(
                    {
                        "dataset_id": ds_num,
                        "doc_id": os.path.splitext(filename)[0],
                        "file_name": res["file_name"],
                        "file_type": res["file_type"],
                        "online_url": get_online_url(ds_num, filename),
                        "text_content": res["text_content"],
                        "audio": None,
                        "image": res["image"],
                        "video": None,
                        "metadata": json.dumps(
                            {
                                "path": "/".join(path_parts),
                                "size": member_size,
                                "source": filename,
                            }
                        ),
                        "error": error_msg,
                    }
                )
            # If PDF failed entirely but returned an error
            if not pdf_results and error_msg:
                results.append(
                    {
                        "dataset_id": ds_num,
                        "doc_id": os.path.splitext(filename)[0],
                        "file_name": filename,
                        "file_type": f_type,
                        "online_url": get_online_url(ds_num, filename),
                        "text_content": None,
                        "audio": None,
                        "image": None,
                        "video": None,
                        "metadata": json.dumps(
                            {"path": "/".join(path_parts), "size": member_size}
                        ),
                        "error": error_msg,
                    }
                )
        elif f_type == "data":
            text_content, error_msg = extract_text_from_data(content, filename)
            results.append(
                {
                    "dataset_id": ds_num,
                    "doc_id": os.path.splitext(filename)[0],
                    "file_name": filename,
                    "file_type": f_type,
                    "online_url": get_online_url(ds_num, filename),
                    "text_content": text_content,
                    "audio": None,
                    "image": None,
                    "video": None,
                    "metadata": json.dumps(
                        {"path": "/".join(path_parts), "size": member_size}
                    ),
                    "error": error_msg,
                }
            )

        return results
    except Exception as e:
        return [
            {
                "dataset_id": ds_num,
                "doc_id": os.path.splitext(filename)[0],
                "file_name": filename,
                "file_type": f_type,
                "online_url": get_online_url(ds_num, filename),
                "text_content": None,
                "audio": None,
                "image": None,
                "video": None,
                "metadata": json.dumps(
                    {"path": "/".join(path_parts), "size": member_size}
                ),
                "error": f"Worker process error: {str(e)}",
            }
        ]


# --- Main Logic ---


def flush_batch(writer: pq.ParquetWriter, buffer: list):
    if not buffer:
        return
    try:
        table = pa.Table.from_pandas(pd.DataFrame(buffer), schema=SCHEMA)
        table = table.replace_schema_metadata(HF_METADATA)
        writer.write_table(table)
    except Exception as e:
        logger.error(f"FATAL WRITE ERROR: {e}")
    finally:
        buffer.clear()
        gc.collect()


def main():
    parser = argparse.ArgumentParser(description="Optimized Epstein Processor")
    parser.add_argument("input_archive", help="Path to .tar.zst")
    parser.add_argument("-o", "--output-dir", default=r"E:\epstein-files/data")
    parser.add_argument("-b", "--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("-w", "--workers", type=int, default=os.cpu_count() or 4)
    args = parser.parse_args()

    input_path = Path(args.input_archive)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if not input_path.exists():
        sys.exit(f"File not found: {input_path}")

    # Initialize stats
    stats = {
        "processed": 0,
        "errors": 0,
        "saved_bytes": 0,  # Bytes of input files where binary was discarded (PDFs/Text)
        "buffered_bytes": 0,
        "last_update": 0,
    }
    batch_buffer = []
    futures = set()

    # --- Helper: Handle Result & Update Stats ---
    def handle_result(results):
        if not isinstance(results, list):
            results = [results]

        source_processed = False
        for res in results:
            batch_buffer.append(res)

            if res["error"]:
                stats["errors"] += 1

            # Buffer size tracking (approx output size)
            if res["text_content"]:
                stats["buffered_bytes"] += len(res["text_content"])
            if res["image"]:
                stats["buffered_bytes"] += len(res["image"]["bytes"])

            # Saved Bytes tracking (Input size of PDF/Text files)
            # We only count the size once per source file (the one with 'source' in metadata or no 'source')
            try:
                if res["metadata"]:
                    meta = json.loads(res["metadata"])
                    if not source_processed:
                        # For PDF/Text, we count the source file size once
                        if "source" in meta or res["file_type"] in ["pdf", "data"]:
                            stats["saved_bytes"] += meta.get("size", 0)
                            source_processed = True
            except Exception:
                pass

        stats["processed"] += 1  # We count source files as 'processed'

    # --- Progress Display Helper ---
    def check_print_progress(force=False):
        if force or (stats["processed"] - stats["last_update"] >= 10):
            saved_mb = stats["saved_bytes"] / (1024 * 1024)
            buffered_mb = stats["buffered_bytes"] / (1024 * 1024)
            print(
                f"Processed: {stats['processed']} | "
                f"Errors: {stats['errors']} | "
                f"Buffered: {buffered_mb:.1f} MB | "
                f"Saved: {saved_mb:.1f} MB",
                end="\r",
            )
            stats["last_update"] = stats["processed"]

    writer = ChunkedParquetWriter(
        output_dir, SCHEMA, target_file_size=500 * 1024 * 1024  # 500 MB chunks
    )

    logger.info(f"Starting processing with {args.workers} workers...")
    logger.info(f"Flow control limit: {MAX_PENDING_FUTURES} pending tasks")

    try:
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            dctx = zstd.ZstdDecompressor()

            with open(input_path, "rb") as fh, dctx.stream_reader(
                fh, read_size=1024 * 1024 * 16
            ) as reader, tarfile.open(
                fileobj=reader, mode="r|", bufsize=1024 * 1024 * 16
            ) as tar:

                for member in tar:
                    if not member.isfile():
                        continue

                    # -- 1. Fast Filtering --
                    path_parts = member.name.split("/")
                    if "metadata" in path_parts:
                        continue
                    if not ("pdfs" in path_parts or "media" in path_parts):
                        continue

                    ds_num = parse_dataset_number(path_parts)
                    if ds_num == -1:
                        continue

                    filename = path_parts[-1]
                    if filename.startswith(".") or filename == "checksums.csv":
                        continue

                    # -- 2. Determine Type & Read --
                    f_type = determine_type(filename)
                    if member.size > MAX_SINGLE_FILE_BYTES:
                        logger.warning(f"Skipping {filename} (>2GB)")
                        continue

                    try:
                        content = tar.extractfile(member).read()
                    except Exception as read_err:
                        logger.error(f"Read error {filename}: {read_err}")
                        continue

                    # -- 3. Dispatch Logic --

                    # CASE A: Media -> Process Locally (Keep binary, don't count as 'Saved')
                    if f_type in ["audio", "image", "video", "other"]:
                        media_data = (
                            {"bytes": content, "path": filename}
                            if f_type != "other"
                            else None
                        )
                        record = {
                            "dataset_id": ds_num,
                            "doc_id": os.path.splitext(filename)[0],
                            "file_name": filename,
                            "file_type": f_type,
                            "online_url": get_online_url(ds_num, filename),
                            "text_content": None,
                            "audio": media_data if f_type == "audio" else None,
                            "image": media_data if f_type == "image" else None,
                            "video": media_data if f_type == "video" else None,
                            "metadata": json.dumps(
                                {"path": member.name, "size": member.size}
                            ),
                            "error": None,
                        }

                        batch_buffer.append(record)
                        stats["processed"] += 1
                        stats["buffered_bytes"] += member.size
                        # Note: We do NOT increment saved_bytes here, as we kept the binary.

                    # CASE B: Text/PDF -> Send to Worker (Discard binary, count as 'Saved')
                    elif f_type in ["pdf", "data"]:
                        future = executor.submit(
                            worker_process_content,
                            content,
                            filename,
                            path_parts,
                            f_type,
                            member.size,
                            ds_num,
                        )
                        futures.add(future)

                    # -- 4. Flow Control & Harvesting --

                    # If queue is full, wait blockingly
                    while len(futures) >= MAX_PENDING_FUTURES:
                        done, _ = wait(futures, return_when=FIRST_COMPLETED)
                        for f in done:
                            futures.remove(f)
                            try:
                                handle_result(f.result())
                            except Exception as e:
                                logger.error(f"Future error: {e}")
                        check_print_progress(force=True)

                    # Opportunistic reaping (non-blocking)
                    done_now = [f for f in futures if f.done()]
                    for f in done_now:
                        futures.remove(f)
                        try:
                            handle_result(f.result())
                        except Exception as e:
                            logger.error(f"Future error: {e}")

                    # -- 5. Flush & Progress --
                    if (
                        len(batch_buffer) >= args.batch_size
                        or stats["buffered_bytes"] >= MAX_BATCH_BYTES
                    ):
                        flush_batch(writer, batch_buffer)
                        stats["buffered_bytes"] = 0

                    check_print_progress()

            # End of loop: Wait for remaining
            for f in as_completed(futures):
                try:
                    handle_result(f.result())
                except Exception as e:
                    logger.error(f"Final future error: {e}")
                check_print_progress(force=True)

            flush_batch(writer, batch_buffer)

    except KeyboardInterrupt:
        logger.warning("\nInterrupted! Flushing buffer...")
        flush_batch(writer, batch_buffer)
    finally:
        if writer:
            writer.close()
        saved_gb = stats["saved_bytes"] / (1024**3)
        print()
        logger.info(
            f"Done. Processed: {stats['processed']} | Errors: {stats['errors']} | Saved: {saved_gb:.2f} GB"
        )


if __name__ == "__main__":
    main()