import io
import requests
import pandas as pd
import matplotlib
matplotlib.use("Agg")  # headless backend for Hugging Face / servers
import matplotlib.pyplot as plt
import gradio as gr

GEOCODE_URL = "https://geocoding-api.open-meteo.com/v1/search"
FORECAST_URL = "https://api.open-meteo.com/v1/forecast"


def geocode_city(city: str):
    city = (city or "").strip()
    if not city:
        raise ValueError("Please enter a city name.")
    r = requests.get(GEOCODE_URL, params={"name": city, "count": 1, "language": "en"}, timeout=10)
    r.raise_for_status()
    data = r.json()
    if not data.get("results"):
        raise ValueError(f"City '{city}' not found.")
    res = data["results"][0]
    return float(res["latitude"]), float(res["longitude"]), f"{res['name']}, {res.get('country','')}"


def fetch_forecast(lat, lon, days, mode, temp_unit, precip_unit):
    params = {
        "latitude": lat,
        "longitude": lon,
        "forecast_days": int(days),
        "timezone": "auto",
        "temperature_unit": "fahrenheit" if temp_unit == "°F" else "celsius",
        "precipitation_unit": "inch" if precip_unit == "in" else "mm",
    }
    if mode == "Daily":
        params["daily"] = "temperature_2m_max,temperature_2m_min,precipitation_sum"
    else:
        params["hourly"] = "temperature_2m,relativehumidity_2m,precipitation"

    r = requests.get(FORECAST_URL, params=params, timeout=15)
    r.raise_for_status()
    return r.json()


def make_plot(x, ys, labels, title, ylabel):
    # ensure x is datetime-like for nice plotting
    try:
        x = pd.to_datetime(x)
    except Exception:
        # fall back to plain strings
        x = list(map(str, x))

    plt.figure(figsize=(8, 3.5))
    for y, lbl in zip(ys, labels):
        plt.plot(x, y, label=lbl)
    if labels:
        plt.legend()
    plt.title(title)
    plt.xlabel("Time")
    plt.ylabel(ylabel)
    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format="png", dpi=120)
    plt.close()
    buf.seek(0)
    return buf.getvalue()


def run(city, mode, days, temp_unit, precip_unit):
    try:
        # validate small things
        days = int(days)
        if days < 1 or days > 14:
            raise ValueError("Days must be between 1 and 14.")

        lat, lon, place = geocode_city(city)
        data = fetch_forecast(lat, lon, days, mode, temp_unit, precip_unit)

        if mode == "Daily":
            if "daily" not in data:
                raise ValueError("Daily data not available for this location.")
            d = data["daily"]
            df = pd.DataFrame({
                "date": d["time"],
                "temp_max": d["temperature_2m_max"],
                "temp_min": d["temperature_2m_min"],
                "precip": d["precipitation_sum"],
            })
            # return types: (markdown, DataFrame, bytes)
            img = make_plot(df["date"], [df["temp_max"], df["temp_min"]],
                            ["Max", "Min"], f"Daily Temperatures — {place}", f"Temp ({temp_unit})")
            md = f"### Daily forecast for **{place}**\nNext {days} day(s)\n\nUnits: Temperature **{temp_unit}**, Precipitation **{precip_unit}**"
            return md, df, img

        else:  # Hourly
            if "hourly" not in data:
                raise ValueError("Hourly data not available for this location.")
            h = data["hourly"]
            # Ensure we don't slice beyond available length
            max_hours = len(h.get("time", []))
            requested_hours = min(days * 24, max_hours)
            df = pd.DataFrame({
                "time": h["time"][:requested_hours]_]()