backend/core/source_resolver.py

"""
Source Resolver for Multi-Source Data Ingestion

Handles resolution of various source URIs:
- Local paths: /path/to/file.geojson
- S3 URIs: s3://bucket/key.shp
- HTTP URLs: https://example.com/data.csv
"""

import logging
import os
import tempfile
import shutil
from pathlib import Path
from typing import Optional, Tuple
from urllib.parse import urlparse
import requests

logger = logging.getLogger(__name__)


class SourceResolver:
    """Resolves various source URIs to local file paths."""
    
    def __init__(self, cache_dir: Optional[Path] = None):
        """
        Initialize the source resolver.
        
        Args:
            cache_dir: Directory for caching downloaded files.
                       Defaults to a temp directory.
        """
        self.cache_dir = cache_dir or Path(tempfile.gettempdir()) / "renewhere_cache"
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        
    def resolve(self, uri: str) -> Tuple[Path, bool]:
        """
        Resolve a URI to a local file path.
        
        Args:
            uri: Local path, S3 URI, or HTTP URL
            
        Returns:
            Tuple of (local_path, is_temporary)
            - local_path: Path to the local file
            - is_temporary: True if file was downloaded and should be cleaned up
        """
        parsed = urlparse(uri)
        
        if parsed.scheme == 's3':
            return self._resolve_s3(uri, parsed)
        elif parsed.scheme in ('http', 'https'):
            return self._resolve_http(uri, parsed)
        else:
            # Assume local path
            return self._resolve_local(uri)
    
    def _resolve_local(self, path: str) -> Tuple[Path, bool]:
        """Resolve a local file path."""
        local_path = Path(path).expanduser().resolve()
        
        if not local_path.exists():
            raise FileNotFoundError(f"Local file not found: {local_path}")
            
        return local_path, False
    
    def _resolve_s3(self, uri: str, parsed) -> Tuple[Path, bool]:
        """
        Resolve an S3 URI to a local file.
        
        Uses boto3 with credentials from environment variables:
        - AWS_ACCESS_KEY_ID
        - AWS_SECRET_ACCESS_KEY
        - AWS_DEFAULT_REGION (optional)
        """
        try:
            import boto3
            from botocore.exceptions import ClientError
        except ImportError:
            raise ImportError(
                "boto3 is required for S3 support. Install with: pip install boto3"
            )
        
        bucket = parsed.netloc
        key = parsed.path.lstrip('/')
        
        if not bucket or not key:
            raise ValueError(f"Invalid S3 URI: {uri}")
        
        # Generate cache filename
        safe_key = key.replace('/', '_')
        cache_path = self.cache_dir / f"s3_{bucket}_{safe_key}"
        
        # Check if already cached
        if cache_path.exists():
            logger.info(f"Using cached S3 file: {cache_path}")
            return cache_path, True
        
        logger.info(f"Downloading from S3: {uri}")
        
        # Download from S3
        s3 = boto3.client('s3')
        try:
            s3.download_file(bucket, key, str(cache_path))
        except ClientError as e:
            raise RuntimeError(f"Failed to download from S3: {e}")
        
        logger.info(f"Downloaded to: {cache_path}")
        return cache_path, True
    
    def _resolve_http(self, uri: str, parsed) -> Tuple[Path, bool]:
        """Resolve an HTTP/HTTPS URL to a local file."""
        # Generate cache filename from URL
        safe_name = parsed.path.split('/')[-1] or 'download'
        # Add hash of full URL to avoid collisions
        url_hash = str(hash(uri))[-8:]
        cache_path = self.cache_dir / f"http_{url_hash}_{safe_name}"
        
        # Check if already cached
        if cache_path.exists():
            logger.info(f"Using cached HTTP file: {cache_path}")
            return cache_path, True
        
        logger.info(f"Downloading from HTTP: {uri}")
        
        # Download with streaming
        try:
            response = requests.get(uri, stream=True, timeout=60)
            response.raise_for_status()
            
            with open(cache_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
                    
        except requests.RequestException as e:
            raise RuntimeError(f"Failed to download from HTTP: {e}")
        
        logger.info(f"Downloaded to: {cache_path}")
        return cache_path, True
    
    def cleanup(self, path: Path):
        """Remove a temporary file from cache."""
        if path.exists() and str(path).startswith(str(self.cache_dir)):
            try:
                path.unlink()
                logger.debug(f"Cleaned up: {path}")
            except Exception as e:
                logger.warning(f"Failed to cleanup {path}: {e}")
    
    def clear_cache(self):
        """Clear all cached files."""
        if self.cache_dir.exists():
            shutil.rmtree(self.cache_dir)
            self.cache_dir.mkdir(parents=True, exist_ok=True)
            logger.info("Cache cleared")


# Singleton instance
_resolver: Optional[SourceResolver] = None


def get_source_resolver() -> SourceResolver:
    """Get the singleton source resolver instance."""
    global _resolver
    if _resolver is None:
        _resolver = SourceResolver()
    return _resolver