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

from flask import (
    Flask, render_template, request, jsonify, session,
    make_response, redirect, url_for
)
from werkzeug.middleware.proxy_fix import ProxyFix
from flask_session import Session
from cachelib import SimpleCache
from dotenv import load_dotenv

load_dotenv()

# Logging config
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# API Keys
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or os.getenv("GEMINI_API")
NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")

# Initialize Gemini if available
if GEMINI_API_KEY:
    import google.generativeai as genai
    genai.configure(api_key=GEMINI_API_KEY)

# Initialize NVIDIA client if available
nvidia_client = None
if NVIDIA_API_KEY:
    try:
        from openai import OpenAI
        nvidia_client = OpenAI(
            base_url="https://integrate.api.nvidia.com/v1",
            api_key=NVIDIA_API_KEY
        )
    except ImportError:
        logger.warning("OpenAI library not installed for NVIDIA fallback")

app = Flask(__name__)
app.secret_key = os.getenv("SECRET_KEY", "pestipedia-secret-key-2024")
app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1)

app.config["SESSION_TYPE"] = "cachelib"
app.config["SESSION_CACHELIB"] = SimpleCache()
app.config["SESSION_PERMANENT"] = True
app.config["PERMANENT_SESSION_LIFETIME"] = timedelta(days=30)

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["SESSION_COOKIE_NAME"] = os.getenv("SESSION_COOKIE_NAME", "hf_app_session")

Session(app)

# Simple in-memory cache
cache = SimpleCache()

# Supported languages
SUPPORTED_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": "Brown Plant Hopper", "type": "pest", "crop": "Rice", "image_url": "/static/images/brown_planthopper.jpg"},
    {"id": 3, "name": "Rice Blast", "type": "disease", "crop": "Rice", "image_url": "/static/images/rice_blast.jpg"},
    {"id": 4, "name": "Bacterial Leaf Blight", "type": "disease", "crop": "Rice", "image_url": "/static/images/bacterial_leaf_blight.jpg"},
    {"id": 5, "name": "Pink Bollworm", "type": "pest", "crop": "Cotton", "image_url": "/static/images/stem_borer.jpg"},
    {"id": 6, "name": "Red Cotton Bug", "type": "pest", "crop": "Cotton", "image_url": "/static/images/jassids.jpg"},
    {"id": 7, "name": "Powdery Mildew", "type": "disease", "crop": "Wheat, Grapes", "image_url": "/static/images/powdery_mildew.jpg"},
    {"id": 8, "name": "Tomato Fruit Borer", "type": "pest", "crop": "Tomato, Cotton", "image_url": "/static/images/aphids.jpg"},
    {"id": 9, "name": "Late Blight", "type": "disease", "crop": "Potato, Tomato", "image_url": "/static/images/yellow_mosaic.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"}
]


def extract_json_from_response(content: str):
    """Extract JSON from model response."""
    content = content.strip()
    json_match = re.search(r'\{[\s\S]*\}', content)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError:
            pass
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        return {}


def call_llm(prompt):
    """Call LLM with Google (Gemini + Gemma) -> NVIDIA (Llama) fallback."""
    
    # 1. Google API: Gemini + Gemma models
    if GEMINI_API_KEY:
        google_models = [
            "gemini-2.5-flash",
            "gemini-2.5-flash-lite",
            "gemma-3-27b-it",
            "gemma-3-12b-it",
            "gemma-3-4b-it",
            "gemma-3-1b-it"
        ]
        for model_name in google_models:
            try:
                logger.info(f"  >> Trying Google {model_name}...")
                model = genai.GenerativeModel(model_name)
                response = model.generate_content(prompt)
                if response and response.text:
                    logger.info(f"  >> SUCCESS with {model_name}")
                    return response.text
            except Exception as e:
                logger.warning(f"  >> {model_name} failed: {e}")
                continue
    
    # 2. NVIDIA API: Llama models only
    if nvidia_client:
        nvidia_models = ["meta/llama-3.1-405b-instruct", "meta/llama-3.1-70b-instruct"]
        for model_name in nvidia_models:
            try:
                logger.info(f"  >> Trying NVIDIA {model_name}...")
                response = nvidia_client.chat.completions.create(
                    model=model_name,
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=4096,
                    temperature=0.3
                )
                logger.info(f"  >> SUCCESS with {model_name}")
                return response.choices[0].message.content
            except Exception as e:
                logger.warning(f"  >> {model_name} failed: {e}")
                continue
    
    return None


@app.route("/")
def index():
    language = session.get("language", "en")
    return render_template(
        "index.html",
        pests_diseases=PESTS_DISEASES,
        languages=SUPPORTED_LANGUAGES,
        current_language=language
    )


@app.route("/set_language", methods=["POST"])
def set_language():
    lang = request.form.get("language", "en")
    if lang in SUPPORTED_LANGUAGES:
        session["language"] = lang
        session.permanent = True
    resp = make_response(redirect(url_for("index")))
    resp.set_cookie("language", lang, max_age=60*60*24*365, samesite="None", secure=not disable_secure)
    return resp


@app.route("/get_details/<int:pest_id>")
def get_details(pest_id):
    lang_param = request.args.get("lang")
    language = lang_param or session.get("language") or request.cookies.get("language") or "en"
    session["language"] = language
    session.permanent = True

    pest = next((p for p in PESTS_DISEASES if p["id"] == pest_id), None)
    if not pest:
        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 pest %s (%s)", pest_id, language)
        return jsonify(cached)

    # Build prompt with EXPLICIT JSON EXAMPLE
    lang_instructions = {
        "en": "Respond in English",
        "hi": "हिंदी में जवाब दें (Respond in Hindi)",
    }
    
    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, Lifecycle, Symptoms, Economic Impact, Management (Organic and Chemical), Prevention.
    
    YOU MUST FORMAT YOUR RESPONSE AS A SINGLE, VALID JSON OBJECT EXACTLY LIKE THIS EXAMPLE:
    {{
        "description": {{"title": "Description", "text": "..."}},
        "lifecycle": {{"title": "Lifecycle", "text": "..."}},
        "symptoms": {{"title": "Symptoms", "text": "..."}},
        "impact": {{"title": "Economic Impact", "text": "..."}},
        "management": {{
            "organic": {{"title": "Organic Control", "text": "Detailed organic and biological control methods..."}},
            "chemical": {{"title": "Chemical Control", "text": "Recommended insecticides and fungicides..."}}
        }},
        "prevention": {{"title": "Prevention", "text": "..."}}
    }}
    
    IMPORTANT: The "management" key MUST have TWO sub-keys: "organic" and "chemical", each with "title" and "text".
    Return ONLY the JSON object, no markdown formatting, no code blocks.
    """

    logger.info(f"[Pestipedia] Fetching details for {pest['name']} in {language}")
    
    response_text = call_llm(prompt)
    
    # --- ROBUST FALLBACK (MOCK DATA) ---
    if not response_text:
        logger.warning("All LLM models failed; returning High-Quality Mock")
        
        mock_details = {
             "description": {
                 "title": "Scientific Diagnosis", 
                 "text": f"The {pest['name']} is a significant agricultural threat in Indian subcontinents. It is characterized by specific patterns on the leaves/fruit. Early detection is critical."
             },
             "lifecycle": {
                 "title": "Life Cycle Analysis", 
                 "text": "The lifecycle typically spans 30-45 days, moving from egg -> larva -> pupa -> adult. Rapid reproduction occurs in humid conditions (25-30°C)."
             },
             "symptoms": {
                 "title": "Clinical Symptoms", 
                 "text": "Look for: 1. Yellowing of leaves. 2. Stunted growth. 3. Visible lesions on the stem or fruit. 4. Reduced yield quality."
             },
             "impact": {
                 "title": "Economic Forecast", 
                 "text": "Can cause yield losses of up to 40-70% if untreated. Reduces market value significantly due to cosmetic damage."
             },
             "management": {
                 "organic": {
                     "title": "Organic Control",
                     "text": "**Biological Control:** Use Neem Oil (5ml/L), Pheromone traps (10/acre), and introduce natural predators like Trichogramma. **Cultural Practices:** Crop rotation, deep summer ploughing, and field sanitation."
                 },
                 "chemical": {
                     "title": "Chemical Control",
                     "text": "**Recommended Insecticides:** Apply Chlorpyrifos 20EC (2ml/L) or Imidacloprid 17.8SL (0.3ml/L) during early infestation. **Application:** Spray in the early morning or late evening for best results. Always follow label instructions."
                 }
             },
             "prevention": {
                 "title": "Preventive Strategy", 
                 "text": "1. Crop rotation. 2. Use resistant varieties. 3. Deep summer ploughing to destroy pupae. 4. Maintain field sanitation."
             }
        }
        
        result = {
            **pest,
            "details": mock_details,
            "language": language,
            "cached_at": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC")
        }
        return jsonify(result)

    detailed_info = extract_json_from_response(response_text)

    # Validate and fix keys (including nested management)
    required_keys = ["description", "lifecycle", "symptoms", "impact", "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": "Details unavailable at this moment."}
    
    # Fix management structure if AI didn't follow instructions
    if "management" not in detailed_info:
        detailed_info["management"] = {}
    
    mgmt = detailed_info["management"]
    
    # If AI returned flat text, convert it
    if "text" in mgmt and "organic" not in mgmt:
        old_text = mgmt.get("text", "")
        detailed_info["management"] = {
            "organic": {"title": "Organic Control", "text": old_text if old_text else "No organic control data available."},
            "chemical": {"title": "Chemical Control", "text": "No chemical control data available."}
        }
    else:
        # Ensure organic and chemical keys exist
        if "organic" not in mgmt or "text" not in mgmt.get("organic", {}):
            mgmt["organic"] = {"title": "Organic Control", "text": "No organic control data available."}
        if "chemical" not in mgmt or "text" not in mgmt.get("chemical", {}):
            mgmt["chemical"] = {"title": "Chemical Control", "text": "No chemical control data available."}

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


if __name__ == "__main__":
    # Changed PORT to 5006 to avoid conflicts
    port = int(os.environ.get("PORT", 5006))
    host = os.environ.get("HOST", "0.0.0.0")
    app.run(host=host, port=port, debug=True)