#!/usr/bin/env python3
"""Download BTS On-Time Performance ZIPs from transtats.bts.gov.

Fetches monthly ZIPs for the configured year range. Skips files already on disk.
"""
import os
import sys
import time
from pathlib import Path

import requests

BASE_URL = "https://transtats.bts.gov/PREZIP/"
FILE_TEMPLATE = (
    "On_Time_Reporting_Carrier_On_Time_Performance_1987_present_{year}_{month}.zip"
)


def download_month(year: int, month: int, output_dir: Path) -> Path | None:
    filename = FILE_TEMPLATE.format(year=year, month=month)
    url = BASE_URL + filename
    output_path = output_dir / filename

    if output_path.exists():
        size_mb = output_path.stat().st_size / (1024 * 1024)
        print(f"  [SKIP] {filename} already exists ({size_mb:.1f} MB)")
        return output_path

    print(f"  Downloading: {year}-{month:02d}... ", end="", flush=True)
    try:
        response = requests.get(url, stream=True, timeout=120)
        response.raise_for_status()
    except requests.RequestException as e:
        print(f"FAILED ({e})")
        return None

    # BTS redirects bad URLs to an HTML error page -- catch that
    content_type = response.headers.get("content-type", "")
    if "html" in content_type.lower():
        print("FAILED (got HTML instead of ZIP, check URL)")
        return None

    total_bytes = 0
    with open(output_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=1024 * 1024):
            f.write(chunk)
            total_bytes += len(chunk)

    size_mb = total_bytes / (1024 * 1024)
    print(f"OK ({size_mb:.1f} MB)")
    return output_path


def main():
    start_year = int(os.getenv("BTS_START_YEAR", "2023"))
    end_year = int(os.getenv("BTS_END_YEAR", "2024"))
    output_dir = Path(os.getenv("BTS_DOWNLOAD_DIR", "./data/raw"))
    output_dir.mkdir(parents=True, exist_ok=True)

    total_months = (end_year - start_year + 1) * 12
    downloaded = 0
    failed = 0

    print(f"Downloading BTS On-Time Performance data: {start_year}-{end_year}")
    print(f"Output directory: {output_dir.resolve()}")
    print(f"Total months to download: {total_months}")
    print()

    for year in range(start_year, end_year + 1):
        print(f"Year {year}:")
        for month in range(1, 13):
            result = download_month(year, month, output_dir)
            if result: 
                downloaded += 1
            else:
                failed += 1

            # Be polite to the server
            time.sleep(2)

        print()

    print(f"Download complete: {downloaded} succeeded, {failed} failed")
    print(f"Files saved to: {output_dir.resolve()}")

    if failed > 0:
        print("\nRe-run this script to retry failed downloads.")
        sys.exit(1)


if __name__ == "__main__":
    main()
