import aiohttp
import json
import os
import asyncio
from typing import Dict, Any, Union, Optional, List

# Assuming Tool_node is available from the HeGFlow framework as per the provided context.
# from HeGFlow.graph.tool_node import Tool_node


class WeatherFetcher(Tool_node):
    """
    A tool for fetching current weather information from OpenWeatherMap.
    It requires an API key for OpenWeatherMap.
    """
    def __init__(self):
        super().__init__(
            name="WeatherFetcher",
            description="Fetches current weather information for a specified city. Requires 'city' parameter, e.g., 'city=\"London\"'."
        )
        # It's highly recommended to load API keys from environment variables
        # rather than hardcoding them for security and flexibility.
        # For this example, replace "YOUR_OPENWEATHERMAP_API_KEY" with an actual key
        # or set it as an environment variable OPENWEATHERMAP_API_KEY.
        self.api_key: str = os.getenv("OPENWEATHERMAP_API_KEY", "YOUR_OPENWEATHERMAP_API_KEY")
        self.base_url: str = "http://api.openweathermap.org/data/2.5/weather"

    async def _execute_tool(self, task: str) -> str:
        """
        Executes a weather data fetch based on the provided `task` string.
        The `task` must specify a city, e.g., "city='London'".
        
        Returns:
            A JSON string containing the raw weather data, or an error message string.
        """
        city: Optional[str] = None
        
        # Attempt to parse the city from the task string
        city_match = re.search(r"city=['\"]([^'\"]+)['\"]", task)

        if city_match:
            city = city_match.group(1)
        else:
            return "Error: No city specified. Please provide a city, e.g., 'city=\"New York\"'."
        
        if self.api_key == "YOUR_OPENWEATHERMAP_API_KEY":
            return "Error: OpenWeatherMap API key is not set. Please replace 'YOUR_OPENWEATHERMAP_API_KEY' or set the 'OPENWEATHERMAP_API_KEY' environment variable."

        params = {
            "q": city,
            "appid": self.api_key,
            "units": "metric"  # Use "imperial" for Fahrenheit, "metric" for Celsius
        }

        async with aiohttp.ClientSession() as session:
            try:
                async with session.get(self.base_url, params=params) as response:
                    data = await response.json()

                    if response.status == 200:
                        return json.dumps(data, indent=2)
                    else:
                        error_message = data.get("message", "Unknown error")
                        return f"Error fetching weather for '{city}': {response.status} - {error_message}"
            except aiohttp.ClientError as e:
                return f"Network or client error during weather API call: {e}"
            except json.JSONDecodeError:
                return f"Error: Failed to decode JSON response from weather API for '{city}'."
            except Exception as e:
                return f"An unexpected error occurred: {type(e).__name__}: {e}"

    def _format_result_to_natural_language(self, raw_result: str, task_description: str) -> str:
        """
        Converts the raw JSON string output from _execute_tool into a human-readable string.
        """
        try:
            parsed_data: Dict[str, Any] = json.loads(raw_result)
        except json.JSONDecodeError:
            # If the raw_result is not valid JSON, it's likely an error message
            return f"Failed to interpret weather results for '{task_description}'. Raw output: {raw_result}"

        if "main" not in parsed_data or "weather" not in parsed_data:
            # This indicates an error message from _execute_tool, or an unexpected API response structure
            return raw_result # Return the raw error message directly

        city_name = parsed_data.get("name", "N/A")
        country_code = parsed_data.get("sys", {}).get("country", "N/A")
        
        temperature = parsed_data["main"].get("temp", "N/A")
        feels_like = parsed_data["main"].get("feels_like", "N/A")
        humidity = parsed_data["main"].get("humidity", "N/A")
        pressure = parsed_data["main"].get("pressure", "N/A")
        
        weather_description = parsed_data["weather"][0].get("description", "N/A")
        wind_speed = parsed_data["wind"].get("speed", "N/A") # in meters/sec if units=metric
        
        formatted_output = (
            f"Current weather for {city_name}, {country_code}:\n"
            f"Description: {weather_description.capitalize()}\n"
            f"Temperature: {temperature}°C (Feels like: {feels_like}°C)\n"
            f"Humidity: {humidity}%\n"
            f"Pressure: {pressure} hPa\n"
            f"Wind Speed: {wind_speed} m/s"
        )
        return formatted_output