"""Pipeline configuration: paths, queries, thresholds."""

import logging
import os
from pathlib import Path

logger = logging.getLogger(__name__)

# ── Directory layout ──────────────────────────────────────────────────────────
ROOT = Path(__file__).resolve().parent.parent  # datasets/checkbox/
PIPELINE_DIR = ROOT / "pipeline"
IMAGES_DIR = ROOT / "images"
ANNOTATIONS_DIR = ROOT / "annotations"
RAW_DIR = ROOT / "raw"
PDFS_DIR = RAW_DIR / "pdfs"
PAGE_ANNOTS_DIR = RAW_DIR / "page_annotations"
METADATA_DIR = ROOT / "metadata"

DOCUMENTS_JSONL = METADATA_DIR / "documents.jsonl"
STATE_JSON = METADATA_DIR / "state.json"

# ── Search queries ────────────────────────────────────────────────────────────
# High-precision checkbox-instruction queries
CHECKBOX_QUERIES = [
    '"check all that apply"',
    '"check the box" OR "check the appropriate box"',
    '"mark with an X"',
    '"yes no" "check"',
    '"check one" OR "check only one"',
    '"select all" "apply"',
]

# Form-type diversity queries
FORM_TYPE_QUERIES = [
    '"incident report" "check"',
    '"inspection" "checklist"',
    '"permit application" "check"',
    '"complaint form" "check"',
    '"enrollment form" "check"',
    '"medical history" "check"',
    '"voter registration"',
    '"tax" "check the box"',
]

ALL_QUERIES = CHECKBOX_QUERIES + FORM_TYPE_QUERIES

# ── Search filters ────────────────────────────────────────────────────────────
MAX_PAGE_COUNT = 20
MIN_PAGE_COUNT = 1
MAX_DOCS_PER_QUERY = 500
MAX_DOCS_PER_ORG = 30

# ── Text pre-filter signals ──────────────────────────────────────────────────
CHECKBOX_SIGNALS = [
    "check all that apply",
    "check the box",
    "check one",
    "mark with an x",
    "yes  no",
    "yes / no",
    "yes no",
    "\u2610",  # ☐ ballot box
    "\u2611",  # ☑ ballot box with check
    "\u2612",  # ☒ ballot box with X
    "\u25a1",  # □ white square
    "\u25a0",  # ■ black square
    "( )",
    "[x]",
    "[ ]",
]

MIN_TEXT_SCORE = 1

# ── Structural filter ────────────────────────────────────────────────────────
MIN_STRUCTURAL_SCORE = 1

# Vector checkbox shape criteria (from VectorCheckboxOptions defaults)
CHECKBOX_MIN_SIZE = 6.0   # PDF points
CHECKBOX_MAX_SIZE = 25.0  # PDF points
CHECKBOX_MAX_ASPECT_RATIO = 1.5

# ── Annotation ────────────────────────────────────────────────────────────────
GEMINI_MODEL = "gemini-3-flash-preview"
GEMINI_RPM_LIMIT = 2000  # Paid tier
MAX_IMAGE_LONGEST_SIDE = 1000  # Resize rendered pages so longest side = this
COCO_BBOX_PADDING = 3  # Pixels to pad each side of COCO bboxes (Gemini boxes are slightly tight)
ANNOTATION_PROMPT_TEMPLATE = """\
You are annotating a document page image for a checkbox detection dataset.

Detect all checkboxes, tick boxes, check marks, selection circles, and radio buttons.
Include both filled/checked and empty/unchecked ones.

For each one, report:
- label: "checkbox_checked" or "checkbox_unchecked"
- box_2d: [ymin, xmin, ymax, xmax] normalized to 0-1000

Rules:
- box_2d coordinates are normalized 0-1000 (0=top/left, 1000=bottom/right)
- "checkbox_checked" = any mark inside (checkmark, X, fill, circle, etc.)
- "checkbox_unchecked" = empty box/circle
- Do NOT include bullet points, list markers, or table cell borders
"""

# ── Validation ────────────────────────────────────────────────────────────────
MIN_BBOX_SIZE_PX = 5
MAX_BBOX_SIZE_PX = 150  # Relative to max-1000px resized images
MAX_BBOX_ASPECT_RATIO = 3.0
DUPLICATE_IOU_THRESHOLD = 0.9
SPOT_CHECK_FRACTION = 0.10  # Cross-engine check on 10% of pages

# ── COCO categories ──────────────────────────────────────────────────────────
COCO_CATEGORIES = [
    {"id": 1, "name": "checkbox_checked", "supercategory": "checkbox"},
    {"id": 2, "name": "checkbox_unchecked", "supercategory": "checkbox"},
]

# ── Train/Val split ──────────────────────────────────────────────────────────
TRAIN_FRACTION = 0.8

# ── Download rate limiting ───────────────────────────────────────────────────
DOWNLOAD_RATE_LIMIT = 10  # requests per second

# ── Training (YOLO) ─────────────────────────────────────────────────────────
YOLO_DIR = ROOT / "yolo"                     # YOLO-format dataset output
MODEL_DIR = ROOT / "model"                   # Final exported models
YOLO_BASE_MODEL = "yolo12n.pt"               # Pretrained nano model
YOLO_IMGSZ = 1024                            # Input image size (larger for small checkboxes)
YOLO_EPOCHS = 200                            # Max epochs (early stopping applies)
YOLO_BATCH = 8                               # Reduced for imgsz=1024 memory
YOLO_PATIENCE = 50                           # Early stopping patience
YOLO_CLS_LOSS = 1.5                          # Classification loss weight (boosted for imbalance)
YOLO_SAVE_PERIOD = 25                        # Checkpoint interval
ONNX_OPSET = 18                              # ONNX opset version

# Oversampling: duplicate images containing checked boxes N times
CHECKED_OVERSAMPLE = 6
ONNX_MODEL_NAME = "checkbox_yolo12n.onnx"    # Output ONNX filename

# COCO category ID → YOLO class index mapping
# COCO: 1=checkbox_checked, 2=checkbox_unchecked
# YOLO: 0=checkbox_checked, 1=checkbox_unchecked
COCO_TO_YOLO_CLASS = {1: 0, 2: 1}
YOLO_CLASS_NAMES = ["checkbox_checked", "checkbox_unchecked"]


# ── DocumentCloud client ────────────────────────────────────────────────────
def get_dc_client():
    """Create a DocumentCloud client, authenticated if credentials are available."""
    try:
        import documentcloud
    except ImportError:
        raise ImportError(
            "python-documentcloud is required. "
            "Install with: pip install python-documentcloud"
        )

    username = os.environ.get("DOCUMENTCLOUD_USERNAME", "js4571@columbia.edu")
    password = os.environ.get("DOCUMENTCLOUD_PASSWORD", "!Nuy8s72zEWTvJm")

    if username and password:
        try:
            client = documentcloud.DocumentCloud(username, password)
            logger.info("Using authenticated DocumentCloud access (%s)", username)
            return client
        except Exception as e:
            logger.warning("Authenticated login failed (%s), falling back to anonymous", e)

    logger.info("Using anonymous DocumentCloud access (rate-limited)")
    return documentcloud.DocumentCloud()


class RateLimitError(Exception):
    """Raised when a rate limit is hit so the pipeline bails immediately."""
    pass


def check_rate_limit(exc: Exception):
    """Re-raise as RateLimitError if exc looks like a rate limit / auth block."""
    msg = str(exc).lower()
    if any(s in msg for s in ["429", "rate limit", "too many requests", "throttl"]):
        raise RateLimitError(f"Rate limited by API — bailing: {exc}") from exc
