import os
import json
from flask import (
    Flask, request, jsonify, render_template,
    Response, stream_with_context, send_file
)
from flask_cors import CORS
from dotenv import load_dotenv

# -----------------------------
# BASIC SETUP
# -----------------------------
load_dotenv()

app = Flask(__name__, static_folder="static", template_folder="templates")
app.secret_key = "supersecretkey"
CORS(app)

# -----------------------------
# SAFE SERVER IMPORTS
# -----------------------------
import pandas as pd
from geopy.distance import geodesic
from gtts import gTTS
from groq import Groq

# -----------------------------
# ENV / API KEYS
# -----------------------------
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    raise RuntimeError("GROQ_API_KEY not set in environment")

client = Groq(api_key=GROQ_API_KEY)

# -----------------------------
# CONFIG
# -----------------------------
CHAT_MODEL = "llama3-8b-8192"

SYSTEM_PROMPT = (
    "You are Kriti, an intelligent and helpful AI assistant created by ISRO. "
    "Respond in the same language as the user. "
    "Keep answers short, clear, and natural (1–2 sentences). "
    "Never mention internal model names."
)

# -----------------------------
# DATA LOAD (SAFE)
# -----------------------------
reservoirs = pd.read_csv("ISROP.csv")

# -----------------------------
# ROUTES (PAGES)
# -----------------------------
@app.route("/")
def home():
    return render_template("index.html")


@app.route("/search")
def search():
    return render_template("index.html", query=request.args.get("query", ""))


@app.route("/reservoir")
def reservoir():
    return render_template("reservoir.html")


@app.route("/climatemap")
def climatemap():
    return render_template("climatemap.html")

# -----------------------------
# REAL-TIME CHAT (STREAMING)
# -----------------------------
@app.route("/chat_stream", methods=["POST"])
def chat_stream():
    data = request.get_json()
    user_input = data.get("message", "").strip()

    if not user_input:
        return jsonify({"error": "Empty message"}), 400

    def generate():
        try:
            completion = client.chat.completions.create(
                model=CHAT_MODEL,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": user_input}
                ],
                stream=True
            )

            for chunk in completion:
                if chunk.choices:
                    delta = chunk.choices[0].delta
                    token = delta.get("content", "")
                    if token:
                        yield f"data: {token}\n\n"

            yield "data: [DONE]\n\n"

        except Exception as e:
            yield f"data: [ERROR] {str(e)}\n\n"

    return Response(
        stream_with_context(generate()),
        mimetype="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no"
        }
    )

# -----------------------------
# NON-STREAM CHAT (FALLBACK)
# -----------------------------
@app.route("/chat", methods=["POST"])
def chat():
    user_input = request.json.get("message", "")

    try:
        completion = client.chat.completions.create(
            model=CHAT_MODEL,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_input}
            ]
        )
        response_text = completion.choices[0].message.content
        return jsonify({"response": response_text})

    except Exception as e:
        return jsonify({"error": str(e)}), 500

# -----------------------------
# TEXT TO SPEECH
# -----------------------------
@app.route("/speak", methods=["POST"])
def speak():
    data = request.json
    text = data.get("text", "")
    lang = data.get("language", "en")

    if not text:
        return "No text provided", 400

    audio_path = "response.mp3"
    gTTS(text=text, lang=lang).save(audio_path)
    return send_file(audio_path, mimetype="audio/mp3")

# -----------------------------
# NEAREST RESERVOIR
# -----------------------------
@app.route("/find_nearest_reservoir", methods=["POST"])
def find_nearest_reservoir():
    req = request.get_json()
    lat, lon = float(req["lat"]), float(req["lon"])

    nearest = None
    min_dist = float("inf")

    for _, row in reservoirs.iterrows():
        dist = geodesic(
            (lat, lon),
            (row["Latitude"], row["Longitude"])
        ).km

        if dist < min_dist:
            min_dist = dist
            nearest = row

    if nearest is None:
        return jsonify({"error": "No reservoir found"})

    return jsonify({
        "name": nearest["ROI_Name"],
        "state": nearest["State"],
        "latitude": nearest["Latitude"],
        "longitude": nearest["Longitude"],
        "distance": round(min_dist, 2)
    })

# -----------------------------
# ENTRY POINT (HF REQUIRED)
# -----------------------------
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)