import plotly.graph_objects as go
import plotly.express as px
import numpy as np
import pandas as pd
from datetime import datetime
import json
from typing import List, Dict, Tuple, Any
import re
from collections import defaultdict
from zoneinfo import ZoneInfo
from src.analysis.geospatial import haversine_distance
from src.config import config
from src.utils import parse_period_string_to_datetime

def format_datetime_for_display(datetime_str: str) -> str:
    """
    Format datetime string from 'YYYY-MM-DD HH:MM:SS' to 'DD-MM-YYYY HH:MM:SS'.
    
    Args:
        datetime_str: Datetime string in format 'YYYY-MM-DD HH:MM:SS'
    
    Returns:
        Formatted datetime string in 'DD-MM-YYYY HH:MM:SS' format
    """
    try:
        if datetime_str and datetime_str != 'N/A':
            dt = datetime.strptime(datetime_str, '%Y-%m-%d %H:%M:%S')
            return dt.strftime('%d-%m-%Y %H:%M:%S')
        return datetime_str
    except:
        return datetime_str

def parse_wkt_linestring(wkt_string: str) -> List[Tuple[float, float]]:
    """
    Parse WKT LINESTRING format to extract coordinates.
    
    Args:
        wkt_string: WKT string in format "LINESTRING(lon1 lat1, lon2 lat2, ...)"
    
    Returns:
        List of (longitude, latitude) tuples
    """
    try:
        # Extract coordinates from LINESTRING format
        # Remove "LINESTRING(" and ")" and split by commas
        coords_str = wkt_string.replace("LINESTRING(", "").replace(")", "")
        coord_pairs = coords_str.split(",")
        
        coordinates = []
        for pair in coord_pairs:
            lon, lat = pair.strip().split()
            coordinates.append((float(lon), float(lat)))
        
        return coordinates
    except Exception as e:
        print(f"Error parsing WKT: {e}")
        return []

def group_storm_data_by_day(storm_data: List[Dict]) -> Dict[str, List[Dict]]:
    """
    Group storm data by day.
    
    Args:
        storm_data: List of storm cell dictionaries
    
    Returns:
        Dictionary with date as key and list of storms as value
    """
    daily_storms = defaultdict(list)
    
    for storm in storm_data:
        # Try different time fields that might exist in storm data
        time_field = storm.get('effective_time') or storm.get('creation_time') or storm.get('expire_time', '')
        if time_field:
            try:
                # Parse time field (format: '2025-08-29T15:08:45.002Z' or '2024-06-22 11:13:00')
                if 'T' in time_field:
                    # ISO format with T
                    storm_date = datetime.strptime(time_field[:10], '%Y-%m-%d').strftime('%d-%m-%Y')
                else:
                    # Standard format
                    storm_date = datetime.strptime(time_field[:10], '%Y-%m-%d').strftime('%d-%m-%Y')
                daily_storms[storm_date].append(storm)
            except:
                continue
    
    return dict(daily_storms)

def group_storm_data_by_month(storm_data: List[Dict]) -> Dict[str, List[Dict]]:
    """
    Group storm data by month.
    
    Args:
        storm_data: List of storm cell dictionaries
    
    Returns:
        Dictionary with month as key and list of storms as value
    """
    monthly_storms = defaultdict(list)
    
    for storm in storm_data:
        # Try different time fields that might exist in storm data
        time_field = storm.get('effective_time') or storm.get('creation_time') or storm.get('expire_time', '')
        if time_field:
            try:
                # Parse time field (format: '2025-08-29T15:08:45.002Z' or '2024-06-22 11:13:00')
                if 'T' in time_field:
                    # ISO format with T
                    storm_month = datetime.strptime(time_field[:10], '%Y-%m-%d').strftime('%Y-%m')
                else:
                    # Standard format
                    storm_month = datetime.strptime(time_field[:10], '%Y-%m-%d').strftime('%Y-%m')
                monthly_storms[storm_month].append(storm)
            except:
                continue
    
    return dict(monthly_storms)

def create_storm_cells_coordinate_plane(storm_data: List[Dict], turbine_df: pd.DataFrame = None, center_lat: float = None, center_lon: float = None) -> go.Figure:
    """
    Create a coordinate plane visualization showing storm cells with turbines and distance rings.
    
    Args:
        storm_data: List of storm cell dictionaries with WKT data
        turbine_df: DataFrame containing turbine locations with 'lat' and 'lng' columns
        center_lat: Center latitude for map (optional)
        center_lon: Center longitude for map (optional)
    
    Returns:
        Plotly figure with storm cells, turbines, and distance rings in coordinate plane view
    """
    # Calculate center if not provided
    if center_lat is None or center_lon is None:
        if turbine_df is not None and not turbine_df.empty:
            # Use turbine centroid as center
            center_lat = turbine_df['lat'].mean()
            center_lon = turbine_df['lng'].mean()
        else:
            # Calculate center from storm data
            all_lats = []
            all_lons = []
            for storm in storm_data:
                coords = parse_wkt_linestring(storm.get('cell_polygon_wkt', ''))
                if coords:
                    lons, lats = zip(*coords)
                    all_lats.extend(lats)
                    all_lons.extend(lons)
            
            if all_lats and all_lons:
                center_lat = np.mean(all_lats)
                center_lon = np.mean(all_lons)
            else:
                center_lat, center_lon = 36.8, 33.8  # Default to Turkey
    
    fig = go.Figure()
    
    # Add distance rings if turbine data is provided
    if turbine_df is not None and not turbine_df.empty:
        from src.analysis.geospatial import create_circle_points
        
        # Add distance rings from turbine centroid
        for radius, color in zip(config.distance_rings, config.ring_colors):
            circle_lats, circle_lons = create_circle_points(center_lat, center_lon, radius)
            fig.add_trace(go.Scatter(
                x=circle_lons,  # X-axis = Longitude
                y=circle_lats,  # Y-axis = Latitude
                mode='lines',
                line=dict(color=color, width=2),
                opacity=0.6,
                name=f'{radius/1000:.0f}km Ring',
                showlegend=True
            ))
        
        # Add turbines colored by risk using fixed intervals
        if 'risk_log' in turbine_df.columns:
            from src.utils import get_turbine_colors_by_fixed_intervals
            turbine_colors = get_turbine_colors_by_fixed_intervals(turbine_df['risk_log'].tolist())
        else:
            turbine_colors = ['red'] * len(turbine_df)
            
        fig.add_trace(go.Scatter(
            x=turbine_df['lng'],  # X-axis = Longitude
            y=turbine_df['lat'],  # Y-axis = Latitude
            mode='markers+text',
            marker=dict(
                size=30, 
                color=turbine_colors, 
                symbol='triangle-down', 
                opacity=0.8,
                line=dict(color='black', width=1)
            ),
            text=turbine_df['name'].tolist(),
            textfont=dict(size=12, color='black'),
            textposition='middle center',
            name='Wind Turbines',
            showlegend=True,
            hovertemplate=(
                "<b>Wind Turbine</b><br>"
                "Name: %{text}<br>"
                "Lat: %{y:.5f}<br>"
                "Lng: %{x:.5f}<br>"
                "<extra></extra>"
            )
        ))
    
    severity_colors = {}
    seen_polygon_key_to_count: Dict[Tuple, int] = {}
    offset_deg_lon = 0.003
    offset_deg_lat = 0.002

    for i, storm in enumerate(storm_data):
        wkt_string = storm.get('cell_polygon_wkt', '')
        if not wkt_string:
            continue

        coords = parse_wkt_linestring(wkt_string)
        if not coords:
            continue

        polygon_key = tuple((round(lon, 6), round(lat, 6)) for lon, lat in coords)
        duplicate_index = 0
        if polygon_key in seen_polygon_key_to_count:
            seen_polygon_key_to_count[polygon_key] += 1
            duplicate_index = seen_polygon_key_to_count[polygon_key] - 1
        else:
            seen_polygon_key_to_count[polygon_key] = 1

        if duplicate_index > 0:
            off_lon = offset_deg_lon * duplicate_index
            off_lat = offset_deg_lat * duplicate_index
            coords = [(lon + off_lon, lat + off_lat) for lon, lat in coords]

        lons, lats = zip(*coords)

        severity = storm.get('lightning_severity', 'Unknown').lower()
        if severity == 'high':
            color = 'purple'
        elif severity == 'medium':
            color = 'orange'
        elif severity == 'low':
            color = 'green'
        else:
            color = 'gray'
        
        # Track severity colors for legend
        severity_colors[severity] = color
        
        # Add storm cell boundary
        fig.add_trace(go.Scatter(
            x=lons,  # X-axis = Longitude
            y=lats,  # Y-axis = Latitude
            mode='lines',
            line=dict(color=color, width=3),
            opacity=1,
            showlegend=False,  # Don't show individual storms in legend
            hovertemplate=(
                f"<b>Storm Cell {i+1}</b><br>"
                f"Severity: {storm.get('lightning_severity', 'Unknown')}<br>"
                f"Effective: {format_datetime_for_display(storm.get('effective_time', 'N/A'))}<br>"
                f"Expire: {format_datetime_for_display(storm.get('expire_time', 'N/A'))}<br>"
                f"Direction: {storm.get('direction', 'N/A')}°<br>"
                f"Speed: {storm.get('speed', 'N/A')} km/h<br>"
                f"<extra></extra>"
            )
        ))
    
    # Add static legend entries for severity levels
    for severity, color in severity_colors.items():
        fig.add_trace(go.Scatter(
            x=[None],  # Invisible trace for legend only
            y=[None],
            mode='lines',
            line=dict(color=color, width=3),
            name=f"{severity.title()} Severity",
            showlegend=True,
            hoverinfo='skip'  # No hover for legend entries
        ))
    
    # Calculate axis limits based on data and distance rings
    max_radius_deg = max(config.distance_rings) / 111000
    
    lat_min = center_lat - max_radius_deg * 1.5
    lat_max = center_lat + max_radius_deg * 1.5
    lon_min = center_lon - max_radius_deg * 1.5
    lon_max = center_lon + max_radius_deg * 1.5
    
    fig.update_layout(
        font=dict(size=18),
        title=dict(text='Storm Cells - Coordinate Plane View', font=dict(size=28)),
        xaxis_title='Longitude',
        yaxis_title='Latitude',
        xaxis=dict(
            range=[lon_min, lon_max],
            showgrid=True,
            gridwidth=1,
            gridcolor='lightgray',
            zeroline=False,
            tickfont=dict(size=22),
            title_font=dict(size=28),
        ),
        yaxis=dict(
            range=[lat_min, lat_max],
            showgrid=True,
            gridwidth=1,
            gridcolor='lightgray',
            zeroline=False,
            tickfont=dict(size=22),
            title_font=dict(size=28),
        ),
        plot_bgcolor='white',
        paper_bgcolor='white',
        showlegend=True,
        legend=dict(
            title=dict(text='Legend', font=dict(size=24)),
            font=dict(size=20),
            orientation='h',
            x=0.5,
            xanchor='center',
            y=-0.18,
            yanchor='top',
            bgcolor='rgba(255, 255, 255, 0.8)',
            bordercolor='black',
            borderwidth=1,
        ),
        width=800,
        height=900,
        margin=dict(l=70, r=40, t=50, b=130),
    )
    
    return fig

def create_storm_cells_map(storm_data: List[Dict], turbine_df: pd.DataFrame = None, center_lat: float = None, center_lon: float = None) -> go.Figure:
    """
    Create a map showing storm cells from the fırtına data with turbines and distance rings.
    Now uses coordinate plane view instead of mapbox.
    
    Args:
        storm_data: List of storm cell dictionaries with WKT data
        turbine_df: DataFrame containing turbine locations with 'lat' and 'lng' columns
        center_lat: Center latitude for map (optional)
        center_lon: Center longitude for map (optional)
    
    Returns:
        Plotly figure with storm cells, turbines, and distance rings in coordinate plane view
    """
    return create_storm_cells_coordinate_plane(storm_data, turbine_df, center_lat, center_lon)

def create_daily_storm_maps(storm_data: List[Dict], max_maps_per_page: int = 2) -> List[go.Figure]:
    """
    Create separate maps for each day with storms.
    
    Args:
        storm_data: List of storm cell dictionaries
        max_maps_per_page: Maximum number of maps to show per page
    
    Returns:
        List of Plotly figures, each containing maps for one or more days
    """
    daily_storms = group_storm_data_by_day(storm_data)
    
    if not daily_storms:
        return []
    
    # Sort days
    sorted_days = sorted(daily_storms.keys())
    
    figures = []
    current_fig = None
    maps_in_current_fig = 0
    
    for day in sorted_days:
        day_storms = daily_storms[day]
        
        if current_fig is None or maps_in_current_fig >= max_maps_per_page:
            # Create new figure
            current_fig = go.Figure()
            maps_in_current_fig = 0
            
            # Set up subplot layout
            if max_maps_per_page == 1:
                # Single map layout
                current_fig.update_layout(
                    mapbox=dict(
                        style='carto-positron',
                        center=dict(lat=36.8, lon=33.8),
                        zoom=8
                    ),
                    margin=dict(l=0, r=0, t=0, b=0),
                    showlegend=True
                )
            else:
                # Multiple maps layout - will be handled in PDF generation
                pass
        
        # Create map for this day
        day_fig = create_storm_cells_map(day_storms)
        
        # Add day title
        day_fig.update_layout(
            title=f"Storm Cells - {day}",
            title_x=0.5,
            title_font_size=16
        )
        
        # If this is the first map in the figure, use it as base
        if maps_in_current_fig == 0:
            current_fig = day_fig
        else:
            # For multiple maps, we'll handle layout in PDF generation
            pass
        
        maps_in_current_fig += 1
        
        # If we've reached the limit, add to figures list
        if maps_in_current_fig >= max_maps_per_page:
            figures.append(current_fig)
            current_fig = None
            maps_in_current_fig = 0
    
    # Add remaining figure if any
    if current_fig is not None:
        figures.append(current_fig)
    
    return figures

def create_monthly_storm_maps(storm_data: List[Dict]) -> Dict[str, go.Figure]:
    """
    Create separate maps for each month with storms.
    
    Args:
        storm_data: List of storm cell dictionaries
    
    Returns:
        Dictionary with month as key and Plotly figure as value
    """
    monthly_storms = group_storm_data_by_month(storm_data)
    
    monthly_figures = {}
    
    for month, month_storms in monthly_storms.items():
        if month_storms:
            fig = create_storm_cells_map(month_storms)
            fig.update_layout(
                title=f"Storm Cells - {month}",
                title_x=0.5,
                title_font_size=16
            )
            monthly_figures[month] = fig
    
    return monthly_figures

def load_storm_data_from_json(json_file_path: str) -> List[Dict]:
    """
    Load storm data from JSON file.
    
    Args:
        json_file_path: Path to the JSON file
    
    Returns:
        List of storm cell dictionaries
    """
    try:
        with open(json_file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Handle different JSON structures
        if isinstance(data, dict):
            # Handle structure with "success" and "data" keys
            if "data" in data and isinstance(data["data"], dict):
                # Convert dictionary of storm cells to list
                storm_list = []
                for storm_id, storm_records in data["data"].items():
                    if isinstance(storm_records, list):
                        storm_list.extend(storm_records)
                    else:
                        storm_list.append(storm_records)
                return storm_list
            # Handle structure where data is directly a list
            elif "data" in data and isinstance(data["data"], list):
                return data["data"]
            # Handle structure where data is directly the list
            else:
                return data
        
        # If data is already a list
        elif isinstance(data, list):
            return data
        
        return []
    except Exception as e:
        print(f"Error loading storm data: {e}")
        return []

def filter_storm_data_by_date_range(storm_data: List[Dict], start_date: str, end_date: str) -> List[Dict]:
    """
    Filter storm data by date range.
    start_date/end_date can be 'DD-MM-YYYY', 'DD-MM-YYYY HH:MM', or ISO (e.g. 2026-01-22T07:00:00Z).
    Storm timestamps are converted to the farm's local timezone (config.timezone) for comparison.
    """
    try:
        start_dt = parse_period_string_to_datetime(start_date)
        end_dt = parse_period_string_to_datetime(end_date)
        if start_dt is None or end_dt is None:
            return storm_data
        if re.fullmatch(r"\d{2}-\d{2}-\d{4}", str(end_date).strip()):
            end_dt = end_dt.replace(hour=23, minute=59, second=59)
        tz_name = getattr(config, 'timezone', None)
        tz = ZoneInfo(tz_name) if tz_name else None
        filtered_data = []
        for storm in storm_data:
            time_field = storm.get('effective_time') or storm.get('creation_time') or storm.get('expire_time', '')
            if time_field:
                try:
                    storm_ts = pd.to_datetime(time_field, utc=True)
                    if tz is not None:
                        storm_ts = storm_ts.tz_convert(tz).tz_localize(None)
                        storm_dt = storm_ts.to_pydatetime()
                    else:
                        storm_dt = storm_ts.to_pydatetime().replace(tzinfo=None)
                    if start_dt <= storm_dt <= end_dt:
                        filtered_data.append(storm)
                except Exception:
                    continue
        return filtered_data
    except Exception as e:
        print(f"Error filtering storm data: {e}")
        return storm_data

def filter_storm_data_by_turbine_proximity(storm_data: List[Dict], turbine_df: pd.DataFrame, max_distance_km: float = None) -> List[Dict]:
    """
    Filter storm data to only include storms within the specified distance from turbines.
    
    Args:
        storm_data: List of storm cell dictionaries
        turbine_df: DataFrame containing turbine locations with 'lat' and 'lng' columns
        max_distance_km: Maximum distance in kilometers. If None, uses the farthest distance ring from config.
    
    Returns:
        Filtered list of storm cell dictionaries
    """
    if max_distance_km is None:
        # Use the farthest distance ring from config (convert meters to km)
        max_distance_km = max(config.distance_rings) / 1000
    
    print(f"🌩️  Filtering storm cells within {max_distance_km} km of turbine locations...")
    
    filtered_storms = []
    
    for storm in storm_data:
        wkt_string = storm.get('cell_polygon_wkt', '')
        if not wkt_string:
            continue
        
        coords = parse_wkt_linestring(wkt_string)
        if not coords:
            continue
        
        # Check if any point in the storm cell is within the distance threshold
        storm_within_range = False
        
        for storm_lon, storm_lat in coords:
            for _, turbine in turbine_df.iterrows():
                turbine_lat = turbine['lat']
                turbine_lon = turbine['lng']
                
                distance_km = haversine_distance(turbine_lat, turbine_lon, storm_lat, storm_lon) / 1000
                
                if distance_km <= max_distance_km:
                    storm_within_range = True
                    break
            
            if storm_within_range:
                break
        
        if storm_within_range:
            filtered_storms.append(storm)
    
    print(f"🌩️  Filtered from {len(storm_data)} to {len(filtered_storms)} storm cells within {max_distance_km} km")
    return filtered_storms

def calculate_storm_cell_centroid(wkt_string: str) -> Tuple[float, float]:
    """
    Calculate the centroid of a storm cell from WKT coordinates.
    
    Args:
        wkt_string: WKT string representing the storm cell boundary
    
    Returns:
        Tuple of (latitude, longitude) for the centroid
    """
    coords = parse_wkt_linestring(wkt_string)
    if not coords:
        return None
    
    # Calculate centroid (simple average of all points)
    lons, lats = zip(*coords)
    centroid_lat = np.mean(lats)
    centroid_lon = np.mean(lons)
    
    return centroid_lat, centroid_lon

def create_storm_cells_summary(storm_data: List[Dict]) -> Dict[str, Any]:
    """
    Create a summary of storm cells data.
    
    Args:
        storm_data: List of storm cell dictionaries
    
    Returns:
        Dictionary with summary statistics
    """
    if not storm_data:
        return {}
    
    # Count by severity
    severity_counts = {}
    total_cells = len(storm_data)
    
    for storm in storm_data:
        severity = storm.get('lightning_severity', 'Unknown')
        severity_counts[severity] = severity_counts.get(severity, 0) + 1
    
    # Calculate average direction and speed
    directions = [storm.get('direction', 0) for storm in storm_data if storm.get('direction') is not None]
    speeds = [storm.get('speed', 0) for storm in storm_data if storm.get('speed') is not None]
    
    avg_direction = np.mean(directions) if directions else 0
    avg_speed = np.mean(speeds) if speeds else 0
    
    # Get date range
    time_fields = []
    for storm in storm_data:
        time_field = storm.get('effective_time') or storm.get('creation_time') or storm.get('expire_time', '')
        if time_field:
            time_fields.append(time_field)
    
    if time_fields:
        try:
            dates = []
            for time_field in time_fields:
                if 'T' in time_field:
                    # ISO format with T
                    dates.append(datetime.strptime(time_field[:10], '%Y-%m-%d'))
                else:
                    # Standard format
                    dates.append(datetime.strptime(time_field[:10], '%Y-%m-%d'))
            start_date = min(dates).strftime('%d-%m-%Y')
            end_date = max(dates).strftime('%d-%m-%Y')
        except:
            start_date = end_date = "Unknown"
    else:
        start_date = end_date = "Unknown"
    
    # Get daily breakdown
    daily_storms = group_storm_data_by_day(storm_data)
    daily_summary = {day: len(storms) for day, storms in daily_storms.items()}
    
    return {
        'total_cells': total_cells,
        'severity_counts': severity_counts,
        'avg_direction': avg_direction,
        'avg_speed': avg_speed,
        'date_range': {'start': start_date, 'end': end_date},
        'daily_breakdown': daily_summary
    }