import json
import litellm
from smolagents import CodeAgent
from tools.tools import (
    get_coordinates,
    get_earthquake_data,
    get_nasa_fire_data,
    get_full_daily_forecast,
    climate_change_data,
    get_full_air_quality_forecast,
    get_full_marine_daily_forecast,
    get_full_flood_daily_forecast,
    get_full_satellite_radiation,
)

class RiskAnalysisAgent:
    """Agent responsible for analyzing risks based on user input and available data."""

    def __init__(self, model):
        self.agent = CodeAgent(
            tools=[
                get_coordinates,
                get_earthquake_data,
                get_nasa_fire_data,
                get_full_daily_forecast,
                climate_change_data,
                get_full_air_quality_forecast,
                get_full_marine_daily_forecast,
                get_full_flood_daily_forecast,
                get_full_satellite_radiation,
            ],
            model=model,
            additional_authorized_imports=["json", "datetime", "math"],
        )

    def analyze_risks(self, user_query: str) -> dict:
        """Analyze risks based on user query."""

        analysis_prompt = f"""
You are an expert climate risk analyst. A user has submitted this query: "{user_query}"

Your task:
1. Extract the location from the query
2. Identify what types of risks they're asking about
3. Gather relevant data using the available tools:
   - Use get_full_daily_forecast for detailed weather data
   - Use climate_change_data for long-term climate projections
   - Use get_full_air_quality_forecast for air quality risks
   - Use get_full_marine_daily_forecast for coastal/marine risks
   - Use get_full_flood_daily_forecast for detailed flood data
   - Use get_full_satellite_radiation for solar radiation data
4. Analyze the risk levels for each identified hazard
5. Return a comprehensive risk analysis

For each risk type you identify, provide:
- Risk level (0-100 scale)
- Key factors contributing to the risk
- Time horizon (immediate, short-term, long-term)
- Confidence level in your assessment
- Historical context (when available)
- Future projections (when available)

Focus on being thorough but concise. Use your judgment to determine which data sources are most relevant.

IMPORTANT: Return a valid Python dictionary (not JSON string) in this exact format:
{{
    "location": {{"city": "CityName", "country": "CountryName", "lat": 0.0, "lon": 0.0}},
    "identified_risks": ["risk1", "risk2"],
    "risk_analysis": {{
        "earthquake": {{
            "risk_level": 25,
            "contributing_factors": ["seismic activity", "building codes"],
            "time_horizon": "long-term",
            "Detailed analysis": "",
            "confidence": "medium",
            "key_insights": "Moderate earthquake risk due to regional seismic activity",
            "historical_context": "Historical earthquake data analysis",
            "future_projections": "Projected seismic activity trends"
        }},
        "wildfire": {{
            "risk_level": 60,
            "contributing_factors": ["dry conditions", "vegetation"],
            "time_horizon": "immediate",
            "Detailed analysis": "",
            "confidence": "high",
            "key_insights": "High wildfire risk during dry season",
            "historical_context": "Past wildfire patterns",
            "future_projections": "Climate change impact on fire risk"
        }},
        "climate": {{
            "risk_level": 45,
            "contributing_factors": ["temperature trends", "precipitation changes"],
            "time_horizon": "long-term",
            "confidence": "high",
            "Detailed analysis": "",
            "key_insights": "Significant climate change impacts expected",
            "future_projections": "Climate model predictions"
        }}
    }},
    "overall_assessment": "Very detailed overall risk summary here"
}}
"""

        try:

            try:
                response = self.agent.run(analysis_prompt)
            except Exception as e:
               print(e)

            #response = litellm.completion(messages=analysis_prompt, model="anthropic/claude-sonnet-4-20250514")

            if isinstance(response, dict):
                return response
            elif isinstance(response, str):
                try:
                    return json.loads(response)
                except json.JSONDecodeError:
                    return {
                        "location": {
                            "city": "Unknown",
                            "country": "Unknown",
                            "lat": 0.0,
                            "lon": 0.0,
                        },
                        "identified_risks": ["general climate risks"],
                        "risk_analysis": {
                            "general": {
                                "risk_level": 30,
                                "contributing_factors": ["climate change"],
                                "time_horizon": "long-term",
                                "confidence": "medium",
                                "key_insights": f"Analysis of query: {user_query}",
                            }
                        },
                        "overall_assessment": f"Climate risk analysis for: {user_query}",
                    }
            else:
                return {
                    "location": {
                        "city": "Unknown",
                        "country": "Unknown",
                        "lat": 0.0,
                        "lon": 0.0,
                    },
                    "identified_risks": ["climate risks"],
                    "risk_analysis": {
                        "general": {
                            "risk_level": 30,
                            "contributing_factors": ["climate factors"],
                            "time_horizon": "medium-term",
                            "confidence": "medium",
                            "key_insights": "General climate risk assessment",
                        }
                    },
                    "overall_assessment": "Basic climate risk analysis completed",
                }
        except Exception as e:
            print(f"Risk analysis error: {e}")
            return {
                "error": f"Risk analysis failed: {str(e)}",
                "location": {
                    "city": "Unknown",
                    "country": "Unknown",
                    "lat": 0.0,
                    "lon": 0.0,
                },
                "identified_risks": [],
                "risk_analysis": {},
                "overall_assessment": "Analysis could not be completed",
            }