from dataclasses import dataclass
import requests, ollama


# ── Dataclasses ───────────────────────────────────────────────

@dataclass
class LocationInfo:
    city: str
    country: str
    lat: float
    lon: float


@dataclass
class WeatherInfo:
    temperature: float
    condition: str
    humidity: int


# ── Tools ─────────────────────────────────────────────────────

def get_user_location() -> LocationInfo:
    """
    Returns the approximate current location of the user
    based on their IP address (city, country, coordinates).
    Use when the user refers to their current position
    with expressions like 'near me', 'here', 'where I am'.
    No arguments required.
    """
    r = requests.get("https://ipinfo.io/json", timeout=5)
    r.raise_for_status()
    d = r.json()
    lat, lon = d["loc"].split(",")
    return LocationInfo(
        city=d.get("city", "Unknown"),
        country=d.get("country", "Unknown"),
        lat=float(lat),
        lon=float(lon)
    )


def get_weather(city: str = None, lat: float = None, lon: float = None) -> WeatherInfo:
    """
    Returns the current weather conditions for a location.
    Use when the user asks about current weather, temperature,
    or atmospheric conditions.
    Args:
        city (str): City name, e.g. "Tokyo", "Rome". Use when the city
                     is known by name. Omit if coordinates are available.
        lat (float): Latitude. Use when coordinates are available (e.g.
                     from get_user_location). Omit if city name is known.
        lon (float): Longitude. Required if lat is provided.
    """
    location = f"{lat},{lon}" if lat is not None and lon is not None else city
    r = requests.get(f"https://wttr.in/{location}?format=j1", timeout=5)
    r.raise_for_status()
    d = r.json()["current_condition"][0]
    return WeatherInfo(
        temperature=float(d["temp_C"]),
        condition=d["weatherDesc"][0]["value"],
        humidity=int(d["humidity"])
    )


# ── Registry ──────────────────────────────────────────────────

TOOLS = {"get_user_location": get_user_location, "get_weather": get_weather}

TOOLS_SCHEMA = [
    {
        "type": "function",
        "function": {
            "name": "get_user_location",
            "description": (
                "Returns the approximate current location of the user "
                "based on IP address (city, country, coordinates). "
                "Use when the user refers to their current position "
                "with expressions like 'near me', 'here', 'where I am'."
            ),
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": (
                "Returns the current weather conditions for a location. "
                "Use when the user asks about current weather, temperature, "
                "or atmospheric conditions."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "City name, e.g. 'Tokyo', 'Rome'. "
                                       "Use when the city is known by name. "
                                       "Omit if coordinates are available."
                    },
                    "lat": {
                        "type": "number",
                        "description": "Latitude. Use when coordinates are available "
                                       "(e.g. from get_user_location). "
                                       "Omit if city name is known."
                    },
                    "lon": {
                        "type": "number",
                        "description": "Longitude. Required if lat is provided."
                    }
                },
                "required": []
            }
        }
    }
]


# ── ReAct Loop ──────────────────────────────────────────────

def run(user_message: str, model: str = "qwen3.5:4b") -> str:
    context = [{"role": "user", "content": user_message}]
    while True:
        msg = ollama.chat(model=model, messages=context, tools=TOOLS_SCHEMA)["message"]
        if not msg.get("tool_calls"):
            return msg["content"]
        context.append(msg)
        for call in msg["tool_calls"]:
            name, args = call["function"]["name"], call["function"]["arguments"]
            print(f"[Action] {name}({args})")
            result = TOOLS[name](**args)
            print(f"[Observation] {result}")
            context.append({"role": "tool", "content": str(result)})


# ── Entry point ───────────────────────────────────────────────

if __name__ == "__main__":
    print(run("What's the weather like near me?"))
    print(run("What's the weather like in Tokyo?"))