# Uncomment to install
# !pip install -q "langchain>=1.0" langchain-openai requests

import requests
from langchain_core.tools import tool


@tool
def get_weather(city: str) -> str:
    """
    Get the CURRENT weather right now for a given city.
    Use this for questions about what to wear or do right now / today.
    Input should be a city name, e.g. 'New Delhi'.
    """
    url = f"https://wttr.in/{city}?format=j1"
    response = requests.get(url, timeout=10)
    response.raise_for_status()
    data = response.json()

    current = data["current_condition"][0]
    temp_c = current["temp_C"]
    feels_like_c = current["FeelsLikeC"]
    description = current["weatherDesc"][0]["value"]

    return (
        f"Current weather in {city}: {temp_c}°C, "
        f"feels like {feels_like_c}°C, conditions: {description}."
    )

@tool
def get_forecast(city: str, days_ahead: int = 1) -> str:
    """
    Get the weather FORECAST for a given city, for a future day.
    Use this for questions about tomorrow, this weekend, or any future day —
    NOT for questions about right now (use get_weather for that).

    Args:
        city: city name, e.g. 'New Delhi'.
        days_ahead: how many days from today (0 = today, 1 = tomorrow, 2 = day after). Max 2.
    """
    days_ahead = max(0, min(days_ahead, 2))  # wttr.in's free tier gives today + 2 days

    url = f"https://wttr.in/{city}?format=j1"
    response = requests.get(url, timeout=10)
    response.raise_for_status()
    data = response.json()

    day = data["weather"][days_ahead]
    date = day["date"]
    max_temp = day["maxtempC"]
    min_temp = day["mintempC"]

    # Average the hourly chance-of-rain readings for a simple daily estimate
    rain_chances = [int(hour["chanceofrain"]) for hour in day["hourly"]]
    avg_rain_chance = sum(rain_chances) // len(rain_chances)

    return (
        f"Forecast for {city} on {date} ({days_ahead} day(s) ahead): "
        f"high {max_temp}°C, low {min_temp}°C, "
        f"average chance of rain: {avg_rain_chance}%."
    )

from langchain.agents import create_agent
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

tools = [get_weather, get_forecast]

SYSTEM_PROMPT = (
    "You are a helpful weather assistant. Always check real weather data before "
    "answering questions about clothing, umbrellas, or outdoor plans. Choose "
    "between the current-weather tool and the forecast tool based on whether "
    "the question is about right now or a future day. Clearly explain your "
    "recommendation using the actual data you observed."
)

agent = create_agent(
    model=llm,
    tools=tools, 
    system_prompt=SYSTEM_PROMPT # OPTIONAL: Agent can figure the tools to use based on docstring
)

def ask(question: str):
    result = agent.invoke({"messages": [{"role": "user", "content": question}]})

    print(f"QUESTION: {question}\n")

    for msg in result["messages"]:
        if type(msg).__name__ == "AIMessage" and getattr(msg, "tool_calls", None):
            for call in msg.tool_calls:
                print(f"  -> Agent chose tool: {call['name']}  args={call['args']}")
        elif type(msg).__name__ == "ToolMessage":
            print(f"  -> Observation: {msg.content}")

    print(f"\nFINAL ANSWER:\n{result['messages'][-1].content}")
    print("\n" + "=" * 70 + "\n")
    return result