import gradio as gr
import requests
import os
from smolagents import tool, CodeAgent, HfApiModel, GoogleSearchTool

# Get API keys from environment
OPENWEATHER_API_KEY = os.environ.get("WEATHER_API_KEY")
SERPER_API_KEY = os.environ.get("SERPER_API_KEY")

@tool
def get_weather(lat: float, lon: float) -> dict | None:
    """
    Fetches current weather data from OpenWeatherMap API using geographic coordinates
    
    Args:
        lat: Latitude of the location (-90 to 90)
        lon: Longitude of the location (-180 to 180)
    
    Returns:
        dict | None: Weather data dictionary containing:
            - condition (str): Weather description
            - temperature (float): Temperature in Celsius
            - humidity (float): Humidity percentage
            - wind_speed (float): Wind speed in m/s
    """
    url = f"https://api.openweathermap.org/data/2.5/weather?lat={lat}&lon={lon}&appid={OPENWEATHER_API_KEY}&units=metric"
    try:
        response = requests.get(url)
        if response.status_code == 200:
            data = response.json()
            return {
                "condition": data["weather"][0]["description"],
                "temperature": data["main"]["temp"],
                "humidity": data["main"]["humidity"],
                "wind_speed": data["wind"]["speed"]
            }
        return None
    except Exception as e:
        print(f"Weather API Error: {e}")
        return None

weather_agent = CodeAgent(
    model=HfApiModel("Qwen/Qwen2.5-Coder-32B-Instruct", max_tokens=8096),
    tools=[GoogleSearchTool(provider="serper"), get_weather],
    name="Weather Expert",
    description="Processes weather queries through coordinate search and API calls"
)

def respond(message, history):
    try:
        system_prompt = """YOU ARE A WEATHER SPECIALIST:
        1. Analyze location from user input
        2. Google search "coordinates of [location]"
        3. Extract LAT/LON numeric values
        4. Call get_weather(LAT, LON)
        5. Respond in Vietnamese markdown
        
        NOTES:
        - Always verify data accuracy
        - Ask for clarification if coordinates not found
        - Units: °C, %, m/s"""
        
        response = weather_agent.run(
            f"[System] {system_prompt}\n[User] {message}"
        )
        
        return response if response else "Could not process request"
    
    except Exception as e:
        print(f"[DEBUG] {str(e)}")
        return "Sorry, I encountered an error processing your request"

with gr.Blocks(theme=gr.themes.Soft(), css="styles.css") as demo:
    gr.Markdown("# 🌦️ Weather Assistant")
    gr.Markdown("Enter any location to check weather conditions")

    chatbot = gr.Chatbot(
        elem_classes=["chatbot"],
    )
    
    # Giao diện chat
    chat_interface = gr.ChatInterface(
        respond,
        chatbot=chatbot,
        additional_inputs=None,
        submit_btn=gr.Button("Submit", variant="primary"),
    )

if __name__ == "__main__":
    demo.launch()