import gradio as gr
from datetime import datetime
from zoneinfo import ZoneInfo
import pycountry
import pytz
from timezonefinder import TimezoneFinder
from geopy.geocoders import Nominatim

# Initialize helpers
geolocator = Nominatim(user_agent="country_datetime_app")
tf = TimezoneFinder()

def get_country_datetime(location_name: str) -> dict:
    """
    Get the current date and time for a given country or city.
    
    - If a country name is provided: returns all its timezones.
    - If a city name is provided: returns the detected country + its local time.
    """

    location_name = location_name.strip()
    if not location_name:
        raise ValueError("Please enter a valid country or city name.")

    # Try to find as a country first
    try:
        country = pycountry.countries.get(name=location_name)
        if not country:
            country = pycountry.countries.search_fuzzy(location_name)[0]
        # Got a valid country — fetch all its timezones
        timezones = pytz.country_timezones.get(country.alpha_2, [])
        if not timezones:
            raise ValueError(f"No timezones found for {country.name}")
        result_type = "country"
    except Exception:
        # Not a country → try to geolocate as a city
        location = geolocator.geocode(location_name)
        if not location:
            raise ValueError(f"Location '{location_name}' not found.")
        tz_name = tf.timezone_at(lng=location.longitude, lat=location.latitude)
        if not tz_name:
            raise ValueError(f"Could not determine timezone for '{location_name}'.")
        try:
            timezone = ZoneInfo(tz_name)
            local_time = datetime.now(timezone)
            country_code = location.raw.get("address", {}).get("country_code", "").upper()
            country = pycountry.countries.get(alpha_2=country_code) if country_code else None
            country_name = country.name if country else "Unknown"
            return {
                "query_type": "city",
                "city": location_name,
                "country": country_name,
                "alpha_2": country_code,
                "datetime_info": [{
                    "timezone": tz_name,
                    "datetime": local_time.strftime("%Y-%m-%d %H:%M:%S"),
                    "utc_offset": local_time.strftime("%z"),
                    "timezone_abbr": local_time.strftime("%Z")
                }]
            }
        except Exception:
            raise ValueError(f"Error retrieving datetime for '{location_name}'.")

    # If it was a valid country
    datetime_info = []
    utc_now = datetime.now(ZoneInfo("UTC"))
    for tz_name in timezones:
        try:
            tz = ZoneInfo(tz_name)
            local_time = utc_now.astimezone(tz)
            datetime_info.append({
                "timezone": tz_name,
                "datetime": local_time.strftime("%Y-%m-%d %H:%M:%S"),
                "utc_offset": local_time.strftime("%z"),
                "timezone_abbr": local_time.strftime("%Z")
            })
        except Exception:
            continue

    return {
        "query_type": "country",
        "country": country.name,
        "alpha_2": country.alpha_2,
        "datetime_info": datetime_info
    }

# -------------------- GRADIO INTERFACE --------------------

def format_result(location):
    try:
        result = get_country_datetime(location)
        lines = []

        if result["query_type"] == "city":
            info = result["datetime_info"][0]
            lines.append(f"🏙️ **City:** {result['city']}")
            lines.append(f"🌍 **Country:** {result['country']} ({result['alpha_2']})")
            lines.append(f"🕒 **Local Time:** {info['datetime']} (UTC{info['utc_offset']}, {info['timezone_abbr']})")

        else:  # country mode
            lines.append(f"🌍 **Country:** {result['country']} ({result['alpha_2']})")
            lines.append("\n**Timezones:**")
            for info in result["datetime_info"]:
                lines.append(
                    f"- {info['timezone']} → {info['datetime']} "
                    f"(UTC{info['utc_offset']}, {info['timezone_abbr']})"
                )

        return "\n".join(lines)

    except ValueError as e:
        return f"❌ {str(e)}"
    except Exception as e:
        return f"⚠️ Unexpected error: {str(e)}"

app = gr.Interface(
    fn=format_result,
    inputs=gr.Textbox(label="Enter Country or City Name", placeholder="e.g. India, Kolkata, New York, Brazil"),
    outputs=gr.Markdown(label="Current Date & Time Information"),
    title="🌎 Smart DateTime Lookup",
    description="Enter any **country** or **city** name to get the current date and time automatically.",
    examples=[["India"], ["Kolkata"], ["Jaipur"], ["New York"], ["Sydney"], ["Russia"]]
)

if __name__ == "__main__":
    app.launch(mcp_server=True)