import gradio as gr
import json
import os
from datetime import datetime
from wave_data_puller import WaveDataPuller
from grib_wave_puller import GRIBWavePuller
import logging
import folium
import numpy as np
from folium import plugins

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def create_global_wave_map(global_data):
    """Create a Folium map with global GRIB wave data visualization"""
    try:
        if not global_data or 'sample_points' not in global_data:
            return None
        
        # Create base map centered on world
        m = folium.Map(
            location=[0, 0],
            zoom_start=2,
            tiles='OpenStreetMap'
        )
        
        sample_points = global_data['sample_points']
        if not sample_points:
            return None
        
        # Color based on wave height (blue to red scale)
        def get_color(height):
            if height < 1.0:
                return 'blue'
            elif height < 2.0:
                return 'green' 
            elif height < 3.0:
                return 'orange'
            elif height < 4.0:
                return 'red'
            else:
                return 'darkred'
        
        # Add sample points to map
        for point in sample_points:
            lat = point['lat']
            lon = point['lon']
            wave_height = point['wave_height']
            
            popup_text = f"""
            <b>Global Wave Data</b><br>
            Lat: {lat}°, Lon: {lon}°<br>
            Wave Height: {wave_height} m<br>
            Source: {global_data.get('data_source', 'N/A')}
            """
            
            folium.CircleMarker(
                location=[lat, lon],
                radius=max(3, wave_height * 2),  # Scale radius with wave height
                popup=popup_text,
                color=get_color(wave_height),
                fill=True,
                fillColor=get_color(wave_height),
                fillOpacity=0.7,
                weight=2
            ).add_to(m)
        
        # Add statistics popup with forecast info
        stats = global_data.get('wave_statistics', {})
        forecast_info = global_data.get('forecast_info', {})
        grid_info = global_data.get('grid_info', {})
        params_found = global_data.get('parameters_found', {})
        
        # Handle regional coverage information
        regions_info = ""
        if 'regions_included' in grid_info:
            regions = grid_info.get('regions_included', [])
            regions_info = f"<b>Regional Coverage:</b> {', '.join(regions)}<br>"
        
        stats_text = f"""
        <b>Global Wave Statistics</b><br>
        Max Wave Height: {stats.get('max_wave_height', 'N/A')} m<br>
        Mean Wave Height: {stats.get('mean_wave_height', 'N/A'):.2f} m<br>
        Data Points: {len(sample_points)}<br>
        {regions_info}
        <b>Data Coverage:</b><br>
        Lat: {grid_info.get('lat_min', 'N/A'):.1f}° to {grid_info.get('lat_max', 'N/A'):.1f}°<br>
        Lon: {grid_info.get('lon_min', 'N/A'):.1f}° to {grid_info.get('lon_max', 'N/A'):.1f}°<br>
        <b>Parameters:</b><br>
        Height: {params_found.get('wave_height', 'N/A')}<br>
        Direction: {params_found.get('wave_direction', 'N/A')}<br>
        Period: {params_found.get('wave_period', 'N/A')}<br>
        Vectors: {'Yes' if params_found.get('has_velocity_components') else 'No'}<br>
        <b>Forecast Info:</b><br>
        Forecast Hour: +{forecast_info.get('forecast_hour', 0)}h<br>
        Model Run: {forecast_info.get('model_run', 'N/A')}<br>
        Valid Time: {forecast_info.get('forecast_valid_time', 'N/A')[:16]}<br>
        Current: {'Yes' if forecast_info.get('is_current') else 'No'}<br>
        Source: {global_data.get('data_source', 'N/A')}
        """
        
        folium.Marker(
            location=[60, -120],  # Top left corner
            popup=stats_text,
            icon=folium.Icon(color='black', icon='info-sign')
        ).add_to(m)
        
        # Add legend
        legend_html = '''
        <div style="position: fixed; 
                    bottom: 50px; left: 50px; width: 180px; height: 140px; 
                    background-color: white; border:2px solid grey; z-index:9999; 
                    font-size:12px; padding: 10px">
        <b>Global Wave Height Legend</b><br>
        <i class="fa fa-circle" style="color:blue"></i> < 1m<br>
        <i class="fa fa-circle" style="color:green"></i> 1-2m<br>
        <i class="fa fa-circle" style="color:orange"></i> 2-3m<br>
        <i class="fa fa-circle" style="color:red"></i> 3-4m<br>
        <i class="fa fa-circle" style="color:darkred"></i> > 4m<br>
        Circle size ∝ wave height
        </div>
        '''
        m.get_root().html.add_child(folium.Element(legend_html))
        
        return m
        
    except Exception as e:
        logger.error(f"Error creating global map: {e}")
        return None

def create_wave_map(data):
    """Create a Folium map with single location wave data visualization"""
    try:
        if not data or 'location' not in data or 'wave_data' not in data:
            return None
        
        lat = data['location']['lat']
        lon = data['location']['lon']
        wave_data = data['wave_data']
        
        # Create base map centered on location
        m = folium.Map(
            location=[lat, lon],
            zoom_start=8,
            tiles='OpenStreetMap'
        )
        
        # Add wave height as circle marker with color intensity
        wave_height = wave_data.get('significant_wave_height', 0)
        wave_period = wave_data.get('peak_wave_period', 0)
        wave_direction = wave_data.get('wave_direction', 0)
        wind_speed = wave_data.get('wind_speed', 0)
        wind_direction = wave_data.get('wind_direction', 0)
        
        # Color based on wave height (blue to red scale)
        def get_color(height):
            if height < 1.0:
                return 'blue'
            elif height < 2.0:
                return 'green'
            elif height < 3.0:
                return 'orange'
            else:
                return 'red'
        
        # Main location marker with wave data
        popup_text = f"""
        <b>Wave Data Location</b><br>
        Lat: {lat}°, Lon: {lon}°<br><br>
        <b>Wave Conditions:</b><br>
        • Height: {wave_height} m<br>
        • Period: {wave_period} s<br>
        • Direction: {wave_direction}°<br><br>
        <b>Wind Conditions:</b><br>
        • Speed: {wind_speed} m/s<br>
        • Direction: {wind_direction}°<br><br>
        <b>Model:</b> {data.get('model', 'N/A')}<br>
        <b>Time:</b> {data.get('timestamp', 'N/A')}
        """
        
        folium.CircleMarker(
            location=[lat, lon],
            radius=max(10, wave_height * 5),  # Scale radius with wave height
            popup=popup_text,
            color=get_color(wave_height),
            fill=True,
            fillColor=get_color(wave_height),
            fillOpacity=0.7,
            weight=3
        ).add_to(m)
        
        # Add wave direction arrow
        # Calculate arrow end point (approximate)
        arrow_length = 0.1  # degrees
        wave_rad = np.radians(wave_direction)
        arrow_end_lat = lat + arrow_length * np.cos(wave_rad)
        arrow_end_lon = lon + arrow_length * np.sin(wave_rad)
        
        folium.PolyLine(
            locations=[[lat, lon], [arrow_end_lat, arrow_end_lon]],
            color='red',
            weight=4,
            opacity=0.8,
            popup=f"Wave Direction: {wave_direction}°"
        ).add_to(m)
        
        # Add wind direction arrow (different color)
        wind_rad = np.radians(wind_direction)
        wind_end_lat = lat + (arrow_length * 0.7) * np.cos(wind_rad)
        wind_end_lon = lon + (arrow_length * 0.7) * np.sin(wind_rad)
        
        folium.PolyLine(
            locations=[[lat, lon], [wind_end_lat, wind_end_lon]],
            color='green',
            weight=3,
            opacity=0.8,
            popup=f"Wind Direction: {wind_direction}°, Speed: {wind_speed} m/s"
        ).add_to(m)
        
        # Add legend
        legend_html = '''
        <div style="position: fixed; 
                    bottom: 50px; left: 50px; width: 150px; height: 120px; 
                    background-color: white; border:2px solid grey; z-index:9999; 
                    font-size:14px; padding: 10px">
        <b>Legend</b><br>
        <i class="fa fa-circle" style="color:blue"></i> Wave Height < 1m<br>
        <i class="fa fa-circle" style="color:green"></i> Wave Height 1-2m<br>
        <i class="fa fa-circle" style="color:orange"></i> Wave Height 2-3m<br>
        <i class="fa fa-circle" style="color:red"></i> Wave Height > 3m<br>
        <span style="color:red">▬</span> Wave Direction<br>
        <span style="color:green">▬</span> Wind Direction
        </div>
        '''
        m.get_root().html.add_child(folium.Element(legend_html))
        
        return m
        
    except Exception as e:
        logger.error(f"Error creating map: {e}")
        return None

def fetch_global_wave_data():
    """Fetch global GRIB wave data and return formatted results with map"""
    try:
        grib_puller = GRIBWavePuller()
        global_data = grib_puller.fetch_global_wave_data()
        
        if global_data:
            # Format the data for display
            formatted_data = json.dumps(global_data, indent=2)
            
            # Save to file
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"global_wave_data_{timestamp}.json"
            filepath = os.path.join(grib_puller.output_dir, filename)
            
            try:
                with open(filepath, 'w') as f:
                    json.dump(global_data, f, indent=2)
                logger.info(f"Global data saved to {filepath}")
            except Exception as save_error:
                logger.error(f"Error saving global data: {save_error}")
            
            # Create global map
            global_map = create_global_wave_map(global_data)
            map_html = global_map._repr_html_() if global_map else "<p>Unable to generate global map</p>"
            
            return formatted_data, "✅ Global wave data fetched successfully!", map_html
        else:
            error_msg = """
❌ No real government wave data available at this time.

Possible reasons:
• NOAA GRIB servers may be temporarily unavailable
• ECMWF free data access limited
• Data files not yet published for requested forecast time
• Server maintenance in progress

Try again later or check NOAA/ECMWF status pages for updates.
Only real government/institutional data is supported - no mock data."""
            
            no_data_html = """
            <div style="padding: 20px; text-align: center; background-color: #f8f9fa; border: 1px solid #dee2e6; border-radius: 5px;">
                <h3 style="color: #6c757d;">⚠️ No Real Wave Data Available</h3>
                <p style="color: #6c757d;">Government wave data servers (NOAA/ECMWF) are currently unavailable.</p>
                <p style="color: #6c757d;">Please try again later when data servers are accessible.</p>
                <p style="color: #6c757d;"><strong>Only real government data is supported.</strong></p>
            </div>
            """
            
            return error_msg, "❌ No real government data available", no_data_html
            
    except Exception as e:
        logger.error(f"Error in fetch_global_wave_data: {e}")
        return f"Error: {str(e)}", "❌ Error occurred", "<p>Error generating global map</p>"

def fetch_wave_data():
    """Fetch single location wave data and return formatted results with map"""
    try:
        puller = WaveDataPuller()
        data = puller.fetch_wave_data()
        
        if data:
            # Format the data for display
            formatted_data = json.dumps(data, indent=2)
            
            # Save to file
            puller.save_data(data)
            
            # Create map
            wave_map = create_wave_map(data)
            map_html = wave_map._repr_html_() if wave_map else "<p>Unable to generate map</p>"
            
            return formatted_data, "✅ Data fetched successfully!", map_html
        else:
            return "No data available", "❌ Failed to fetch data", "<p>No map data available</p>"
            
    except Exception as e:
        logger.error(f"Error in fetch_wave_data: {e}")
        return f"Error: {str(e)}", "❌ Error occurred", "<p>Error generating map</p>"

def get_recent_files():
    """Get list of recent wave data files"""
    try:
        data_dir = "/tmp/wave_data"
        if not os.path.exists(data_dir):
            return "No data files found"
        
        files = [f for f in os.listdir(data_dir) if f.endswith('.json')]
        files.sort(reverse=True)  # Most recent first
        
        if not files:
            return "No data files found"
        
        file_list = []
        for file in files[:10]:  # Show last 10 files
            filepath = os.path.join(data_dir, file)
            try:
                with open(filepath, 'r') as f:
                    data = json.load(f)
                timestamp = data.get('timestamp', 'Unknown')
                file_list.append(f"📄 {file} (Generated: {timestamp})")
            except:
                file_list.append(f"📄 {file}")
        
        return "\n".join(file_list)
        
    except Exception as e:
        return f"Error reading files: {str(e)}"

def view_file_content(filename):
    """View content of a specific data file with map"""
    try:
        if not filename:
            return "Please enter a filename", "<p>No file selected</p>"
        
        # Clean filename (remove emoji and extra text)
        clean_filename = filename.split()[1] if " " in filename else filename
        
        filepath = os.path.join("/tmp/wave_data", clean_filename)
        if not os.path.exists(filepath):
            return f"File {clean_filename} not found", "<p>File not found</p>"
        
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        # Create map for this data
        wave_map = create_wave_map(data)
        map_html = wave_map._repr_html_() if wave_map else "<p>Unable to generate map</p>"
        
        return json.dumps(data, indent=2), map_html
        
    except Exception as e:
        return f"Error reading file: {str(e)}", "<p>Error generating map</p>"

# Create Gradio interface
with gr.Blocks(title="NWPS SWAN Wave Data Puller") as demo:
    gr.Markdown("""
    # 🌊 Global Wave Data Puller
    
    This application fetches wave data from multiple sources:
    - **Global GRIB Data**: Worldwide wave heights from ECMWF/NOAA GRIB files
    - **Single Location**: NWPS SWAN model data for specific locations
    """)
    
    with gr.Tab("Global GRIB Data"):
        gr.Markdown("### Fetch Global Wave GRIB Data")
        gr.Markdown("Downloads global wave height data from ECMWF/NOAA GRIB files covering the entire world.")
        
        global_fetch_button = gr.Button("🌍 Fetch Global Wave Data", variant="primary")
        global_status_output = gr.Textbox(label="Status", interactive=False)
        
        with gr.Row():
            with gr.Column():
                global_data_output = gr.Textbox(label="Global Wave Data (JSON)", lines=20, interactive=False)
            with gr.Column():
                global_map_output = gr.HTML(label="Global Wave Map")
        
        global_fetch_button.click(
            fn=fetch_global_wave_data,
            outputs=[global_data_output, global_status_output, global_map_output]
        )
        
    with gr.Tab("Single Location Data"):
        gr.Markdown("### Fetch Single Location Wave Data")
        gr.Markdown("Fetches wave data for a specific location (original functionality).")
        
        fetch_button = gr.Button("🌊 Fetch Location Wave Data", variant="primary")
        status_output = gr.Textbox(label="Status", interactive=False)
        
        with gr.Row():
            with gr.Column():
                data_output = gr.Textbox(label="Wave Data (JSON)", lines=20, interactive=False)
            with gr.Column():
                map_output = gr.HTML(label="Wave Data Map")
        
        fetch_button.click(
            fn=fetch_wave_data,
            outputs=[data_output, status_output, map_output]
        )
    
    with gr.Tab("Data Files"):
        gr.Markdown("### Recent Data Files")
        
        refresh_button = gr.Button("🔄 Refresh File List")
        files_output = gr.Textbox(label="Recent Files", lines=10, interactive=False)
        
        gr.Markdown("### View File Content")
        filename_input = gr.Textbox(label="Filename (copy from list above)", placeholder="wave_data_20250826_203857.json")
        view_button = gr.Button("👁️ View File")
        
        with gr.Row():
            with gr.Column():
                content_output = gr.Textbox(label="File Content", lines=15, interactive=False)
            with gr.Column():
                file_map_output = gr.HTML(label="File Data Map")
        
        refresh_button.click(
            fn=get_recent_files,
            outputs=files_output
        )
        
        view_button.click(
            fn=view_file_content,
            inputs=filename_input,
            outputs=[content_output, file_map_output]
        )
        
        # Removed auto-load on startup to prevent hanging during initialization
        # demo.load(fn=get_recent_files, outputs=files_output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)