from fastapi import FastAPI, File, UploadFile, Form
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional
from ultralytics import YOLO
import uvicorn
import shutil
import os

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], 
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ==========================================
#  LOAD MODELS (Dual-Architecture)
# ==========================================
print("⏳ Loading Models...")

try:
    disease_model = YOLO('best.pt')
    print("✅ Goyam Disease Classifier loaded!")
except Exception as e:
    print(f"❌ Error loading Disease Model: {e}")

try:
    gatekeeper_model = YOLO('yolov8n-cls.pt') 
    print("✅ Plant Gatekeeper loaded!")
except Exception as e:
    print(f"❌ Error loading Gatekeeper: {e}")


# ==========================================
# 🛡️ THE SMART GATEKEEPER 
# ==========================================
def is_likely_plant(image_path):
    try:
        results = gatekeeper_model(image_path, verbose=False)
        top5_indices = results[0].probs.top5
        top5_names = [results[0].names[i].lower() for i in top5_indices]
        
        print(f"Gatekeeper sees: {top5_names}")

        #  ALLOW LIST: Botanical terms AND known macro-photography hallucinations
        allowed_keywords = [
            # True Botanical
            'plant', 'leaf', 'grass', 'flower', 'tree', 'fern', 'moss', 'weed', 
            'crop', 'agriculture', 'field', 'greenhouse', 'pot', 'earth', 'soil', 
            'vegetation', 'forest', 'valley', 'daisy', 'corn', 'acorn', 'paddy', 'cardoon', 'reed',
            # Common ImageNet Hallucinations for extreme close-up leaf textures
            'damselfly', 'chameleon', 'lizard', 'ear', 'lacewing', 'spider', 'insect', 
            'bug', 'mantis', 'paintbrush', 'broom', 'bow', 'nematode', 'slug', 'snail', 'snake'
        ]

        for predicted_item in top5_names:
            for good_word in allowed_keywords:
                if good_word in predicted_item:
                    print(f"✅ Passed: Gatekeeper authorized based on '{predicted_item}'")
                    return True

        print(f" Blocked: No natural or agricultural features detected.")
        return False

    except Exception as e:
        print(f" Gatekeeper Error: {e}")
        return True 

def get_recommendation(disease_name):
    recommendations = {
        "Leaf Blast": "Use Tricyclazole 75 WP. Avoid applying excess nitrogen fertilizer.",
        "Sheath Blight": "Drain the field immediately. Apply validamycin or carbendazim.",
        "Brown Spot": "Improve soil fertility. Apply potassium and phosphorus.",
        "Healthy Rice Leaf": "No disease detected. Keep maintaining optimal water levels!"
    }
    for key, value in recommendations.items():
        if key.lower() in disease_name.lower():
            return value
    return "Consult your local agricultural extension officer for treatment."

@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    latitude: Optional[str] = Form(None),
    longitude: Optional[str] = Form(None)
):
    print(f" Receiving image: {file.filename}")
    
    temp_filename = f"temp_{file.filename}"
    try:
        with open(temp_filename, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        # 🛑 STEP 1: RUN THE SMART GATEKEEPER
        if not is_likely_plant(temp_filename):
            return {
                "filename": file.filename,
                "disease": "Invalid Image",
                "confidence": "0%",
                "recommendation": "This image does not appear to be a plant or paddy field. Please upload a clear photo of a rice leaf.",
                "latitude": float(latitude) if latitude else None,
                "longitude": float(longitude) if longitude else None
            }

        # STEP 2: RUN DISEASE CLASSIFICATION
        results = disease_model(temp_filename, verbose=False)
        
        top_idx = results[0].probs.top1
        confidence_score = float(results[0].probs.top1conf)
        detected_name = results[0].names[top_idx]
        
        response_data = {
            "filename": file.filename,
            "disease": detected_name,
            "confidence": f"{int(confidence_score * 100)}%",
            "recommendation": get_recommendation(detected_name),
            "latitude": float(latitude) if latitude else None,
            "longitude": float(longitude) if longitude else None
        }

        return response_data

    except Exception as e:
        print(f" API Error: {e}")
        return {"error": str(e)}
        
    finally:
        if os.path.exists(temp_filename):
            os.remove(temp_filename)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)