import os
import gradio as gr
import numpy as np
import plotly.graph_objects as go
import requests
import cv2
from PIL import Image
import qrcode
from fpdf import FPDF
from sentinelhub import SHConfig
from groq import Groq
import google.generativeai as genai
import tempfile
import time  # For rate limit handling

# -------------------- ENVIRONMENT VARIABLES --------------------
HF_API_KEY = os.getenv("HF_API_KEY")
GROQ_API_KEY = "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"
GEMINI_API_KEY = "AIzaSyCqPnhDNwBP6Tsw1wkLGdXCIVDnNO44swY"
SENTINEL_CLIENT_ID = os.getenv("SENTINEL_CLIENT_ID")
SENTINEL_CLIENT_SECRET = os.getenv("SENTINEL_CLIENT_SECRET")

# -------------------- SENTINEL CONFIG --------------------
config = SHConfig()
if SENTINEL_CLIENT_ID and SENTINEL_CLIENT_SECRET:
    config.client_id = SENTINEL_CLIENT_ID
    config.client_secret = SENTINEL_CLIENT_SECRET

# -------------------- AI FUNCTIONS --------------------
def gemini_summary(text):
    try:
        if not GEMINI_API_KEY: return None, "Missing Key"
        genai.configure(api_key=GEMINI_API_KEY)
        model = genai.GenerativeModel('gemini-1.5-flash') 
        response = model.generate_content(text)
        return response.text, None
    except Exception as e:
        return None, str(e)

def groq_summary(text):
    try:
        if not GROQ_API_KEY: return None, "Missing Key"
        client = Groq(api_key=GROQ_API_KEY)
        completion = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "user", "content": text}]
        )
        return completion.choices[0].message.content, None
    except Exception as e:
        return None, str(e)

def hf_summary(text):
    try:
        url = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
        headers = {"Authorization": f"Bearer {HF_API_KEY}"}
        payload = {
            "inputs": f"<|system|>You are a scientist.</s><|user|>{text}</s><|assistant|>",
            "parameters": {"max_new_tokens": 800}
        }
        r = requests.post(url, headers=headers, json=payload, timeout=25)
        if r.status_code == 200:
            return r.json()[0]["generated_text"].split("<|assistant|>")[-1], None
        else:
            return None, f"Status {r.status_code}: {r.text}"
    except Exception as e:
        return None, str(e)

def smart_summary(text):
    errors = []
    out, err = groq_summary(text)
    if out: return out
    errors.append(f"Groq: {err}")
    out, err = gemini_summary(text)
    if out: return out
    errors.append(f"Gemini: {err}")
    if HF_API_KEY:
        out, err = hf_summary(text)
        if out: return out
        errors.append(f"HF: {err}")
    return "⚠ SYSTEM FAILURE. DEBUG LOG:\n" + "\n".join(errors)

# -------------------- AUDIO FUNCTION (IMPROVED) --------------------
def generate_audio_report(text):
    try:
        from gtts import gTTS
    except ImportError:
        raise gr.Error("❌ Library Missing! Add 'gTTS' to requirements.txt")

    if not text or len(text.strip()) < 5:
        raise gr.Error("❌ Report is too short or empty. Generate a report first!")

    # Clean text of Markdown symbols for better speech
    clean_text = text.replace("**", "").replace("#", "")

    try:
        # Retry logic for rate limits
        for i in range(3):
            try:
                tts = gTTS(text=clean_text[:1000], lang='en')
                with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as f:
                    tts.save(f.name)
                    return f.name
            except Exception as e:
                if "429" in str(e) and i < 2:
                    time.sleep(2) # Wait 2 seconds and try again
                    continue
                raise e
    except Exception as e:
        raise gr.Error(f"Speech Generation Error: {str(e)}")

# -------------------- MATH & LOGIC --------------------
def calculate_wqi(pH, do, nutrients):
    wqi = (7 - abs(7 - pH)) * 0.2 + (do/14) * 0.5 + (10 - nutrients) * 0.3
    wqi_score = max(0, min(100, int(wqi*10)))
    return wqi_score

def calculate_hsi(flow_rate, temp, sediment):
    hsi = 100 - abs(flow_rate-50)*0.5 - abs(temp-20)*2 - sediment*1.5
    return max(0, min(100, int(hsi)))

def calculate_erosion(sediment, construction):
    score = sediment*1.5 + construction*2
    return max(0, min(100, int(score)))

def potability_status(wqi):
    if wqi > 80: return "Safe"
    elif wqi > 50: return "Boil Required"
    else: return "Toxic"

def river_stability(wqi, hsi, erosion):
    return int((wqi*0.4 + hsi*0.4 + (100-erosion)*0.2))

def analyze_satellite_image(img):
    if img is None: return 0
    img_array = np.array(img.convert("L"))
    turbidity_score = int(np.mean(img_array)/2.55)
    return turbidity_score

# -------------------- VISUALS & INSIGHTS --------------------
def create_plots(wqi, hsi, erosion, turbidity):
    fig = go.Figure()
    colors = ['#0061ff', '#60efff', '#ff4b4b', '#ffb347']
    fig.add_trace(go.Bar(name="Metrics", x=["WQI", "HSI", "Erosion", "Turbidity"],
                          y=[wqi, hsi, erosion, turbidity], marker_color=colors))
    fig.update_layout(title="River Health Metrics", yaxis=dict(range=[0,100]), template="plotly_white")
    return fig

def generate_graph_insights(wqi, hsi, erosion, turbidity):
    text = "### 📉 Graph Analysis\n\n"
    if wqi > 70: text += f"🔵 **Water Quality:** {wqi}/100. Excellent condition.\n\n"
    elif wqi > 40: text += f"🔵 **Water Quality:** {wqi}/100. Moderate pollution.\n\n"
    else: text += f"🔵 **Water Quality:** {wqi}/100. **CRITICAL**.\n\n"
    
    if hsi > 70: text += f"🟢 **Habitat:** {hsi}/100. Good biodiversity.\n\n"
    else: text += f"🟢 **Habitat:** {hsi}/100. Poor conditions.\n\n"
    return text

# -------------------- PDF ENGINE --------------------
def generate_pdf(wqi, hsi, erosion, turbidity, summary_text):
    pdf = FPDF()
    pdf.add_page()
    qr = qrcode.QRCode(box_size=3)
    qr.add_data(f"Verified FlumenIntel Report | WQI: {wqi}")
    qr.make(fit=True)
    img = qr.make_image(fill_color="black", back_color="white")
    
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
        img.save(tmp.name)
        pdf.image(tmp.name, x=165, y=10, w=30)
    
    pdf.set_y(15)
    pdf.set_font("Arial", "B", 24)
    pdf.set_text_color(0, 97, 255)
    pdf.cell(0, 10, "FlumenIntel", ln=True, align='L')
    pdf.ln(10)
    pdf.set_font("Arial", "", 12)
    pdf.set_text_color(0, 0, 0)
    
    clean_text = summary_text.encode('latin-1', 'replace').decode('latin-1')
    pdf.multi_cell(0, 6, clean_text)
    
    report_path = os.path.join(tempfile.gettempdir(), "FlumenIntel_Report.pdf")
    pdf.output(report_path)
    return report_path

# -------------------- MAIN PROCESSOR --------------------
def process_data(flow_rate, water_temp, sediment, construction, pH, do, nutrients, sat_img):
    try:
        wqi = calculate_wqi(pH, do, nutrients)
        hsi = calculate_hsi(flow_rate, water_temp, sediment)
        erosion = calculate_erosion(sediment, construction)
        turbidity = analyze_satellite_image(sat_img)
        stability = river_stability(wqi, hsi, erosion)
        potability = potability_status(wqi)
        
        prompt = f"Write a professional health report for a river. WQI: {wqi}, HSI: {hsi}, Erosion: {erosion}, Turbidity: {turbidity}. Potability: {potability}."
        summary = smart_summary(prompt)
        
        fig = create_plots(wqi, hsi, erosion, turbidity)
        graph_text = generate_graph_insights(wqi, hsi, erosion, turbidity)
        
        pdf_path = generate_pdf(wqi, hsi, erosion, turbidity, summary)
        
        status_text = f"Stability Index: {stability}/100\nStatus: {potability}"
        return status_text, fig, graph_text, summary, pdf_path

    except Exception as e:
        return str(e), None, "", "", None

def run_app(flow, temp, sediment, construction, ph, do, nutrients, sat_img):
    return process_data(flow, temp, sediment, construction, ph, do, nutrients, sat_img)

# -------------------- UI DESIGN (WITH SCROLLBAR) --------------------
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;600&display=swap');
* { font-family: 'Poppins', sans-serif !important; }
#title-box { background: linear-gradient(135deg, #0061ff 0%, #60efff 100%); color: white; padding: 20px; border-radius: 12px; text-align: center;}
#analyze-btn { background: #0061ff; color: white; border: none; font-weight: bold; cursor: pointer; border-radius: 8px;}

/* Custom Scroll Wheel Styling */
#scroll-area textarea {
    overflow-y: scroll !important;
    scrollbar-width: thin;
    scrollbar-color: #0061ff #2d2d2d;
}
#scroll-area textarea::-webkit-scrollbar {
    width: 8px;
}
#scroll-area textarea::-webkit-scrollbar-track {
    background: #1a1a1a;
}
#scroll-area textarea::-webkit-scrollbar-thumb {
    background-color: #0061ff;
    border-radius: 10px;
}
"""

with gr.Blocks(title="FlumenIntel") as demo:
    gr.HTML(f"<style>{custom_css}</style>")
    with gr.Column(elem_id="title-box"):
        gr.Markdown("# FlumenIntel 🌊\n### Advanced River Health Analytics")

    with gr.Tabs():
        with gr.TabItem("🚀 Dashboard"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### 1. Hydrological Data")
                    flow = gr.Number(label="Flow Rate", value=45)
                    temp = gr.Number(label="Temperature", value=18)
                    sediment = gr.Slider(0, 10, label="Sediment", value=2)
                    construction = gr.Slider(0, 10, label="Construction", value=0)
                    
                    gr.Markdown("### 2. Chemical Data")
                    ph = gr.Number(label="pH Level", value=7.2)
                    do = gr.Number(label="Dissolved Oxygen", value=9.5)
                    nutrients = gr.Slider(0, 10, label="Nutrient Load", value=1)
                    
                    gr.Markdown("### 3. Visual Analysis")
                    sat_img = gr.Image(label="Satellite Image", type="pil")
                    
                    analyze_btn = gr.Button("GENERATE REPORT", elem_id="analyze-btn")

                with gr.Column(scale=2):
                    status_box = gr.Textbox(label="System Status", interactive=False)
                    
                    with gr.Tabs():
                        with gr.TabItem("📊 Visual Analytics"):
                            plot_output = gr.Plot(label="Metric Visualization")
                            graph_summary_box = gr.Markdown("### Insights...")
                        
                        with gr.TabItem("📄 Official Report"):
                            ai_summary = gr.Textbox(
                                label="Scientist's Assessment", 
                                lines=15, 
                                interactive=False, 
                                elem_id="scroll-area" # Added specifically for the CSS
                            )
                            
                            with gr.Row():
                                audio_btn = gr.Button("🔊 Listen to Report (gTTS)")
                                audio_out = gr.Audio(label="Player", type="filepath")
                            
                            audio_btn.click(
                                fn=generate_audio_report, 
                                inputs=ai_summary, 
                                outputs=audio_out
                            )

                        with gr.TabItem("📥 Export"):
                            pdf_output = gr.File(label="Download Official FlumenIntel Report")

        with gr.TabItem("👤 About Me"):
             gr.Markdown("## Abdullah\nComputer Engineering Undergraduate | AI & Hardware Enthusiast")

    analyze_btn.click(
        run_app,
        inputs=[flow, temp, sediment, construction, ph, do, nutrients, sat_img],
        outputs=[status_box, plot_output, graph_summary_box, ai_summary, pdf_output]
    )

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