putnam 0.1.0

# download_tests.py — CORRECTED FINAL VERSION (Dec 2025) — Matches exact 8-column table
import os
import re
import requests
from bs4 import BeautifulSoup
from urllib.parse import urljoin, urlparse

BASE_URL = "https://kskedlaya.org/putnam-archive/"
DOWNLOAD_ROOT = "putnam_downloads"

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
                  "AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/131.0 Safari/537.36"
}

def get_year_from_cell(cell):
    text = cell.get_text(strip=True)
    match = re.search(r'\b(19|20)\d{2}\b', text)
    return match.group(0) if match else None

def parse_table():
    print("Fetching and parsing the Putnam archive page...")
    r = requests.get(BASE_URL, headers=HEADERS, timeout=30)
    r.raise_for_status()
    soup = BeautifulSoup(r.text, "html.parser")

    # Grab the first (and only) table
    table = soup.find("table")
    if not table:
        raise RuntimeError("Could not find the main table")

    downloads = []
    rows = table.find_all("tr")

    for row_idx, row in enumerate(rows):
        cells = row.find_all(["td", "th"])  # Handle both td and th (header)
        if len(cells) < 8:
            continue  # Skip incomplete rows (including header)

        # First cell = year
        year = get_year_from_cell(cells[0])
        if not year:
            continue

        # Column 1: Problems TeX
        cell = cells[1]
        for a in cell.find_all("a", href=True):
            href = a["href"].strip()
            if not re.search(r'\.(pdf|tex|html?)$', href, re.I):
                continue
            url = urljoin(BASE_URL, href)
            ext = os.path.splitext(href)[1].lower()
            filename = f"{year}{ext}"
            downloads.append({
                "year": year,
                "category": "Problems",
                "url": url,
                "filename": filename
            })

        # Column 2: Problems PDF
        cell = cells[2]
        for a in cell.find_all("a", href=True):
            href = a["href"].strip()
            if not re.search(r'\.(pdf|tex|html?)$', href, re.I):
                continue
            url = urljoin(BASE_URL, href)
            ext = os.path.splitext(href)[1].lower()
            filename = f"{year}{ext}"
            downloads.append({
                "year": year,
                "category": "Problems",
                "url": url,
                "filename": filename
            })

        # Column 3: Solutions TeX
        cell = cells[3]
        for a in cell.find_all("a", href=True):
            href = a["href"].strip()
            if not re.search(r'\.(pdf|tex|html?)$', href, re.I):
                continue
            url = urljoin(BASE_URL, href)
            ext = os.path.splitext(href)[1].lower()
            filename = f"{year}s{ext}"
            downloads.append({
                "year": year,
                "category": "Solutions",
                "url": url,
                "filename": filename
            })

        # Column 4: Solutions PDF
        cell = cells[4]
        for a in cell.find_all("a", href=True):
            href = a["href"].strip()
            if not re.search(r'\.(pdf|tex|html?)$', href, re.I):
                continue
            url = urljoin(BASE_URL, href)
            ext = os.path.splitext(href)[1].lower()
            filename = f"{year}s{ext}"
            downloads.append({
                "year": year,
                "category": "Solutions",
                "url": url,
                "filename": filename
            })

        # Column 5: Winners HTML
        cell = cells[5]
        for a in cell.find_all("a", href=True):
            href = a["href"].strip()
            if not re.search(r'\.(pdf|tex|html?)$', href, re.I):
                continue
            url = urljoin(BASE_URL, href)
            ext = os.path.splitext(href)[1].lower()
            orig_name = os.path.basename(urlparse(url).path)
            if "announcementofwinners" in orig_name.lower():
                filename = f"AnnouncementOfWinners{year}{ext}"
            else:
                filename = f"putnam{year}results{ext}"
            downloads.append({
                "year": year,
                "category": "Winners",
                "url": url,
                "filename": filename
            })

        # Column 6: Winners PDF
        cell = cells[6]
        for a in cell.find_all("a", href=True):
            href = a["href"].strip()
            if not re.search(r'\.(pdf|tex|html?)$', href, re.I):
                continue
            url = urljoin(BASE_URL, href)
            ext = os.path.splitext(href)[1].lower()
            filename = f"AnnouncementOfWinners{year}{ext}"
            downloads.append({
                "year": year,
                "category": "Winners",
                "url": url,
                "filename": filename
            })

        # Column 7: Scores HTML
        cell = cells[7]
        for a in cell.find_all("a", href=True):
            href = a["href"].strip()
            if not re.search(r'\.(pdf|tex|html?)$', href, re.I):
                continue
            url = urljoin(BASE_URL, href)
            ext = os.path.splitext(href)[1].lower()
            filename = f"putnam{year}stats{ext}"
            downloads.append({
                "year": year,
                "category": "Scores",
                "url": url,
                "filename": filename
            })

    return downloads


def download(url, path):
    print(f"  Downloading {os.path.basename(path)}")
    try:
        r = requests.get(url, headers=HEADERS, stream=True, timeout=60)
        r.raise_for_status()
        with open(path, "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"  Saved: {os.path.basename(path)}")
    except Exception as e:
        print(f"  Failed {url}: {e}")


def main():
    os.makedirs(DOWNLOAD_ROOT, exist_ok=True)
    files = parse_table()

    print(f"Found {len(files)} files across {len(set(f['year'] for f in files))} years\n")

    downloaded_count = 0
    for item in files:
        year_dir = os.path.join(DOWNLOAD_ROOT, item["year"], item["category"])
        os.makedirs(year_dir, exist_ok=True)
        filepath = os.path.join(year_dir, item["filename"])

        if os.path.exists(filepath) and os.path.getsize(filepath) > 1000:
            print(f"  Skipping (exists): {item['filename']}")
            continue

        download(item["url"], filepath)
        downloaded_count += 1

    print(f"\nDownloaded {downloaded_count} new files. Total Putnam archive ready in {DOWNLOAD_ROOT}!")


if __name__ == "__main__":
    main()