import requests
import gradio as gr
from groq import Groq
import os

# =========================
# Load GROQ API Key from environment (Hugging Face Secrets)
# =========================
GROQ_API_KEY = os.getenv("GROQ_API_KEY")  # set in HF Secrets
if not GROQ_API_KEY:
    raise ValueError("❌ GROQ_API_KEY not found in environment!")

client = Groq(api_key=GROQ_API_KEY)

# =========================
# Common country aliases
# =========================
ALIASES = {
    "usa": "united states",
    "us": "united states",
    "uk": "united kingdom",
    "south korea": "korea",
    "north korea": "korea, north",
    "russia": "russian federation",
}

# =========================
# Fetch Country Info (Robust)
# =========================
def fetch_country(country_name):
    country_name = country_name.strip().lower()
    country_name = ALIASES.get(country_name, country_name)

    try:
        url = f"https://restcountries.com/v3.1/name/{country_name}"
        response = requests.get(url, timeout=10)
        data = response.json()

        if not data:
            return None, "❌ No country data returned."

        # Try exact match first
        country = None
        for c in data:
            if c["name"]["common"].lower() == country_name:
                country = c
                break
        if not country:
            # fallback to first with valid population
            for c in data:
                if c.get("population"):
                    country = c
                    break
        if not country:
            country = data[0]  # final fallback

        population = country.get("population")
        if not population:
            population = "N/A"
        else:
            population = f"{population:,}"

        return {
            "name": country["name"]["common"],
            "capital": country.get("capital", ["N/A"])[0],
            "population": population,
            "region": country.get("region", "N/A")
        }, None

    except Exception as e:
        return None, f"❌ API Error: {str(e)}"

# =========================
# Generate AI Summary
# =========================
def generate_summary(info):
    try:
        prompt = f"""
        Write a short, friendly travel-style summary for:
        Country: {info['name']}
        Capital: {info['capital']}
        Region: {info['region']}
        """

        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=[{"role": "user", "content": prompt}]
        )

        return response.choices[0].message.content.strip()

    except Exception as e:
        return f"❌ AI Error: {str(e)}"

# =========================
# Main Logic
# =========================
def search_country(query):
    if not query.strip():
        return "⚠️ Please enter a country name."

    info, error = fetch_country(query)
    if error:
        return error

    summary = generate_summary(info)

    return f"""
### 🌍 {info['name']}

**🏛 Capital:** {info['capital']}  
**👥 Population:** {info['population']}  
**🗺 Region:** {info['region']}

---

### 🧠 AI Travel Summary
{summary}
"""

# =========================
# Gradio UI
# =========================
with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown(
        "<h1 style='text-align:center;'>🌍 AI Country Explorer</h1>"
        "<p style='text-align:center; font-size:16px;'>Discover countries with real data + AI-powered summaries</p>"
    )

    with gr.Row():
        country_input = gr.Textbox(
            placeholder="Type a country name (e.g. Malaysia, Indonesia, USA)",
            label="Country"
        )

    search_btn = gr.Button("🔍 Explore Country")
    output = gr.Markdown()

    search_btn.click(search_country, country_input, output)
    country_input.submit(search_country, country_input, output)

demo.launch()