#!/usr/bin/env python3
"""
Bulk Resume Uploader
--------------------

Batch local resumes into the TalentPrism bulk parsing API.

The script scans a directory of individual resume files, builds the manifest
CSV and ZIP in memory, uploads row- and ZIP-size-limited batches, and polls
each upload until it reaches a terminal status. It writes a checkpoint file
named ``.bulk_resume_uploader_checkpoint.json`` in the resume directory so a
crashed run can be resumed safely.

Usage:
    export TP_API_KEY='Api-Key sk_...'

    # Preview deterministic batch numbers without uploading.
    python bulk_resume_uploader.py --api-key="$TP_API_KEY" --directory=/path/to/resumes --dry-run

    # Start or resume a run. Reusing --resume skips completed checkpointed batches.
    python bulk_resume_uploader.py --api-key="$TP_API_KEY" --directory=/path/to/resumes --resume

    # If a timed-out POST later appears as a completed upload, continue at the next batch.
    python bulk_resume_uploader.py --api-key="$TP_API_KEY" --directory=/path/to/resumes --start-batch=4 --resume
"""

import argparse
import csv
import io
import json
import sys
import time
import zipfile
from pathlib import Path
from typing import Iterable, List

import requests


MAX_ROWS_PER_BATCH = 200
MAX_ZIP_BYTES_PER_BATCH = 20 * 1024 * 1024
DEFAULT_RATE_LIMIT_UPLOADS = 10
SUPPORTED_SUFFIXES = {".pdf", ".doc", ".docx", ".png", ".jpg", ".jpeg", ".webp", ".tif", ".tiff"}
ACTIVE_STATUSES = {"queued", "processing"}
TERMINAL_STATUSES = {"completed", "completed_with_errors", "failed"}


def format_bytes(size: int) -> str:
    if size < 1024:
        return f"{size} B"
    if size < 1024 * 1024:
        return f"{size / 1024:.1f} KB"
    return f"{size / (1024 * 1024):.1f} MB"


def estimate_zip_entry_size(path: Path) -> int:
    # PDFs and DOCX files are usually already compressed, so raw bytes are a close upper estimate.
    return path.stat().st_size + len(path.name.encode("utf-8")) + 256


def estimate_batch_size(rows: List[Path]) -> int:
    return sum(estimate_zip_entry_size(path) for path in rows)


def chunked_by_limits(files: List[Path], max_rows: int, max_zip_bytes: int) -> Iterable[List[Path]]:
    batch: List[Path] = []
    estimated_size = 0

    for path in files:
        entry_size = estimate_zip_entry_size(path)
        if entry_size > max_zip_bytes:
            raise ValueError(
                f"{path} is too large for one upload. Max ZIP size is {format_bytes(max_zip_bytes)}; "
                f"estimated file contribution is {format_bytes(entry_size)}."
            )

        would_exceed_rows = len(batch) >= max_rows
        would_exceed_bytes = batch and estimated_size + entry_size > max_zip_bytes
        if would_exceed_rows or would_exceed_bytes:
            yield batch
            batch = []
            estimated_size = 0

        batch.append(path)
        estimated_size += entry_size

    if batch:
        yield batch


def build_manifest(rows: List[Path]) -> bytes:
    buffer = io.StringIO()
    writer = csv.DictWriter(buffer, fieldnames=["resume_url", "zip_path", "first_name", "last_name", "email", "phone"])
    writer.writeheader()
    for path in rows:
        writer.writerow({"resume_url": "", "zip_path": path.name})
    return buffer.getvalue().encode("utf-8")


def build_zip(rows: List[Path]) -> bytes:
    memfile = io.BytesIO()
    with zipfile.ZipFile(memfile, "w", compression=zipfile.ZIP_DEFLATED) as archive:
        for path in rows:
            archive.write(path, arcname=path.name)
    memfile.seek(0)
    return memfile.read()


def load_checkpoint(path: Path) -> dict:
    if not path.exists():
        return {"version": 1, "batches": {}}
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise ValueError(f"Checkpoint file is invalid JSON: {path}") from exc


def save_checkpoint(path: Path, checkpoint: dict):
    path.write_text(json.dumps(checkpoint, indent=2, sort_keys=True), encoding="utf-8")


def batch_file_names(rows: List[Path]) -> list[str]:
    return [path.name for path in rows]


def checkpoint_batch(args, batch_index: int, rows: List[Path], **updates):
    checkpoint = load_checkpoint(args.checkpoint_path)
    checkpoint.setdefault("version", 1)
    checkpoint["directory"] = str(args.directory_path)
    checkpoint["base_url"] = args.base_url.rstrip("/")
    checkpoint["max_rows"] = args.max_rows
    checkpoint["max_zip_bytes"] = args.max_zip_bytes

    batches = checkpoint.setdefault("batches", {})
    record = batches.setdefault(str(batch_index), {})
    record["file_names"] = batch_file_names(rows)
    record.update(updates)
    save_checkpoint(args.checkpoint_path, checkpoint)


def find_checkpoint_record(args, batch_index: int, rows: List[Path]) -> dict | None:
    record = load_checkpoint(args.checkpoint_path).get("batches", {}).get(str(batch_index))
    if not record:
        return None
    if record.get("file_names") != batch_file_names(rows):
        print(f"[batch {batch_index}] Checkpoint file list no longer matches. Treating this as a new batch.")
        return None
    return record


def upload_batch(batch_index: int, rows: List[Path], args, headers):
    manifest_bytes = build_manifest(rows)
    zip_bytes = build_zip(rows)
    if len(zip_bytes) > args.max_zip_bytes:
        raise ValueError(
            f"Built ZIP is {format_bytes(len(zip_bytes))}, above the configured "
            f"{format_bytes(args.max_zip_bytes)} limit. Retry with fewer files."
        )

    files = {
        "manifest": ("manifest.csv", manifest_bytes, "text/csv"),
        "resumes_zip": ("resumes.zip", zip_bytes, "application/zip"),
    }
    data = {}
    if args.schema_key:
        data["schema_key"] = args.schema_key
    if args.source:
        data["source"] = args.source

    url = f"{args.base_url.rstrip('/')}/candidate/api/bulk-parse-resumes/"
    print(f"[batch {batch_index}] Uploading {len(rows)} resumes ({format_bytes(len(zip_bytes))}) -> {url}")
    # Increased timeout to 300s (5 minutes) as bulk uploads with files can be slow
    try:
        response = requests.post(url, headers=headers, files=files, data=data, timeout=300)
    except requests.exceptions.ReadTimeout:
        print(f"[batch {batch_index}] ERROR: Upload timed out after 300s. The server might still be processing it.")
        print(f"[batch {batch_index}] Check the TalentPrism dashboard or retry with fewer files.")
        raise

    if response.status_code != 202:
        print(f"[batch {batch_index}] ERROR {response.status_code}: {response.text}")
        response.raise_for_status()
    payload = response.json()
    upload_id = payload["upload_id"]
    print(f"[batch {batch_index}] Queued as upload_id {upload_id}")
    checkpoint_batch(
        args,
        batch_index,
        rows,
        upload_id=upload_id,
        status="queued",
        processed_rows=0,
        total_rows=len(rows),
        success_count=0,
        failure_count=0,
    )
    wait_for_completion(upload_id, batch_index, rows, args, headers)


def wait_for_completion(upload_id: int, batch_index: int, rows: List[Path], args, headers):
    status_url = f"{args.base_url.rstrip('/')}/candidate/api/bulk-parse-resumes/{upload_id}/"
    retry_count = 0
    max_retries = 5
    poll_error_count = 0
    max_poll_errors = 10
    backoff = args.poll_interval

    while True:
        try:
            resp = requests.get(status_url, headers=headers, timeout=30)
        except requests.exceptions.RequestException as exc:
            poll_error_count += 1
            if poll_error_count > max_poll_errors:
                print(f"[batch {batch_index}] Polling failed too many times. Last error: {exc}")
                raise
            wait_seconds = min(backoff, 300)
            print(
                f"[batch {batch_index}] Polling connection error: {exc}. "
                f"Retrying in {wait_seconds}s ({poll_error_count}/{max_poll_errors})..."
            )
            time.sleep(wait_seconds)
            backoff *= 2
            continue

        if resp.status_code == 429:
            retry_count += 1
            if retry_count > max_retries:
                print(f"[batch {batch_index}] Rate limit hit too many times. Aborting poll.")
                resp.raise_for_status()

            wait_seconds = int(resp.headers.get("Retry-After", backoff))
            print(
                f"[batch {batch_index}] Rate limited. Waiting {wait_seconds}s before retry ({retry_count}/{max_retries})..."
            )
            time.sleep(wait_seconds)
            # Increase backoff for next time just in case
            backoff *= 2
            continue

        if resp.status_code != 200:
            print(f"[batch {batch_index}] status poll failed ({resp.status_code}): {resp.text}")
            resp.raise_for_status()

        # Reset retry count on success
        retry_count = 0
        poll_error_count = 0
        data = resp.json()
        print(
            f"[batch {batch_index}] status={data['status']} "
            f"processed={data['processed_rows']}/{data['total_rows']} "
            f"success={data['success_count']} failure={data['failure_count']}"
        )
        checkpoint_batch(
            args,
            batch_index,
            rows,
            upload_id=upload_id,
            status=data["status"],
            processed_rows=data["processed_rows"],
            total_rows=data["total_rows"],
            success_count=data["success_count"],
            failure_count=data["failure_count"],
        )
        if data["status"] not in ACTIVE_STATUSES:
            if data["failure_count"]:
                for row in data["rows"]:
                    if row["status"] == "failed":
                        print(f"  - Row {row['row_number']} failed: {row['error_message']}")
            break
        time.sleep(args.poll_interval)


def main():
    parser = argparse.ArgumentParser(
        description="Upload local resumes in row- and ZIP-size-limited batches to TalentPrism."
    )
    parser.add_argument("--directory", required=True, help="Directory containing resume files.")
    parser.add_argument("--api-key", required=True, help="TalentPrism API key (format: Api-Key ...).")
    parser.add_argument("--base-url", default="https://talentprism.ai", help="API base URL.")
    parser.add_argument("--schema-key", help="Override resume parsing schema.")
    parser.add_argument("--source", default="bulk-script", help="Tracking label stored with the upload.")
    parser.add_argument("--poll-interval", type=int, default=60, help="Seconds between status polls.")
    parser.add_argument("--max-uploads", type=int, default=DEFAULT_RATE_LIMIT_UPLOADS, help="Safety cap per run.")
    parser.add_argument("--max-rows", type=int, default=MAX_ROWS_PER_BATCH, help="Maximum resumes per upload.")
    parser.add_argument("--start-batch", type=int, default=1, help="Skip batches before this number.")
    parser.add_argument(
        "--resume", action="store_true", help="Resume from the checkpoint file and skip completed batches."
    )
    parser.add_argument("--dry-run", action="store_true", help="Print computed batches without uploading anything.")
    parser.add_argument(
        "--checkpoint",
        help="Path to checkpoint JSON. Defaults to .bulk_resume_uploader_checkpoint.json in the resume directory.",
    )
    parser.add_argument(
        "--max-zip-mb",
        type=float,
        default=MAX_ZIP_BYTES_PER_BATCH / (1024 * 1024),
        help="Maximum ZIP size per upload in MB.",
    )

    args = parser.parse_args()
    if args.start_batch <= 0:
        sys.exit("--start-batch must be greater than 0.")
    args.max_zip_bytes = int(args.max_zip_mb * 1024 * 1024)
    directory = Path(args.directory).expanduser().resolve()
    if not directory.is_dir():
        sys.exit(f"Directory not found: {directory}")
    args.directory_path = directory
    args.checkpoint_path = (
        Path(args.checkpoint).expanduser().resolve()
        if args.checkpoint
        else directory / ".bulk_resume_uploader_checkpoint.json"
    )

    all_files = sorted([path for path in directory.iterdir() if path.is_file()])
    files = [path for path in all_files if path.suffix.lower() in SUPPORTED_SUFFIXES]
    ignored_files = [path for path in all_files if path.suffix.lower() not in SUPPORTED_SUFFIXES]
    if not files:
        sys.exit("No files found in the directory.")
    if ignored_files:
        print(f"Ignoring {len(ignored_files)} unsupported/non-resume file(s), including: {ignored_files[0].name}")
    print(f"Using checkpoint: {args.checkpoint_path}")

    batches = list(chunked_by_limits(files, args.max_rows, args.max_zip_bytes))
    if args.dry_run:
        print(f"Found {len(files)} resume file(s) across {len(batches)} batch(es).")
        for batch_index, batch in enumerate(batches, start=1):
            print(
                f"[batch {batch_index}] count={len(batch)} "
                f"estimated_zip={format_bytes(estimate_batch_size(batch))} "
                f"first={batch[0].name} last={batch[-1].name}"
            )
        return

    headers = {"Authorization": args.api_key}
    upload_count = 0

    for batch_index, batch in enumerate(batches, start=1):
        if batch_index < args.start_batch:
            print(f"[batch {batch_index}] Skipping because --start-batch={args.start_batch}.")
            continue

        if args.resume:
            record = find_checkpoint_record(args, batch_index, batch)
            if record:
                status = record.get("status")
                upload_id = record.get("upload_id")
                if status in TERMINAL_STATUSES:
                    print(f"[batch {batch_index}] Skipping checkpointed {status} upload_id {upload_id}.")
                    continue
                if upload_id and status in ACTIVE_STATUSES:
                    print(f"[batch {batch_index}] Resuming status polling for upload_id {upload_id}.")
                    wait_for_completion(upload_id, batch_index, batch, args, headers)
                    continue

        if upload_count >= args.max_uploads:
            print("Hit max uploads for this run; stopping to respect rate limits.")
            break
        try:
            upload_batch(batch_index, batch, args, headers)
            upload_count += 1
        except Exception as exc:
            print(f"[batch {batch_index}] Aborting due to error: {exc}")
            raise
        time.sleep(args.poll_interval)  # small pause before next batch

    print(f"Done. Submitted {upload_count} upload(s).")


if __name__ == "__main__":
    main()