import os
import re
import requests
import torch
import gradio as gr
from transformers import AutoTokenizer, pipeline

# --- Model Setup
MODEL_ID = "Aaayushiii/mt5-crop-lora"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# Prefer GPU if available
device = 0 if torch.cuda.is_available() else -1

# Build the pipeline (text2text for mT5)
pipe = pipeline(
    "text2text-generation",
    model=MODEL_ID,
    tokenizer=tokenizer,
    device=device,
)

# --- Weather (OpenWeather key or fallback)
OPENWEATHER_API_KEY = os.getenv("OPENWEATHER_API_KEY", "").strip()

def get_weather(location: str) -> str | None:
    if OPENWEATHER_API_KEY:
        try:
            url = (
                "http://api.openweathermap.org/data/2.5/weather"
                f"?q={location}&appid={OPENWEATHER_API_KEY}&units=metric"
            )
            r = requests.get(url, timeout=10).json()
            if r.get("cod") == 200:
                return f"📍 {r['name']} | 🌡️ {r['main']['temp']}°C | {r['weather'][0]['description'].capitalize()}"
        except Exception:
            pass
    # Fallback Open-Meteo
    try:
        geo = requests.get(
            f"https://geocoding-api.open-meteo.com/v1/search?name={location}&count=1",
            timeout=10,
        ).json()
        if not geo.get("results"):
            return None
        lat, lon = geo["results"][0]["latitude"], geo["results"][0]["longitude"]
        city = geo["results"][0]["name"]
        wx = requests.get(
            f"https://api.open-meteo.com/v1/forecast?latitude={lat}&longitude={lon}&current_weather=true",
            timeout=10,
        ).json()
        cw = wx.get("current_weather")
        if cw:
            return f"📍 {city} | 🌡️ {cw['temperature']}°C | 💨 {cw['windspeed']} km/h"
    except Exception:
        pass
    return None

import requests

def extract_location(text: str) -> str | None:
    caps = re.findall(r"\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b", text)
    return max(caps, key=len) if caps else None



def crop_recommendation(query: str, weather: str | None = None) -> str:
    prompt = f"User: {query}\n"
    if weather:
        prompt += f"(Weather: {weather})\n"
    prompt += "Assistant: Provide crop recommendation including top options, weather suitability, and profitability."
    out = pipe(
        prompt,
        max_new_tokens=150,
        num_beams=3,
        do_sample=True,
        temperature=0.7,
    )[0]["generated_text"]
    return out

def agent(query: str) -> str:
    loc = extract_location(query)
    if "weather" in query.lower():
        if not loc:
            return "⚠️ Please specify a location."
        return get_weather(loc) or f"⚠️ Weather not found for {loc}."
    if loc:
        weather = get_weather(loc)
        return crop_recommendation(query, weather)
    return crop_recommendation(query)

# --- UI
img_path = "crop.png"
examples = [
    "Weather in Delhi",
    "Best crop in Agra during Kharif",
]

dark_css = """
body {
    background: #121212;
    font-family: 'Inter', sans-serif;
    color: #f5f5f5;
    display: flex;
    justify-content: center;
}
.main-card {
    max-width: 800px;
    margin: 40px auto;
    padding: 24px;
    background: #1e1e1e;
    border-radius: 20px;
    box-shadow: 0 8px 30px rgba(0,0,0,0.8);
    border: 1px solid #333;
}
.chatbox {
    background: #2a2a2a;
    border-radius: 16px;
    padding: 12px;
    border: 1px solid #444;
    min-height: 350px;
    color: #f5f5f5;
}
.gr-button {
    background: linear-gradient(to right bottom, #444, #777) !important;
    color: white !important;
    border-radius: 14px;
    font-weight: 600;
    padding: 10px 18px;
    transition: all 0.25s ease-in-out;
    margin-top: 10px;
}
.gr-button:hover {
    transform: translateY(-2px) scale(1.05);
    box-shadow: 0 0 12px rgba(200,200,200,0.4);
}
.gr-textbox textarea {
    border: 2px solid #555;
    border-radius: 14px;
    padding: 12px;
    font-size: 15px;
    background: #1a1a1a;
    color: #f5f5f5;
}
h1 {
    font-size: 32px;
    font-weight: 700;
    margin-bottom: 20px;
    text-align: center;
    background: linear-gradient(to right bottom, #bbb, #eee);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
#example-row .examples {
    display: flex;
    justify-content: center;
    flex-wrap: wrap;
    gap: 10px;
    margin: 15px 0 20px 0;
}
#example-row .example {
    border-radius: 20px;
    background: #333;
    color: #f5f5f5;
    font-size: 14px;
    padding: 6px 16px;
    cursor: pointer;
    transition: all 0.25s ease;
    box-shadow: 0 4px 12px rgba(0,0,0,0.2);
}
#example-row .example:hover {
    background: #555;
    transform: translateY(-2px) scale(1.06);
    box-shadow: 0 0 15px rgba(200,200,200,0.4);
}
"""
with gr.Blocks(css=dark_css) as demo:
    with gr.Column(elem_classes="main-card"):
        gr.Markdown("<h1>🌱 AgriBot – Crop + Weather Assistant</h1>")
        chatbot = gr.Chatbot(label="Chat with AgriBot", elem_classes="chatbox", height=400)
        user_in = gr.Textbox(placeholder="Ask something, e.g., 'Weather in Jaipur?'")
        gr.Examples(examples=examples, inputs=user_in, elem_id="example-row")
        submit = gr.Button("Ask")

        if os.path.exists(img_path):
            gr.Image(img_path, show_label=False)
        gr.Markdown("""
        ### 🍀 Features
        - 🌦️ Live weather lookup
        - 🌾 Crop recommendations
        - 💰 Profitability insights
        """)

    def respond(history, msg):
        reply = agent(msg)
        history.append((msg, reply))
        return history, None

    submit.click(respond, [chatbot, user_in], [chatbot, user_in])
    user_in.submit(respond, [chatbot, user_in], [chatbot, user_in])

if __name__ == "__main__":
    demo.launch()