#!/usr/bin/env python3
"""
Request and download Ohio lidar tiles based on coordinates.

Usage:
    ./request_tiles.py --coords -82.4438 40.0547 --coords -82.4440 40.0548 --output tiles/
    echo "-82.4438,40.0547" | ./request_tiles.py --output tiles/
    ./request_tiles.py --coords-file locations.txt --output tiles/

This tool:
1. Converts lon/lat coordinates to Web Mercator (EPSG:3857)
2. Queries Ohio's ArcGIS tile service to get tile metadata
3. Downloads tiles from OGRIP as .zip files
4. Outputs JSON summary of downloaded tiles to stdout
5. Respects rate limiting between requests

Output format (JSON to stdout):
{
    "requested": 5,
    "downloaded": 4,
    "skipped": 1,
    "failed": 0,
    "tiles": [
        {
            "tile_name": "BS19820747",
            "county": "LIC",
            "year": "2020",
            "block": "4",
            "url": "https://gis1.oit.ohio.gov/...",
            "output_path": "tiles/BS19820747.zip",
            "status": "downloaded"
        },
        ...
    ]
}
"""

import sys
import time
import json
import logging
import argparse
import requests
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Any

try:
    from pyproj import Transformer
except ImportError:
    print("Error: pyproj not installed. Run: pip install pyproj", file=sys.stderr)
    sys.exit(1)

# Configure logging to stderr only (stdout is for JSON output)
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    stream=sys.stderr
)
logger = logging.getLogger(__name__)

# Ohio ArcGIS tile service
TILE_SERVICE_URL = "https://maps.ohio.gov/arcgis/rest/services/OGRIP/3DepTiles/MapServer/0/query"

# OGRIP download URL template
DOWNLOAD_URL_TEMPLATE = "https://gis1.oit.ohio.gov/ZIPARCHIVES_III/ELEVATION/3DEP/LIDAR/{county}/{tile_name}.zip"


def lonlat_to_webmercator(lon: float, lat: float) -> Tuple[float, float]:
    """Convert lon/lat (WGS84) to Web Mercator (EPSG:3857)."""
    transformer = Transformer.from_crs("EPSG:4326", "EPSG:3857", always_xy=True)
    x, y = transformer.transform(lon, lat)
    return x, y


def query_tile_info(x: float, y: float) -> Optional[Dict[str, Any]]:
    """
    Query Ohio ArcGIS service for tile information at given Web Mercator coordinates.
    
    Returns dict with keys: tile_name, county, year, block, note
    Returns None if no tile found or on error.
    """
    params = {
        'f': 'json',
        'returnGeometry': 'false',
        'spatialRel': 'esriSpatialRelIntersects',
        'geometry': json.dumps({
            'x': x,
            'y': y,
            'spatialReference': {'wkid': 3857}
        }),
        'geometryType': 'esriGeometryPoint',
        'inSR': '3857',
        'outFields': '*',
        'outSR': '3857'
    }
    
    try:
        response = requests.get(TILE_SERVICE_URL, params=params, timeout=10)
        response.raise_for_status()
        data = response.json()
        
        if 'features' not in data or len(data['features']) == 0:
            return None
        
        # Get first feature (should only be one)
        attrs = data['features'][0]['attributes']
        
        return {
            'tile_name': attrs.get('TileName'),
            'county': attrs.get('County'),
            'year': attrs.get('Year'),
            'block': attrs.get('Block'),
            'note': attrs.get('note')
        }
    
    except requests.RequestException as e:
        logger.error(f"Failed to query tile service: {e}")
        return None
    except (KeyError, json.JSONDecodeError) as e:
        logger.error(f"Failed to parse tile service response: {e}")
        return None


def download_tile(tile_name: str, county: str, output_dir: Path) -> Tuple[bool, Optional[Path]]:
    """
    Download a tile zip file from OGRIP.
    
    Returns (success, output_path)
    """
    url = DOWNLOAD_URL_TEMPLATE.format(county=county, tile_name=tile_name)
    output_path = output_dir / f"{tile_name}.zip"
    
    # Check if already exists
    if output_path.exists():
        logger.info(f"Tile {tile_name} already exists, skipping download")
        return True, output_path
    
    try:
        logger.info(f"Downloading {tile_name} from {county}...")
        response = requests.get(url, timeout=30, stream=True)
        response.raise_for_status()
        
        # Download with progress
        output_dir.mkdir(parents=True, exist_ok=True)
        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        size_mb = output_path.stat().st_size / (1024 * 1024)
        logger.info(f"Downloaded {tile_name} ({size_mb:.2f} MB)")
        return True, output_path
    
    except requests.RequestException as e:
        logger.error(f"Failed to download {tile_name}: {e}")
        # Clean up partial download
        if output_path.exists():
            output_path.unlink()
        return False, None


def process_coordinates(
    coords: List[Tuple[float, float]],
    output_dir: Path,
    min_delay: float
) -> Dict[str, Any]:
    """
    Process a list of coordinates and download corresponding tiles.
    
    Returns summary dict with download results.
    """
    results = {
        'requested': len(coords),
        'downloaded': 0,
        'skipped': 0,
        'failed': 0,
        'tiles': []
    }
    
    # Track unique tiles to avoid duplicate downloads
    seen_tiles = set()
    
    for i, (lon, lat) in enumerate(coords):
        logger.info(f"Processing coordinate {i+1}/{len(coords)}: ({lon:.6f}, {lat:.6f})")
        
        # Convert to Web Mercator
        x, y = lonlat_to_webmercator(lon, lat)
        logger.debug(f"Web Mercator: ({x:.2f}, {y:.2f})")
        
        # Query tile info
        tile_info = query_tile_info(x, y)
        
        if not tile_info:
            logger.warning(f"No tile found for coordinates ({lon:.6f}, {lat:.6f})")
            results['failed'] += 1
            results['tiles'].append({
                'lon': lon,
                'lat': lat,
                'status': 'not_found'
            })
            continue
        
        tile_name = tile_info['tile_name']
        county = tile_info['county']
        
        # Check if we've already processed this tile
        if tile_name in seen_tiles:
            logger.info(f"Tile {tile_name} already processed in this run, skipping")
            results['skipped'] += 1
            continue
        
        seen_tiles.add(tile_name)
        
        # Download tile
        output_path = output_dir / f"{tile_name}.zip"
        
        if output_path.exists():
            logger.info(f"Tile {tile_name} already exists on disk")
            results['skipped'] += 1
            status = 'already_exists'
            success = True
        else:
            success, download_path = download_tile(tile_name, county, output_dir)
            
            if success:
                results['downloaded'] += 1
                status = 'downloaded'
            else:
                results['failed'] += 1
                status = 'download_failed'
        
        # Add to results
        tile_result = {
            'tile_name': tile_name,
            'county': county,
            'year': tile_info.get('year'),
            'block': tile_info.get('block'),
            'url': DOWNLOAD_URL_TEMPLATE.format(county=county, tile_name=tile_name),
            'output_path': str(output_path) if success else None,
            'status': status
        }
        results['tiles'].append(tile_result)
        
        # Rate limiting between downloads (skip if this was already on disk)
        if status == 'downloaded' and i < len(coords) - 1:
            logger.debug(f"Waiting {min_delay}s before next request...")
            time.sleep(min_delay)
    
    return results


def parse_coords_from_stdin() -> List[Tuple[float, float]]:
    """Parse coordinates from stdin (one per line, comma or space separated)."""
    coords = []
    for line in sys.stdin:
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        
        # Try comma separator first, then space
        parts = line.replace(',', ' ').split()
        if len(parts) >= 2:
            try:
                lon = float(parts[0])
                lat = float(parts[1])
                coords.append((lon, lat))
            except ValueError:
                logger.warning(f"Skipping invalid line: {line}")
    
    return coords


def parse_coords_from_file(filepath: Path) -> List[Tuple[float, float]]:
    """Parse coordinates from a file."""
    coords = []
    with open(filepath, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            
            parts = line.replace(',', ' ').split()
            if len(parts) >= 2:
                try:
                    lon = float(parts[0])
                    lat = float(parts[1])
                    coords.append((lon, lat))
                except ValueError:
                    logger.warning(f"Skipping invalid line: {line}")
    
    return coords


def main():
    parser = argparse.ArgumentParser(
        description='Download Ohio lidar tiles based on coordinates',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Download tiles for specific coordinates
  %(prog)s --coords -82.4438 40.0547 --coords -82.4440 40.0548 --output tiles/
  
  # Read coordinates from stdin
  echo "-82.4438,40.0547" | %(prog)s --output tiles/
  
  # Read from file
  %(prog)s --coords-file locations.txt --output tiles/
  
  # Custom rate limiting
  %(prog)s --coords -82.4438 40.0547 --output tiles/ --delay 2.0

Coordinate format:
  Longitude, Latitude in decimal degrees (WGS84)
  Example: -82.4438 40.0547 (Newark Earthworks area)
        """
    )
    
    parser.add_argument(
        '--coords',
        nargs=2,
        type=float,
        action='append',
        metavar=('LON', 'LAT'),
        help='Coordinate pair (longitude latitude). Can be specified multiple times.'
    )
    
    parser.add_argument(
        '--coords-file',
        type=Path,
        metavar='FILE',
        help='File containing coordinates (one pair per line, comma or space separated)'
    )
    
    parser.add_argument(
        '--output',
        type=Path,
        required=True,
        metavar='DIR',
        help='Output directory for downloaded tile zip files'
    )
    
    parser.add_argument(
        '--delay',
        type=float,
        default=1.0,
        metavar='SECONDS',
        help='Minimum delay between tile downloads (default: 1.0 seconds)'
    )
    
    parser.add_argument(
        '--verbose',
        action='store_true',
        help='Enable verbose debug logging'
    )
    
    args = parser.parse_args()
    
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    
    # Collect coordinates from all sources
    coords = []
    
    if args.coords:
        coords.extend(args.coords)
    
    if args.coords_file:
        if not args.coords_file.exists():
            logger.error(f"Coordinates file not found: {args.coords_file}")
            sys.exit(1)
        coords.extend(parse_coords_from_file(args.coords_file))
    
    # Check stdin if no coordinates provided
    if not coords and not sys.stdin.isatty():
        coords.extend(parse_coords_from_stdin())
    
    if not coords:
        logger.error("No coordinates provided. Use --coords, --coords-file, or pipe to stdin.")
        parser.print_help(sys.stderr)
        sys.exit(1)
    
    logger.info(f"Processing {len(coords)} coordinate(s)")
    logger.info(f"Output directory: {args.output}")
    logger.info(f"Rate limit: {args.delay}s between downloads")
    
    # Process coordinates and download tiles
    results = process_coordinates(coords, args.output, args.delay)
    
    # Output JSON to stdout
    print(json.dumps(results, indent=2))
    
    # Log summary to stderr
    logger.info("="*50)
    logger.info(f"Summary: {results['downloaded']} downloaded, {results['skipped']} skipped, {results['failed']} failed")
    
    if results['failed'] > 0:
        logger.error(f"{results['failed']} tile(s) failed to download or locate")
        sys.exit(1)


if __name__ == '__main__':
    main()