import requests
import numpy as np
from PIL import Image
import io
from typing import Optional, List
import time

class WMSTileCapture:
    """Captures WMS tiles for processing - similar to how Folium fetches them."""
    
    def __init__(self):
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        })
    
    def capture_folium_style_tile(self, wms_url: str, layer: str, bbox: List[float], 
                                 width: int = 512, height: int = 512, zoom: int = 6) -> Optional[np.ndarray]:
        """Capture WMS tile using the exact same method as Folium's WmsTileLayer."""
        try:
            # Use identical parameters to Folium's WmsTileLayer
            params = {
                'SERVICE': 'WMS',
                'VERSION': '1.3.0',
                'REQUEST': 'GetMap',
                'LAYERS': layer,
                'STYLES': '',
                'FORMAT': 'image/png',
                'TRANSPARENT': 'TRUE',
                'WIDTH': str(width),
                'HEIGHT': str(height),
                'CRS': 'EPSG:4326'
            }
            
            # For WMS 1.3.0 with EPSG:4326, bbox order is miny,minx,maxy,maxx (south,west,north,east)
            west, south, east, north = bbox
            bbox_str = f"{south},{west},{north},{east}"
            params['BBOX'] = bbox_str
            
            print(f"Capturing WMS tile: {layer}")
            print(f"Bbox: {bbox} -> {bbox_str}")
            print(f"URL: {wms_url}")
            
            response = self.session.get(wms_url, params=params, timeout=30)
            
            print(f"Response status: {response.status_code}")
            print(f"Content-Type: {response.headers.get('content-type', 'unknown')}")
            print(f"Content-Length: {len(response.content)} bytes")
            
            if response.status_code != 200:
                print(f"HTTP Error: {response.text[:500]}")
                return None
            
            # Check if it's actually an image
            content_type = response.headers.get('content-type', '').lower()
            if 'image' not in content_type:
                print(f"Not an image: {response.text[:300]}")
                return None
            
            # Load image
            image = Image.open(io.BytesIO(response.content))
            image_array = np.array(image)
            
            print(f"Captured image: {image_array.shape}, dtype: {image_array.dtype}")
            
            # Check if image has any radar data (non-transparent pixels)
            if len(image_array.shape) >= 3 and image_array.shape[2] >= 4:
                non_transparent = np.sum(image_array[:, :, 3] > 0)
                print(f"Non-transparent pixels: {non_transparent}")
            
            return image_array
            
        except Exception as e:
            print(f"Error capturing WMS tile: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def get_north_america_radar(self, wms_url: str, layer: str) -> List[dict]:
        """Get radar coverage that matches Tab 1's exact parameters."""
        
        # Strategy 1: Try the exact same bounds that Folium uses for Tab 1
        # These bounds should match what's displayed in the working Tab 1
        
        print(f"Attempting to match Tab 1 radar coverage...")
        
        # Use the same geographic bounds as the original Canadian radar map
        # This should align perfectly with Tab 1
        canada_bbox = [-141.0, 41.6751, -52.6194, 83.1139]  # [west, south, east, north]
        
        print(f"Capturing with Canada bounds: {canada_bbox}")
        
        # Capture at maximum resolution to preserve detail
        tile = self.capture_folium_style_tile(wms_url, layer, canada_bbox, 2048, 2048)
        
        if tile is not None and len(tile.shape) >= 3:
            non_transparent = np.sum(tile[:, :, 3] > 0) if tile.shape[2] >= 4 else tile.shape[0] * tile.shape[1]
            print(f"High-res capture successful: {tile.shape}, {non_transparent} data pixels")
            return [{'image': tile, 'bbox': canada_bbox, 'description': 'Full Canada Coverage (High-Res)'}]
        
        # Strategy 2: Multiple high-resolution tiles covering the same area
        print("High-res single tile failed, trying tiled approach...")
        
        # Divide the Canada bounds into quadrants for better resolution
        west, south, east, north = canada_bbox
        mid_lon = (west + east) / 2
        mid_lat = (south + north) / 2
        
        quadrants = [
            {'bbox': [west, south, mid_lon, mid_lat], 'name': 'Southwest Canada'},
            {'bbox': [mid_lon, south, east, mid_lat], 'name': 'Southeast Canada'},
            {'bbox': [west, mid_lat, mid_lon, north], 'name': 'Northwest Canada'},
            {'bbox': [mid_lon, mid_lat, east, north], 'name': 'Northeast Canada'},
        ]
        
        tiles = []
        for quadrant in quadrants:
            bbox = quadrant['bbox']
            print(f"Capturing quadrant: {quadrant['name']} {bbox}")
            tile = self.capture_folium_style_tile(wms_url, layer, bbox, 1024, 1024)
            if tile is not None:
                tiles.append({'image': tile, 'bbox': bbox, 'description': quadrant['name']})
            time.sleep(0.2)
        
        return tiles
    
    def merge_radar_tiles(self, tile_data: List[dict]) -> Optional[dict]:
        """Merge multiple radar tiles and return the best one with metadata."""
        if not tile_data:
            return None
        
        # Find the tile with the most radar data
        best_tile_info = None
        max_data = 0
        
        for tile_info in tile_data:
            tile = tile_info['image']
            if len(tile.shape) >= 3 and tile.shape[2] >= 4:
                non_transparent = np.sum(tile[:, :, 3] > 0)
                print(f"Tile {tile_info['description']}: {non_transparent} radar pixels")
                if non_transparent > max_data:
                    max_data = non_transparent
                    best_tile_info = tile_info
        
        if best_tile_info is None and tile_data:
            # Fallback to first tile if no good data found
            best_tile_info = tile_data[0]
            
        return best_tile_info