"""Step 3: Download PDFs for pre-filtered documents."""

import logging
import time
from pathlib import Path
from typing import List

import requests

from .config import DOWNLOAD_RATE_LIMIT, PDFS_DIR
from .search import load_documents
from .state import PipelineState

logger = logging.getLogger(__name__)


def download(state: PipelineState) -> int:
    """Download PDFs for documents that passed pre-filter.

    Returns number of newly downloaded PDFs.
    """
    documents = load_documents()
    PDFS_DIR.mkdir(parents=True, exist_ok=True)

    downloaded = 0
    skipped = 0
    failed = 0
    min_interval = 1.0 / DOWNLOAD_RATE_LIMIT
    last_request = 0.0

    for doc_record in documents:
        doc_id = str(doc_record["id"])

        # Skip if not pre-filtered or no candidate pages
        if not state.is_done(doc_id, "prefiltered"):
            continue
        candidate_pages = state.get(doc_id, "candidate_pages", [])
        if not candidate_pages:
            continue

        # Skip if already downloaded
        pdf_path = PDFS_DIR / f"{doc_id}.pdf"
        if pdf_path.exists():
            state.mark_done(doc_id, "downloaded")
            skipped += 1
            continue

        # Rate limit
        elapsed = time.monotonic() - last_request
        if elapsed < min_interval:
            time.sleep(min_interval - elapsed)

        pdf_url = doc_record.get("pdf_url", "")
        if not pdf_url:
            # Construct URL from document ID
            pdf_url = f"https://www.documentcloud.org/documents/{doc_id}.pdf"

        max_retries = 3
        for attempt in range(max_retries):
            try:
                last_request = time.monotonic()
                resp = requests.get(pdf_url, timeout=60, stream=True)
                resp.raise_for_status()

                # Validate content type
                ctype = resp.headers.get("Content-Type", "").lower()
                if ctype and "pdf" not in ctype and "octet-stream" not in ctype:
                    logger.warning("Skipping %s: unexpected Content-Type %s", doc_id, ctype)
                    break

                # Write to temp file, rename on success to avoid partial downloads
                tmp_path = pdf_path.with_suffix(".tmp")
                with open(tmp_path, "wb") as f:
                    for chunk in resp.iter_content(chunk_size=8192):
                        f.write(chunk)

                # Basic size check — a real PDF is at least ~100 bytes
                if tmp_path.stat().st_size < 100:
                    logger.warning("Skipping %s: downloaded file too small (%d bytes)", doc_id, tmp_path.stat().st_size)
                    tmp_path.unlink(missing_ok=True)
                    break

                tmp_path.replace(pdf_path)
                state.mark_done(doc_id, "downloaded")
                downloaded += 1
                logger.debug("Downloaded %s (%s)", doc_id, doc_record.get("title", ""))
                break

            except Exception as e:
                if attempt < max_retries - 1:
                    wait = 2 ** (attempt + 1)
                    logger.warning("Download attempt %d failed for %s: %s (retry in %ds)", attempt + 1, doc_id, e, wait)
                    time.sleep(wait)
                else:
                    logger.warning("Failed to download %s after %d attempts: %s", doc_id, max_retries, e)
                    # Clean up partial temp file
                    pdf_path.with_suffix(".tmp").unlink(missing_ok=True)
                    failed += 1

    logger.info(
        "Download complete. %d new, %d skipped, %d failed.",
        downloaded, skipped, failed,
    )
    return downloaded


def get_pdf_path(doc_id: str) -> Path:
    """Return the local path for a downloaded PDF."""
    return PDFS_DIR / f"{doc_id}.pdf"


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    st = PipelineState()
    download(st)
