import os
import json
import re
import logging
from datetime import datetime, timedelta

from flask import (
    Flask, render_template, request, jsonify, session,
    make_response, redirect, url_for
)
from werkzeug.middleware.proxy_fix import ProxyFix
import google.generativeai as genai
from dotenv import load_dotenv
from cachelib import SimpleCache

# Load .env
load_dotenv()

# Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("app")

# Flask app
app = Flask(__name__, static_folder="static", template_folder="templates")

# Secret key
app.secret_key = os.getenv("FLASK_SECRET_KEY", os.urandom(24))

# ProxyFix - trust one proxy layer (adjust numbers if you sit behind more proxies)
app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_port=1, x_prefix=1)

# Cookie / session policy (sane defaults for Hugging Face Spaces)
# By default we set SECURE cookies and SameSite=None which is required for cross-site in modern browsers.
# If you are running locally over HTTP, set DISABLE_SECURE_COOKIE=1 in your environment (for development only).
disable_secure = os.getenv("DISABLE_SECURE_COOKIE", "0") in ("1", "true", "True")
app.config["SESSION_COOKIE_SAMESITE"] = "None"
app.config["SESSION_COOKIE_SECURE"] = False if disable_secure else True
app.config["PERMANENT_SESSION_LIFETIME"] = timedelta(days=7)
app.config["SESSION_PERMANENT"] = True
app.config["SESSION_COOKIE_NAME"] = os.getenv("SESSION_COOKIE_NAME", "hf_app_session")

# Simple in-memory cache (for details from the Gemini API)
cache = SimpleCache(threshold=200, default_timeout=60 * 60 * 24 * 7)  # 7 days default timeout

# Configure Gemini (Gemini API key must be set in environment as GEMINI_API_KEY)
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if GEMINI_API_KEY:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
    except Exception as e:
        logger.exception("Failed to configure Gemini client: %s", e)
else:
    logger.warning("GEMINI_API_KEY not provided. API calls will not work until you set this env var.")

# Supported languages and list of pests/diseases (kept from your original list)
LANGUAGES = {
    "en": "English", "hi": "हिंदी (Hindi)", "bn": "বাংলা (Bengali)", "te": "తెలుగు (Telugu)",
    "mr": "मराठी (Marathi)", "ta": "தமிழ் (Tamil)", "gu": "ગુજરાતી (Gujarati)", "ur": "اردو (Urdu)",
    "kn": "ಕನ್ನಡ (Kannada)", "or": "ଓଡ଼ିଆ (Odia)", "ml": "മലയാളം (Malayalam)"
}

PESTS_DISEASES = [
    {"id": 1, "name": "Fall Armyworm", "type": "pest", "crop": "Maize, Sorghum", "image_url": "/static/images/fall_armyworm.jpg"},
    {"id": 2, "name": "Rice Blast", "type": "disease", "crop": "Rice", "image_url": "/static/images/rice_blast.jpg"},
    {"id": 3, "name": "Aphids", "type": "pest", "crop": "Various crops", "image_url": "/static/images/aphids.jpg"},
    {"id": 4, "name": "Powdery Mildew", "type": "disease", "crop": "Wheat, Vegetables", "image_url": "/static/images/powdery_mildew.jpg"},
    {"id": 5, "name": "Stem Borer", "type": "pest", "crop": "Rice, Sugarcane", "image_url": "/static/images/stem_borer.jpg"},
    {"id": 6, "name": "Yellow Mosaic Virus", "type": "disease", "crop": "Pulses", "image_url": "/static/images/yellow_mosaic.jpg"},
    {"id": 7, "name": "Brown Planthopper", "type": "pest", "crop": "Rice", "image_url": "/static/images/brown_planthopper.jpg"},
    {"id": 8, "name": "Bacterial Leaf Blight", "type": "disease", "crop": "Rice", "image_url": "/static/images/bacterial_leaf_blight.jpg"},
    {"id": 9, "name": "Jassids", "type": "pest", "crop": "Cotton, Maize", "image_url": "/static/images/jassids.jpg"},
    {"id": 10, "name": "Downy Mildew", "type": "disease", "crop": "Grapes, Onion", "image_url": "/static/images/downy_mildew.jpg"},
    {"id": 11, "name": "Whitefly", "type": "pest", "crop": "Tomato, Cotton", "image_url": "/static/images/whitefly.jpg"},
    {"id": 12, "name": "Fusarium Wilt", "type": "disease", "crop": "Banana, Tomato", "image_url": "/static/images/fusarium_wilt.jpg"}
]

# Utility: robust JSON extraction from potentially noisy model output
def extract_json_from_response(content: str):
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        # Attempt to extract codeblock containing JSON
        json_match = re.search(r'```json(.*?)```', content, re.DOTALL)
        if json_match:
            json_str = json_match.group(1).strip()
        else:
            # Fallback: try to extract first {...} block
            brace_match = re.search(r'(\{(?:.|\n)*\})', content)
            json_str = brace_match.group(1) if brace_match else content

        # Remove stray markdown markers and weird trailing commas
        json_str = json_str.replace('```json', '').replace('```', '').strip()
        json_str = re.sub(r',(\s*[\]\}])', r'\1', json_str)

        try:
            return json.loads(json_str)
        except json.JSONDecodeError as e:
            logger.error("JSON parsing error after cleanup: %s\nOriginal content: %s", e, content[:1000])
            raise ValueError("Failed to parse JSON response from API")

# Root page
@app.route("/")
def index():
    # Determine language: prefer session, fallback to cookie, then default 'en'.
    language = session.get("language") or request.cookies.get("language") or "en"
    # Keep session in sync with cookie if needed
    session["language"] = language
    session.permanent = True
    return render_template("index.html",
                           pests_diseases=PESTS_DISEASES,
                           languages=LANGUAGES,
                           current_language=language)

# Language setter (Post-Redirect-Get)
@app.route("/set_language", methods=["POST"])
def set_language():
    language = request.form.get("language", "en")
    if language not in LANGUAGES:
        logger.warning("Attempt to set unsupported language: %s", language)
        language = "en"

    # Save in server-side session
    session["language"] = language
    session.permanent = True
    logger.info("Language successfully changed to %s in session.", language)

    # Also set a client cookie as reliable fallback for stateless platforms
    resp = make_response(redirect(url_for("index")))
    # cookie lifetime 7 days
    max_age = 60 * 60 * 24 * 7
    secure_flag = False if disable_secure else True
    resp.set_cookie("language", language, max_age=max_age, secure=secure_flag, samesite="None", httponly=False)
    return resp

# Endpoint to fetch details for a pest/disease (uses cache + Gemini)
@app.route("/get_details/<int:pest_id>")
def get_details(pest_id: int):
    # Resolve language (session first, cookie fallback)
    language = session.get("language") or request.cookies.get("language") or "en"

    pest = next((p for p in PESTS_DISEASES if p["id"] == pest_id), None)
    if pest is None:
        return jsonify({"error": "Not found"}), 404

    cache_key = f"pest_{pest_id}_{language}"
    cached = cache.get(cache_key)
    if cached:
        logger.info("Cache hit for %s (%s)", pest_id, language)
        return jsonify(cached)

    if not GEMINI_API_KEY:
        logger.warning("Gemini API key missing; returning placeholder result.")
        result = {
            **pest,
            "details": {
                "description": {"title": "Description", "text": f"Placeholder description for {pest['name']} ({language})."},
                "lifecycle": {"title": "Lifecycle", "text": "N/A"},
                "symptoms": {"title": "Symptoms", "text": "N/A"},
                "impact": {"title": "Economic Impact", "text": "N/A"},
                "management": {"title": "Management", "text": "N/A"},
                "prevention": {"title": "Prevention", "text": "N/A"},
            },
            "language": language,
            "cached_at": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
        }
        cache.set(cache_key, result)
        return jsonify(result)

    # Build prompt with language-specific instruction
    lang_instructions = {
        "en": "Respond in English",
        "hi": "हिंदी में जवाब दें (Respond in Hindi)",
        "bn": "বাংলায় উত্তর দিন (Respond in Bengali)",
        "te": "తెలుగులో సమాధానం ఇవ్వండి (Respond in Telugu)",
        "mr": "मराठीत उत्तर द्या (Respond in Marathi)",
        "ta": "தமிழில் பதிலளிக்கவும் (Respond in Tamil)",
        "gu": "ગુજરાતીમાં જવાબ આપો (Respond in Gujarati)",
        "ur": "اردو میں جواب دیں (Respond in Urdu)",
        "kn": "ಕನ್ನಡದಲ್ಲಿ ಉತ್ತರಿಸಿ (Respond in Kannada)",
        "or": "ଓଡ଼ିଆରେ ଉତ୍ତର ଦିଅନ୍ତୁ (Respond in Odia)",
        "ml": "മലയാളത്തിൽ മറുപടി നൽകുക (Respond in Malayalam)"
    }
    prompt = f"""
{lang_instructions.get(language, 'Respond in English')}
Provide detailed information about the agricultural {pest['type']} '{pest['name']}' affecting '{pest['crop']}' in India.
Include these sections: Description and identification, Lifecycle and spread, Damage symptoms, Economic impact, Management and control (organic and chemical), Preventive measures.
Format the response as a single, strictly valid JSON object with keys: "description", "lifecycle", "symptoms", "impact", "management", "prevention".
Each key must contain an object with "title" and "text" fields.
"""

    try:
        model = genai.GenerativeModel("gemini-1.5-flash")
        response = model.generate_content(prompt)
        text = getattr(response, "text", "") or str(response)
        detailed_info = extract_json_from_response(text)

        # Ensure keys exist
        required_keys = ["description", "lifecycle", "symptoms", "impact", "management", "prevention"]
        for key in required_keys:
            if key not in detailed_info or "text" not in detailed_info.get(key, {}):
                detailed_info[key] = {"title": key.capitalize(), "text": "Information could not be generated for this section."}
                logger.warning("API response missing or malformed for key: %s", key)

        result = {
            **pest,
            "details": detailed_info,
            "language": language,
            "cached_at": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
        }
        cache.set(cache_key, result)
        logger.info("Cached details for pest %s (%s)", pest_id, language)
        return jsonify(result)

    except Exception as e:
        logger.exception("Error fetching details from Gemini API: %s", e)
        return jsonify({"error": "Failed to fetch information", "message": str(e)}), 500


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    host = os.environ.get("HOST", "0.0.0.0")
    app.run(host=host, port=port)