# alert.py

import json
import os
import re
from typing import Any, Dict, List
from weatherr import get_weather
from dotenv import load_dotenv
from openai import OpenAI
import requests
load_dotenv()


class WeatherAlertAPI:
    def __init__(self):
        self.api_url = "http://sachet.ndma.gov.in/cap_public_website/FetchAllAlertDetails"
    
    def fetch_alerts(self) -> List[Dict[str, Any]]:
        """Fetch all weather alerts from the NDMA API"""
        try:
            response = requests.get(self.api_url, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            if isinstance(data, list):
                return data
            elif isinstance(data, dict) and "nowcastDetails" in data:
                return data["nowcastDetails"]
            else:
                return []
                
        except Exception as e:
            print(f"Error fetching alerts: {e}")
            return []
    
    def normalize_text(self, text: str) -> str:
        """Normalize text for better matching"""
        return re.sub(r'\s+', ' ', text.strip().lower())
    
    def location_matches(self, area_description: str, search_terms: List[str]) -> bool:
        """Check if any search terms match the area description"""
        normalized_area = self.normalize_text(area_description)
        for term in search_terms:
            normalized_term = self.normalize_text(term)
            if normalized_term in normalized_area:
                return True
        return False

# Initialize the weather API
weather_api = WeatherAlertAPI()

# OpenAI Tool Definitions (Updated format)
WEATHER_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_weather_alerts",
            "description": "Get weather alerts for one or more locations in India. Can search by city, state, district, or region name.",
            "parameters": {
                "type": "object",
                "properties": {
                    "locations": {
                        "type": "array",
                        "items": {
                            "type": "string"
                        },
                        "description": "List of locations to search for (up to 5). Can be city names, state names, districts, etc.",
                        "maxItems": 5,
                        "minItems": 1
                    },
                    "include_details": {
                        "type": "boolean",
                        "description": "Whether to include detailed alert information like warning messages and coordinates",
                        "default": True
                    }
                },
                "required": ["locations"]
            }
        }
    },{
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Fetch current weather using address, coordinates, or district/state.",
            "parameters": {
                "type": "object",
                "properties": {
                    "latitude": {"type": "number", "description": "Latitude coordinate"},
                    "longitude": {"type": "number", "description": "Longitude coordinate"},
                    "address": {"type": "string", "description": "Full address or location name"},
                    "district": {"type": "string", "description": "District name"},
                    "state": {"type": "string", "description": "State name"}
                },
                "required": []  # none are strictly required; function will validate
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_alert_summary",
            "description": "Get a summary of all current weather alerts by severity level",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_available_locations",
            "description": "Get a list of all locations that currently have weather alerts",
            "parameters": {
                "type": "object",
                "properties": {
                    "limit": {
                        "type": "integer",
                        "description": "Maximum number of locations to return",
                        "default": 50,
                        "minimum": 1,
                        "maximum": 200
                    }
                },
                "required": []
            }
        }
    }
]

def get_current_weather(latitude: float = None, longitude: float = None, address: str = None, district: str = None, state: str = None):
    return get_weather(latitude=latitude, longitude=longitude, address=address, district=district, state=state)

# Function Implementations
def get_weather_alerts(locations: List[str], include_details: bool = True) -> Dict[str, Any]:
    """
    Get weather alerts for specified locations
    """
    try:
        all_alerts = weather_api.fetch_alerts()
        
        if not all_alerts:
            return {
                "status": "error",
                "message": "No alerts found or error fetching data",
                "data": {}
            }
        
        results = {}
        total_alerts = 0
        
        for location in locations:
            search_terms = [term.strip() for term in location.split(',') if term.strip()]
            location_alerts = []
            
            for alert in all_alerts:
                area_desc = alert.get("area_description", "")
                if weather_api.location_matches(area_desc, search_terms):
                    if include_details:
                        # Full alert details
                        alert_data = {
                            "alert_id": alert.get("identifier"),
                            "severity": alert.get("severity"),
                            "severity_level": alert.get("severity_level"),
                            "severity_color": alert.get("severity_color"),
                            "disaster_type": alert.get("disaster_type"),
                            "area_description": alert.get("area_description"),
                            "start_time": alert.get("effective_start_time"),
                            "end_time": alert.get("effective_end_time"),
                            "source": alert.get("alert_source"),
                            "area_covered": alert.get("area_covered"),
                            "language": alert.get("actual_lang"),
                            "warning_message": alert.get("warning_message"),
                            "coordinates": alert.get("centroid"),
                            "disseminated": alert.get("disseminated")
                        }
                    else:
                        # Summary only
                        alert_data = {
                            "alert_id": alert.get("identifier"),
                            "severity": alert.get("severity"),
                            "disaster_type": alert.get("disaster_type"),
                            "area_description": alert.get("area_description"),
                            "start_time": alert.get("effective_start_time"),
                            "end_time": alert.get("effective_end_time")
                        }
                    
                    location_alerts.append(alert_data)
            
            results[location] = location_alerts
            total_alerts += len(location_alerts)
        
        return {
            "status": "success",
            "total_alerts": total_alerts,
            "locations_searched": len(locations),
            "data": results
        }
        
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error processing request: {str(e)}",
            "data": {}
        }

def get_alert_summary() -> Dict[str, Any]:
    """
    Get summary of all current alerts by severity
    """
    try:
        all_alerts = weather_api.fetch_alerts()
        
        if not all_alerts:
            return {
                "status": "error",
                "message": "No alerts found or error fetching data",
                "summary": {}
            }
        
        summary = {"WARNING": 0, "ALERT": 0, "WATCH": 0, "Yellow": 0, "Orange": 0, "Other": 0}
        disaster_types = {}
        states = {}
        
        for alert in all_alerts:
            # Count by severity
            severity = alert.get("severity", "Other")
            if severity in summary:
                summary[severity] += 1
            else:
                summary["Other"] += 1
            
            # Count by disaster type
            disaster_type = alert.get("disaster_type", "Unknown")
            disaster_types[disaster_type] = disaster_types.get(disaster_type, 0) + 1
            
            # Count by state (extract from area_description)
            area_desc = alert.get("area_description", "")
            if "districts of" in area_desc:
                state = area_desc.split("districts of")[-1].strip()
                states[state] = states.get(state, 0) + 1
        
        total = sum(summary.values())
        
        return {
            "status": "success",
            "total_alerts": total,
            "severity_breakdown": summary,
            "top_disaster_types": dict(sorted(disaster_types.items(), key=lambda x: x[1], reverse=True)[:10]),
            "top_affected_states": dict(sorted(states.items(), key=lambda x: x[1], reverse=True)[:10])
        }
        
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error getting summary: {str(e)}",
            "summary": {}
        }

def get_available_locations(limit: int = 50) -> Dict[str, Any]:
    """
    Get list of locations with current alerts
    """
    try:
        all_alerts = weather_api.fetch_alerts()
        
        if not all_alerts:
            return {
                "status": "error",
                "message": "No alerts found or error fetching data",
                "locations": []
            }
        
        locations = set()
        for alert in all_alerts:
            area_desc = alert.get("area_description", "").strip()
            if area_desc:
                locations.add(area_desc)
        
        sorted_locations = sorted(list(locations))[:limit]
        
        return {
            "status": "success",
            "total_available": len(locations),
            "returned": len(sorted_locations),
            "locations": sorted_locations
        }
        
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error getting locations: {str(e)}",
            "locations": []
        }

# Function dispatcher for OpenAI (Updated for tool_calls)
def execute_function(function_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
    """
    Execute the appropriate function based on OpenAI's tool call
    """
    if function_name == "get_weather_alerts":
        return get_weather_alerts(**arguments)
    elif function_name == "get_weather":
        return get_weather(**arguments)
    elif function_name == "get_alert_summary":
        return get_alert_summary(**arguments)
    elif function_name == "get_available_locations":
        return get_available_locations(**arguments)
    else:
        return {
            "status": "error",
            "message": f"Unknown function: {function_name}"
        }

def chat_with_weather_assistant(user_query: str = "", loop: bool = False, api_key: str = os.getenv("API_KEY"), api_url: str = os.getenv("API_URL"), model_name: str = os.getenv("MODEL_NAME")) -> None:
    """
    Chat with the weather alert assistant.
    
    Args:
        user_query (str, optional): The user's query about weather alerts. If empty, will prompt for input.
        loop (bool, optional): Whether to continue the conversation in a loop. 
                             If True, will keep asking for new queries until 'quit' is entered.
                             If False, will process a single query and return. Defaults to False.
    """
    if not api_key or not api_url or not model_name:
        raise ValueError("API key, API URL, and model name are required.")
        
    client = OpenAI(api_key=api_key, base_url=api_url)
    
    while True:
        # Get user input if not provided or in loop mode
        if not user_query:
            user_input = input("\nAsk about weather alerts (or 'quit' to exit): ")
            if user_input.lower() == 'quit':
                break
        else:
            user_input = user_query
            user_query = None  # Clear the query after first use
        
        try:
            # Initial API call
            response = client.chat.completions.create(
                model=model_name,
                messages=[
                    {
                        "role": "system", 
                        "content": "You are a weather information/alert assistant for India. Use the provided functions to help users get weather alert information. Always provide helpful and accurate information about weather alerts. Your respose should be in same language as user query. Response only should in paragraph only there should not be any table or list till user say to show table or list and response should be in easy language which can understand by uneducated person."
                    },
                    {"role": "user", "content": user_input}
                ],
                tools=WEATHER_TOOLS,
                tool_choice="auto"
            )
            
            message = response.choices[0].message
            
            if hasattr(message, 'tool_calls') and message.tool_calls:
                # Prepare messages for function execution
                messages = [
                    {
                        "role": "system", 
                        "content": "You are a weather information/alert assistant for India. Use the provided functions to help users get weather alert information. Always provide helpful and accurate information about weather alerts. Your response should be in the same language as the user query. Response should be in paragraph format only (no tables or lists unless explicitly requested) and in simple language that can be understood by everyone."
                    },
                    {"role": "user", "content": user_input},
                    message  # Add the assistant message with tool calls
                ]
                
                # Execute all tool calls
                for tool_call in message.tool_calls:
                    function_name = tool_call.function.name
                    function_args = json.loads(tool_call.function.arguments)
                    function_result = execute_function(function_name, function_args)
                    
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": json.dumps(function_result)
                    })
                
                # Get final response
                final_response = client.chat.completions.create(
                    model=model_name,
                    messages=messages
                )
                
                final_message = final_response.choices[0].message.content
                print(f"\n🤖 Assistant: {final_message}")
                
            else:
                print(f"\n🤖 Assistant: {message.content}")
            
            # If not in loop mode, break after one iteration
            if not loop:
                break
                
        except Exception as e:
            print(f"❌ Error: {e}")
            if not loop:
                break

if __name__ == "__main__":
    chat_with_weather_assistant()
    # print(WEATHER_TOOLS)
# get me weather alert of mumbai