#!/usr/bin/env python3
"""CLI orchestrator for the checkbox dataset acquisition pipeline.

Usage:
    python -m datasets.checkbox.pipeline.run [steps...] [options]

Steps (in order):
    search          Search DocumentCloud for candidate documents
    prefilter       Screen pages by text signals
    download        Download PDFs
    structural      Structural pre-filter with natural-pdf
    annotate        Gemini annotation of checkbox locations
    validate        Automated QC checks
    build           Build COCO dataset
    train           Train YOLO12 detector and export ONNX
    visualize       Generate spot-check visualizations
    all             Run all steps in order

Options:
    --resume        Resume from last checkpoint (default behavior)
    --stats         Show pipeline statistics and exit
    --viz-count N   Number of pages to visualize (default: 10)
"""

import argparse
import logging
import sys
from pathlib import Path

from .state import PipelineState

logger = logging.getLogger(__name__)

STEP_ORDER = [
    "search",
    "prefilter",
    "download",
    "structural",
    "annotate",
    "validate",
    "build",
    "train",
    "visualize",
]


def print_stats(state: PipelineState):
    """Print pipeline progress statistics."""
    summary = state.summary
    print("\n=== Pipeline Status ===")
    print(f"  Documents searched:         {summary.get('searched', 0)}")
    print(f"  Documents pre-filtered:     {summary.get('prefiltered', 0)}")
    print(f"  Documents downloaded:       {summary.get('downloaded', 0)}")
    print(f"  Documents struct-filtered:  {summary.get('structural_filtered', 0)}")
    print(f"  Documents annotated:        {summary.get('annotated', 0)}")
    print(f"  Documents validated:        {summary.get('validated', 0)}")
    print()

    # Count total pages at each stage
    total_candidate = 0
    total_qualifying = 0
    total_annotated = 0
    for doc_id in state.all_doc_ids():
        total_candidate += len(state.get(doc_id, "candidate_pages", []))
        total_qualifying += len(state.get(doc_id, "qualifying_pages", []))
        total_annotated += len(state.get(doc_id, "annotated_pages", []))

    print(f"  Total candidate pages:      {total_candidate}")
    print(f"  Total qualifying pages:     {total_qualifying}")
    print(f"  Total annotated pages:      {total_annotated}")
    print()


def run_step(step: str, state: PipelineState, args: argparse.Namespace):
    """Run a single pipeline step."""
    print(f"\n{'='*60}")
    print(f"  Step: {step}")
    print(f"{'='*60}\n")

    if step == "search":
        from .search import search
        count = search(state)
        print(f"  Found {count} new documents")

    elif step == "prefilter":
        from .prefilter import prefilter
        results = prefilter(state)
        total_pages = sum(len(p) for p in results.values())
        print(f"  {len(results)} docs with {total_pages} candidate pages")

    elif step == "download":
        from .download import download
        count = download(state)
        print(f"  Downloaded {count} new PDFs")

    elif step == "structural":
        from .structural_filter import structural_filter
        results = structural_filter(state)
        total_pages = sum(len(p) for p in results.values())
        print(f"  {len(results)} docs with {total_pages} qualifying pages")

    elif step == "annotate":
        from .annotate import annotate
        count = annotate(state)
        print(f"  Annotated {count} new pages")

    elif step == "validate":
        from .validate import validate
        stats = validate(state)
        print(f"  {stats.get('total_checkboxes_valid', 0)}/{stats.get('total_checkboxes_in', 0)} "
              f"checkboxes valid across {stats.get('total_pages', 0)} pages")

    elif step == "build":
        from .build_coco import build_coco
        stats = build_coco(state)
        for k, v in stats.items():
            print(f"  {k}: {v}")

    elif step == "train":
        from .train import train
        stats = train()
        for k, v in stats.items():
            print(f"  {k}: {v}")

    elif step == "visualize":
        from .visualize import visualize_random
        count = getattr(args, "viz_count", 10)
        save_dir = Path("temp/checkbox_viz")
        results = visualize_random(state, count=count, save_dir=save_dir)
        print(f"  Saved {len(results)} visualizations to {save_dir}")

    else:
        print(f"  Unknown step: {step}")
        return False

    return True


def main():
    from dotenv import load_dotenv
    load_dotenv()

    parser = argparse.ArgumentParser(
        description="Checkbox dataset acquisition pipeline",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "steps",
        nargs="*",
        default=["all"],
        help="Pipeline steps to run (default: all)",
    )
    parser.add_argument(
        "--stats",
        action="store_true",
        help="Show pipeline statistics and exit",
    )
    parser.add_argument(
        "--viz-count",
        type=int,
        default=10,
        dest="viz_count",
        help="Number of pages to visualize (default: 10)",
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Enable debug logging",
    )

    args = parser.parse_args()

    # Configure logging
    level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
        datefmt="%H:%M:%S",
    )

    state = PipelineState()

    if args.stats:
        print_stats(state)
        return

    # Resolve steps
    steps = []
    for s in args.steps:
        if s == "all":
            steps = STEP_ORDER[:]
            break
        elif s in STEP_ORDER:
            steps.append(s)
        else:
            print(f"Unknown step: {s}")
            print(f"Available steps: {', '.join(STEP_ORDER)}")
            sys.exit(1)

    if not steps:
        print("No steps specified. Use 'all' to run the full pipeline.")
        parser.print_help()
        sys.exit(1)

    print(f"Running steps: {', '.join(steps)}")
    print_stats(state)

    for step in steps:
        try:
            success = run_step(step, state, args)
            state.flush()  # Ensure all state is written after each step
            if not success:
                print(f"\nStep '{step}' failed. Stopping.")
                sys.exit(1)
        except KeyboardInterrupt:
            state.flush()
            print(f"\n\nInterrupted during step '{step}'. Progress saved.")
            print_stats(state)
            sys.exit(130)
        except Exception as e:
            state.flush()
            logger.exception("Step '%s' failed with error", step)
            print(f"\nStep '{step}' failed: {e}")
            print("Progress up to this point has been saved. Re-run to resume.")
            sys.exit(1)

    print("\n" + "=" * 60)
    print("  Pipeline complete!")
    print("=" * 60)
    print_stats(state)


if __name__ == "__main__":
    main()
