import gradio as gr
import requests
import folium
import json
import time
import os
from typing import Dict, List, Optional, Tuple
import pandas as pd

class AirQualityMapper:
    """Class to handle AirNow API interactions and map generation"""
    
    def __init__(self):
        self.base_url = "https://www.airnowapi.org"
        self.aqi_colors = {
            "Good": "#00E400",
            "Moderate": "#FFFF00", 
            "Unhealthy for Sensitive Groups": "#FF7E00",
            "Unhealthy": "#FF0000",
            "Very Unhealthy": "#8F3F97",
            "Hazardous": "#7E0023"
        }
        self.aqi_ranges = {
            (0, 50): "Good",
            (51, 100): "Moderate", 
            (101, 150): "Unhealthy for Sensitive Groups",
            (151, 200): "Unhealthy",
            (201, 300): "Very Unhealthy",
            (301, 500): "Hazardous"
        }
    
    def get_aqi_category(self, aqi_value: int) -> str:
        """Get AQI category based on value"""
        for (min_val, max_val), category in self.aqi_ranges.items():
            if min_val <= aqi_value <= max_val:
                return category
        return "Unknown"
    
    def get_aqi_color(self, category: str) -> str:
        """Get color for AQI category"""
        return self.aqi_colors.get(category, "#808080")
    
    def fetch_airnow_data(self, api_key: str) -> Tuple[List[Dict], str]:
        """
        Fetch air quality data from AirNow API
        Returns: (data_list, status_message)
        """
        if not api_key or api_key.strip() == "":
            return [], "❌ Please enter a valid AirNow API key"
        
        try:
            # Get data for major US cities and regions
            # We'll use a comprehensive list of state capitals and major cities
            locations = [
                ("90210", "California"), ("10001", "New York"), ("60601", "Illinois"),
                ("75201", "Texas"), ("33101", "Florida"), ("30301", "Georgia"),
                ("98101", "Washington"), ("97201", "Oregon"), ("80201", "Colorado"),
                ("85001", "Arizona"), ("89101", "Nevada"), ("84101", "Utah"),
                ("59601", "Montana"), ("58501", "North Dakota"), ("57501", "South Dakota"),
                ("68501", "Nebraska"), ("66601", "Kansas"), ("73101", "Oklahoma"),
                ("55101", "Minnesota"), ("50301", "Iowa"), ("65101", "Missouri"),
                ("72201", "Arkansas"), ("70801", "Louisiana"), ("39201", "Mississippi"),
                ("35201", "Alabama"), ("37201", "Tennessee"), ("40601", "Kentucky"),
                ("25301", "West Virginia"), ("23219", "Virginia"), ("27601", "North Carolina"),
                ("29201", "South Carolina"), ("32301", "Florida"), ("01501", "Massachusetts"),
                ("06101", "Connecticut"), ("02901", "Rhode Island"), ("03301", "New Hampshire"),
                ("05601", "Vermont"), ("04330", "Maine"), ("19901", "Delaware"),
                ("21201", "Maryland"), ("17101", "Pennsylvania"), ("07001", "New Jersey"),
                ("12201", "New York"), ("43215", "Ohio"), ("46201", "Indiana"),
                ("48601", "Michigan"), ("53201", "Wisconsin"), ("99501", "Alaska"),
                ("96801", "Hawaii")
            ]
            
            all_data = []
            
            for zipcode, state in locations:
                try:
                    # Current observations endpoint
                    url = f"{self.base_url}/aq/observation/zipCode/current/"
                    params = {
                        "format": "application/json",
                        "zipCode": zipcode,
                        "distance": 50,  # 50 mile radius
                        "API_KEY": api_key
                    }
                    
                    response = requests.get(url, params=params, timeout=10)
                    
                    if response.status_code == 200:
                        data = response.json()
                        if data:  # If data is not empty
                            for observation in data:
                                observation['source_state'] = state
                                observation['source_zipcode'] = zipcode
                            all_data.extend(data)
                    
                    # Add delay to respect rate limits
                    time.sleep(0.5)
                    
                except requests.exceptions.RequestException as e:
                    continue  # Skip this location and continue with others
            
            if not all_data:
                return [], "⚠️ No air quality data found. Please check your API key or try again later."
            
            # Remove duplicates based on reporting area
            seen_areas = set()
            unique_data = []
            for item in all_data:
                area_key = (item.get('ReportingArea', ''), item.get('StateCode', ''))
                if area_key not in seen_areas:
                    seen_areas.add(area_key)
                    unique_data.append(item)
            
            return unique_data, f"✅ Successfully loaded {len(unique_data)} monitoring locations"
            
        except Exception as e:
            return [], f"❌ Error fetching data: {str(e)}"
    
    def create_map(self, data: List[Dict]) -> str:
        """Create an interactive map with air quality data"""
        if not data:
            # Create a basic US map if no data
            m = folium.Map(location=[39.8283, -98.5795], zoom_start=4)
            folium.Marker(
                [39.8283, -98.5795],
                popup="No data available. Please check your API key.",
                icon=folium.Icon(color='red', icon='info-sign')
            ).add_to(m)
            return m._repr_html_()
        
        # Calculate center point of all data
        lats = [item['Latitude'] for item in data if 'Latitude' in item]
        lons = [item['Longitude'] for item in data if 'Longitude' in item]
        
        if lats and lons:
            center_lat = sum(lats) / len(lats)
            center_lon = sum(lons) / len(lons)
        else:
            center_lat, center_lon = 39.8283, -98.5795  # Center of US
        
        # Create map
        m = folium.Map(location=[center_lat, center_lon], zoom_start=4)
        
        # Add markers for each monitoring location
        for item in data:
            try:
                lat = item.get('Latitude')
                lon = item.get('Longitude')
                aqi = item.get('AQI', 0)
                parameter = item.get('ParameterName', 'Unknown')
                area = item.get('ReportingArea', 'Unknown Area')
                state = item.get('StateCode', 'Unknown')
                category = item.get('Category', {}).get('Name', self.get_aqi_category(aqi))
                
                if lat is None or lon is None:
                    continue
                
                # Get color based on AQI category
                color = self.get_aqi_color(category)
                
                # Create popup content
                popup_content = f"""
                <div style="width: 200px;">
                    <h4>{area}, {state}</h4>
                    <p><b>AQI:</b> {aqi} ({category})</p>
                    <p><b>Parameter:</b> {parameter}</p>
                    <p><b>Location:</b> {lat:.3f}, {lon:.3f}</p>
                    <p><b>Last Updated:</b> {item.get('DateObserved', 'Unknown')} {item.get('HourObserved', '')}:00</p>
                </div>
                """
                
                # Determine marker color based on AQI
                if aqi <= 50:
                    marker_color = 'green'
                elif aqi <= 100:
                    marker_color = 'yellow'
                elif aqi <= 150:
                    marker_color = 'orange'
                elif aqi <= 200:
                    marker_color = 'red'
                elif aqi <= 300:
                    marker_color = 'purple'
                else:
                    marker_color = 'darkred'
                
                # Add marker
                folium.Marker(
                    [lat, lon],
                    popup=folium.Popup(popup_content, max_width=250),
                    tooltip=f"{area}: AQI {aqi}",
                    icon=folium.Icon(color=marker_color, icon='cloud')
                ).add_to(m)
                
            except Exception as e:
                continue  # Skip problematic markers
        
        # Add legend
        legend_html = """
        <div style="position: fixed; 
                    bottom: 50px; left: 50px; width: 150px; height: 180px; 
                    background-color: white; border:2px solid grey; z-index:9999; 
                    font-size:14px; padding: 10px">
        <h4>AQI Legend</h4>
        <p><i class="fa fa-circle" style="color:green"></i> Good (0-50)</p>
        <p><i class="fa fa-circle" style="color:yellow"></i> Moderate (51-100)</p>
        <p><i class="fa fa-circle" style="color:orange"></i> Unhealthy for Sensitive (101-150)</p>
        <p><i class="fa fa-circle" style="color:red"></i> Unhealthy (151-200)</p>
        <p><i class="fa fa-circle" style="color:purple"></i> Very Unhealthy (201-300)</p>
        <p><i class="fa fa-circle" style="color:darkred"></i> Hazardous (301+)</p>
        </div>
        """
        m.get_root().html.add_child(folium.Element(legend_html))
        
        return m._repr_html_()
    
    def create_data_table(self, data: List[Dict]) -> pd.DataFrame:
        """Create a data table from the air quality data"""
        if not data:
            return pd.DataFrame()
        
        # Extract relevant columns
        table_data = []
        for item in data:
            table_data.append({
                'Reporting Area': item.get('ReportingArea', 'Unknown'),
                'State': item.get('StateCode', 'Unknown'),
                'AQI': item.get('AQI', 0),
                'Category': item.get('Category', {}).get('Name', self.get_aqi_category(item.get('AQI', 0))),
                'Parameter': item.get('ParameterName', 'Unknown'),
                'Date': item.get('DateObserved', 'Unknown'),
                'Hour': item.get('HourObserved', 'Unknown'),
                'Latitude': item.get('Latitude', 'Unknown'),
                'Longitude': item.get('Longitude', 'Unknown')
            })
        
        df = pd.DataFrame(table_data)
        return df.sort_values('AQI', ascending=False)

# Initialize the mapper
mapper = AirQualityMapper()

def update_map(api_key: str):
    """Update the map with fresh air quality data"""
    if not api_key.strip():
        return "Please enter your AirNow API key above.", pd.DataFrame()
    
    # Fetch data
    data, status = mapper.fetch_airnow_data(api_key)
    
    # Create map
    map_html = mapper.create_map(data)
    
    # Create data table
    df = mapper.create_data_table(data)
    
    return map_html, df

# Create Gradio interface
with gr.Blocks(title="AirNow Air Quality Sensor Map", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🌬️ AirNow Air Quality Sensor Map
        
        This interactive map displays real-time air quality data from EPA's AirNow network of over 2,000 monitoring stations across the United States.
        
        ## How to use:
        1. **Get an API Key**: Register for a free API key at [docs.airnowapi.org](https://docs.airnowapi.org/)
        2. **Enter your API key** in the field below
        3. **Click "Load Air Quality Data"** to fetch current readings
        4. **Explore the map**: Click on markers to see detailed information about each monitoring station
        
        ## About the Data:
        - Data is updated hourly from state, local, tribal, and federal air quality agencies
        - Colors indicate Air Quality Index (AQI) levels from Good (green) to Hazardous (dark red)
        - AQI values tell you how clean or polluted the air is and associated health effects
        
        **⚠️ Note**: This data is preliminary and should not be used for regulatory decisions. For official data, visit [EPA's AirData](https://www.epa.gov/outdoor-air-quality-data).
        """
    )
    
    with gr.Row():
        with gr.Column(scale=3):
            api_key_input = gr.Textbox(
                label="AirNow API Key",
                placeholder="Enter your AirNow API key here...",
                type="password",
                info="Get your free API key at docs.airnowapi.org"
            )
        with gr.Column(scale=1):
            load_button = gr.Button("Load Air Quality Data", variant="primary", size="lg")
    
    status_text = gr.Markdown("Enter your API key and click 'Load Air Quality Data' to begin.")
    
    with gr.Tabs():
        with gr.TabItem("Interactive Map"):
            map_output = gr.HTML(label="Air Quality Map", height=600)
        
        with gr.TabItem("Data Table"):
            data_table = gr.Dataframe(
                label="Air Quality Monitoring Stations",
                height=500,
                interactive=False
            )
    
    gr.Markdown(
        """
        ## AQI Health Guidelines:
        
        - **Good (0-50)**: Air quality is satisfactory for everyone
        - **Moderate (51-100)**: Air quality is acceptable for most people
        - **Unhealthy for Sensitive Groups (101-150)**: Members of sensitive groups may experience health effects
        - **Unhealthy (151-200)**: Everyone may begin to experience health effects
        - **Very Unhealthy (201-300)**: Health warnings of emergency conditions
        - **Hazardous (301+)**: Health alert - everyone may experience serious health effects
        
        ## Data Sources:
        - **AirNow API**: Real-time air quality data from EPA's monitoring network
        - **Monitoring Agencies**: 120+ local, state, tribal, and federal government agencies
        - **Update Frequency**: Hourly observations, daily forecasts
        
        ## Links:
        - [AirNow.gov](https://www.airnow.gov) - Official air quality information
        - [AirNow API Documentation](https://docs.airnowapi.org/) - API documentation and registration
        - [EPA AirData](https://www.epa.gov/outdoor-air-quality-data) - Official regulatory air quality data
        """
    )
    
    # Set up event handler
    load_button.click(
        fn=update_map,
        inputs=[api_key_input],
        outputs=[map_output, data_table]
    ).then(
        fn=lambda: "Map updated with latest air quality data! 🌍",
        outputs=[status_text]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True)