import os
import requests
from io import BytesIO
from PyPDF2 import PdfReader
import streamlit as st
import folium
from streamlit_folium import st_folium
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from groq import Groq
from dotenv import load_dotenv
load_dotenv()

# ------------------------------
# 🔐 API Keys
# ------------------------------
OPENWEATHER_API_KEY = os.environ.get("OPENWEATHER_API_KEY")
AQICN_TOKEN = os.environ.get("AQICN_TOKEN")
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

# ------------------------------
# ⚙️ Groq Client
# ------------------------------
@st.cache_resource
def get_groq_client():
    return Groq(api_key=GROQ_API_KEY)

client = get_groq_client()

# ------------------------------
# 📚 RAG Functions
# ------------------------------
@st.cache_data
def get_drive_links():
    return [
        "https://drive.google.com/file/d/1zEM3MxxpWcK9oaWc1GcmhamrGqgPjGig/view?usp=sharing",
        "https://drive.google.com/file/d/1Km2a_JzOPLQAyPvpkF3NjuXVU5RXwuZg/view?usp=sharing",
        "https://drive.google.com/file/d/1aj6QUMOy-6-idm3242yqxrgnt_7zjgjo/view?usp=sharing",
        "https://drive.google.com/file/d/1fmyAoEh_qRiHYvT78cmj4jWLcOcV4Y3K/view?usp=sharing",
        "https://drive.google.com/file/d/17-ZSypeJlhIc6B2hsJ_qm5xMMZP2xEWL/view?usp=sharing"
    ]

def download_pdf_from_drive(link):
    file_id = link.split("/d/")[1].split("/")[0]
    url = f"https://drive.google.com/uc?id={file_id}&export=download"
    response = requests.get(url)
    if response.status_code == 200:
        return BytesIO(response.content)
    raise Exception("Failed to download PDF.")

def extract_text_from_pdf(pdf_stream):
    reader = PdfReader(pdf_stream)
    text = ""
    for page in reader.pages:
        extracted = page.extract_text()
        if extracted:
            text += extracted
    return text

def chunk_text(text, chunk_size=500, chunk_overlap=50):
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size, chunk_overlap=chunk_overlap
    )
    return splitter.split_text(text)

@st.cache_data(show_spinner="📚 Loading and processing all PDFs...")
def process_all_documents():
    links = get_drive_links()
    all_chunks = []
    for link in links:
        pdf = download_pdf_from_drive(link)
        text = extract_text_from_pdf(pdf)
        chunks = chunk_text(text)
        all_chunks.extend(chunks)
    return all_chunks

@st.cache_resource(show_spinner="🔍 Creating FAISS vector store...")
def build_vector_db(chunks):
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    return FAISS.from_texts(chunks, embedding=embeddings)

def query_vector_db(query, vector_db):
    docs = vector_db.similarity_search(query, k=3)
    context = "\n".join([doc.page_content for doc in docs])
    completion = client.chat.completions.create(
        model="llama3-8b-8192",
        messages=[
            {"role": "system", "content": f"Use the following context:\n{context}"},
            {"role": "user", "content": query},
        ]
    )
    return completion.choices[0].message.content

# ------------------------------
# 🛰 Weather & Risk Functions
# ------------------------------
def get_coordinates(city):
    url = f"http://api.openweathermap.org/geo/1.0/direct?q={city}&limit=1&appid={OPENWEATHER_API_KEY}"
    res = requests.get(url).json()
    if res:
        return res[0]['lat'], res[0]['lon']
    return None, None

def get_weather(lat, lon):
    url = f"https://api.openweathermap.org/data/2.5/weather?lat={lat}&lon={lon}&appid={OPENWEATHER_API_KEY}&units=metric"
    res = requests.get(url).json()
    return res['main']['temp'], res['weather'][0]['description']

def get_aqi(lat, lon):
    url = f"https://api.waqi.info/feed/geo:{lat};{lon}/?token={AQICN_TOKEN}"
    res = requests.get(url).json()
    if res["status"] == "ok":
        return res["data"]["aqi"]
    return "Unavailable"

def get_glacier_risk(city):
    data = {
        "Skardu": ("High glacier melt risk", "July–August"),
        "Hunza": ("Moderate glacier melt risk", "June–August"),
        "Chitral": ("GLOF potential during peak summer", "July"),
    }
    return data.get(city.title(), ("No glacier melt alert", "N/A"))

def get_flood_risk(city):
    data = {
        "Karachi": ("Urban flood risk (monsoon)", "July–September"),
        "Lahore": ("Localized flooding in low-lying areas", "July–August"),
        "Sukkur": ("Flood risk due to Indus overflow", "August–September"),
    }
    return data.get(city.title(), ("No flood alert", "N/A"))

def get_climate_advice(city, season, age, interest):
    prompt = f"""As a climate adaptation advisor, suggest realistic, long-term actions someone can take in {city} during the {season} season to reduce their climate vulnerability.
The person is {age} years old and interested in {interest}. Consider water scarcity, heat, flooding, and energy challenges in {city}.
Make the advice practical, location-aware, and future-focused."""

    chat = client.chat.completions.create(
        model="llama3-8b-8192",
        messages=[{"role": "user", "content": prompt}]
    )
    return chat.choices[0].message.content

st.title("🌿 PakClimate (Climate Change Awareness App)")
# ------------------------------
# 🧭 Tabs
# ------------------------------
tabs = st.tabs(["🏠 Home", "🤖 RAG Chatbot", "🌍 Climate Intelligence Map", "🧠 Climate Adaptation Advisor"])

# ------------------------------
# 🏠 Home Tab
# ------------------------------
with tabs[0]:
    st.markdown("""
    ### 🧠 A platform to raise awareness about climate change in Pakistan 🇵🇰

    Pakistan is facing serious climate challenges:
    - 🌡️ Rising temperatures & heatwaves
    - 🌊 Urban flooding & glacier melt
    - 🚱 Water stress
    - 💨 Poor air quality

    These climate threats are:
    - 🧒 Endangering public health, especially for children and the elderly
    - 🌾 Disrupting agriculture and food security, reducing crop yields and threatening livelihoods
    - 🏙️ Overwhelming urban infrastructure, damaging homes, roads, and drainage systems
    - 💧 Depleting freshwater resources, making clean water scarce for millions
    - ⚡ Straining energy systems, increasing demand while reducing hydropower capacity
    - 🏞️ Threatening biodiversity and ecosystems, from mountains to coastlines
    - 🏠 Displacing communities, especially in flood-prone and drought-hit regions

    This app informs and empowers citizens to understand and act on climate risks using real-time data and long-term guidance.

    Let's make a difference!
    """)

    st.markdown("### The impact is powerfull...")

    image_files = ["img1.png", "img2.png", "img3.png", "img4.png", "img5.png"]
    cols = st.columns(5)
    for col, img in zip(cols, image_files):
        if os.path.exists(img):
            col.image(img, use_container_width=True)
        else:
            col.warning(f"Missing: {img}")

# ------------------------------
# 🤖 RAG Chatbot Tab
# ------------------------------
with tabs[1]:
    st.title("📄 Let's Learn about Climate Change")

    VECTOR_DB_PATH = "vector_store"

    @st.cache_resource(show_spinner="🔍 Loading FAISS vector store or building it if not found...")
    def get_or_load_vector_db():
        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
        if os.path.exists(VECTOR_DB_PATH):
            return FAISS.load_local(VECTOR_DB_PATH, embeddings, allow_dangerous_deserialization=True)
        chunks = process_all_documents()
        db = FAISS.from_texts(chunks, embedding=embeddings)
        db.save_local(VECTOR_DB_PATH)
        return db

    if "vector_db" not in st.session_state:
        st.session_state.vector_db = get_or_load_vector_db()

    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    st.markdown("#### 💡 Try asking one of these FAQs:")
    faqs = [
        "What is climate change?",
        "What are the recent disasters caused by climate change in Pakistan?",
        "What are the effects of climate change on natural resources of Pakistan?",
        "How does Government of Pakistan see the issue of climate change?"
    ]
    faq_cols = st.columns(2)
    for i, faq in enumerate(faqs):
        if faq_cols[i % 2].button(faq):
            st.session_state.faq_query = faq

    query = st.text_input("💬 Any other...", value=st.session_state.get("faq_query", ""))

    col1, col2 = st.columns([1, 4])
    with col1:
        if st.button("🧹 Clear Chat"):
            st.session_state.chat_history = []
            st.session_state.faq_query = ""
            query = ""

    if query:
        response = query_vector_db(query, st.session_state.vector_db)
        st.session_state.chat_history.append(("You", query))
        st.session_state.chat_history.append(("AI", response))
        st.session_state.faq_query = ""

    if st.session_state.chat_history:
        st.subheader("💬 Chat Conversation")
        for role, msg in st.session_state.chat_history:
            if role == "You":
                st.markdown(f"**🧑 {role}:** {msg}")
            else:
                st.markdown(f"**🤖 {role}:** {msg}")

# ------------------------------
# 🌍 Climate Map Tab
# ------------------------------
with tabs[2]:
    st.subheader("📍 Real-time Climate Risk & Weather Map")
    city = st.text_input("🏙️ Enter a city in Pakistan (e.g., Lahore, Skardu, Karachi):")

    if city:
        lat, lon = get_coordinates(city)
        if lat is not None:
            temp, weather = get_weather(lat, lon)
            aqi = get_aqi(lat, lon)
            flood_alert, flood_period = get_flood_risk(city)
            glacier_alert, glacier_period = get_glacier_risk(city)

            st.markdown(f"""
            **📍 City:** {city.title()}
            **🌡️ Temperature:** {temp}°C
            **🌤️ Weather:** {weather}
            **🌫️ AQI:** {aqi}

            **🌊 Flood Risk:** {flood_alert}
            ⏰ *Active during:* {flood_period}

            **🧊 Glacier Melt Risk:** {glacier_alert}
            ⏰ *Active during:* {glacier_period}
            """)

            m = folium.Map(location=[lat, lon], zoom_start=6, tiles="cartodbpositron")

            folium.raster_layers.WmsTileLayer(
                url='https://firms.modaps.eosdis.nasa.gov/wms/',
                name='FIRMS - Active Fires (MODIS)',
                layers='fires_modis',
                fmt='image/png',
                transparent=True,
                attr="NASA FIRMS"
            ).add_to(m)

            folium.raster_layers.WmsTileLayer(
                url='https://gibs.earthdata.nasa.gov/wms/epsg4326/best/wms.cgi?',
                name='MODIS Snow/Ice Daily',
                layers='MODIS_Terra_Snow_Cover_Daily',
                fmt='image/png',
                transparent=True,
                attr="NASA GIBS"
            ).add_to(m)

            folium.raster_layers.WmsTileLayer(
                url='https://sedac.ciesin.columbia.edu/geoserver/wms',
                name='Flood Hazard Frequency',
                layers='ndh:flood-hazard-frequency-distribution',
                fmt='image/png',
                transparent=True,
                attr="SEDAC"
            ).add_to(m)

            folium.Marker([lat, lon], tooltip=f"{city.title()} | Temp: {temp}°C | AQI: {aqi}",
                          icon=folium.Icon(color="blue")).add_to(m)

            for z in ["Skardu", "Hunza", "Chitral"]:
                z_lat, z_lon = get_coordinates(z)
                folium.Marker([z_lat, z_lon], tooltip=z, icon=folium.Icon(color="red")).add_to(m)

            for z in ["Karachi", "Lahore", "Sukkur"]:
                z_lat, z_lon = get_coordinates(z)
                folium.Marker([z_lat, z_lon], tooltip=z, icon=folium.Icon(color="orange")).add_to(m)

            folium.LayerControl().add_to(m)
            st_folium(m, width=700, height=500)
        else:
            st.error("❌ City not found. Please check the spelling.")

# ------------------------------
# 🧠 Climate Advisor Tab
# ------------------------------
with tabs[3]:
    st.markdown("### 🧠 Climate Adaptation Advisor")

    col1, col2 = st.columns(2)
    with col1:
        city = st.text_input("City", "Lahore")
        season = st.selectbox("Season", ["Summer", "Winter", "Monsoon", "Spring", "Autumn"])
    with col2:
        age = st.slider("Your Age", 10, 80, 25)
        interest = st.text_input("Your Interest", "Health")

    if st.button("Generate Advice"):
        advice = get_climate_advice(city, season, age, interest)
        st.success("✅ Climate Advice Generated:")
        st.write(advice)