#!/usr/bin/env python3
"""
Downloads OSM PBF files from Geofabrik if they don't already exist.

Usage:
    python -m scripts.download_pbf ohio
    python -m scripts.download_pbf egypt india thailand
    python -m scripts.download_pbf --list
"""

import argparse
import sys
from pathlib import Path
import requests
from tqdm import tqdm

from scripts.geofabrik_registry import get_geofabrik_url, list_available_locations


def download_pbf(location: str, output_dir: Path = None, force: bool = False) -> Path:
    """
    Download a PBF file from Geofabrik if it doesn't exist.
    
    Args:
        location: Location name (e.g., 'ohio', 'egypt')
        output_dir: Directory to save the file (default: output/{location}/)
        force: Force re-download even if file exists
        
    Returns:
        Path to the downloaded file
    """
    location = location.lower()
    
    # Get URL from registry
    try:
        url = get_geofabrik_url(location)
    except ValueError as e:
        print(f"❌ {e}")
        sys.exit(1)
    
    # Determine output location
    if output_dir is None:
        base_dir = Path.cwd()
        output_dir = base_dir / "output" / location
    
    output_dir.mkdir(parents=True, exist_ok=True)
    filename = output_dir / f"{location}.osm.pbf"
    
    # Check if file already exists
    if filename.exists() and not force:
        file_size = filename.stat().st_size / (1024 * 1024)  # MB
        print(f"✅ File already exists: {filename} ({file_size:.1f} MB)")
        return filename
    
    # Download the file
    print(f"⬇️  Downloading OSM data for {location.capitalize()}...")
    print(f"🔗 URL: {url}")
    
    try:
        with requests.get(url, stream=True, timeout=120) as r:
            r.raise_for_status()
            total_size = int(r.headers.get("content-length", 0))
            block_size = 8192
            
            with open(filename, "wb") as f, tqdm(
                total=total_size,
                unit="B",
                unit_scale=True,
                unit_divisor=1024,
                desc=f"{location}.osm.pbf"
            ) as bar:
                for chunk in r.iter_content(chunk_size=block_size):
                    f.write(chunk)
                    bar.update(len(chunk))
        
        # Verify download
        if filename.stat().st_size < 1024 * 100:  # Less than 100 KB
            raise ValueError("Downloaded file seems too small — possible download error.")
        
        file_size = filename.stat().st_size / (1024 * 1024)  # MB
        print(f"✅ Download complete: {filename} ({file_size:.1f} MB)")
        return filename
        
    except requests.exceptions.RequestException as e:
        print(f"❌ Download failed: {e}")
        if filename.exists():
            filename.unlink()  # Clean up partial download
        sys.exit(1)
    except Exception as e:
        print(f"❌ Error: {e}")
        if filename.exists():
            filename.unlink()
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        description="Download OSM PBF files from Geofabrik",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  Download single location:
    python -m scripts.download_pbf ohio
    
  Download multiple locations:
    python -m scripts.download_pbf ohio kentucky new-york
    python -m scripts.download_pbf egypt india thailand
    
  Force re-download:
    python -m scripts.download_pbf ohio --force
    
  List available locations:
    python -m scripts.download_pbf --list
        """
    )
    
    parser.add_argument(
        "locations",
        nargs="*",
        help="Location names to download (e.g., ohio, egypt, india)"
    )
    
    parser.add_argument(
        "--list",
        action="store_true",
        help="List all available locations in the registry"
    )
    
    parser.add_argument(
        "--force",
        action="store_true",
        help="Force re-download even if file exists"
    )
    
    parser.add_argument(
        "--output-dir",
        type=Path,
        help="Custom output directory (default: output/{location}/)"
    )
    
    args = parser.parse_args()
    
    # Handle --list flag
    if args.list:
        list_available_locations()
        return
    
    # Require at least one location
    if not args.locations:
        parser.print_help()
        print("\n❌ Error: Please provide at least one location name")
        sys.exit(1)
    
    # Download each location
    print(f"\n📥 Downloading {len(args.locations)} location(s)...\n")
    
    downloaded = []
    failed = []
    
    for location in args.locations:
        try:
            filepath = download_pbf(
                location,
                output_dir=args.output_dir,
                force=args.force
            )
            downloaded.append((location, filepath))
            print()  # Add spacing between downloads
        except Exception as e:
            print(f"❌ Failed to download {location}: {e}\n")
            failed.append(location)
    
    # Summary
    print("=" * 60)
    print(f"✅ Successfully downloaded: {len(downloaded)}/{len(args.locations)}")
    
    if downloaded:
        print("\n📁 Downloaded files:")
        for location, filepath in downloaded:
            file_size = filepath.stat().st_size / (1024 * 1024)
            print(f"  • {location}: {filepath} ({file_size:.1f} MB)")
    
    if failed:
        print(f"\n❌ Failed: {len(failed)}")
        for location in failed:
            print(f"  • {location}")
    
    print("=" * 60)


if __name__ == "__main__":
    main()