from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
from PIL import Image
import base64
import io
import requests

app = FastAPI(title="STOA Plant Disease API")

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

# --- MODEL LOADING ---
print("Booting Agricultural Node. Loading MobileNetV2 Plant model...")
# THE FIX: Explicitly borrow the Google MobileNetV2 image processor
pipe = pipeline(
    "image-classification", 
    model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification",
    image_processor="google/mobilenet_v2_1.0_224"
)
print("Agent Ready!")

# --- REQUEST SCHEMA ---
class PredictRequest(BaseModel):
    image: str | None = None
    image_url: str | None = None

# --- ENDPOINTS ---
@app.get("/health")
def health_check():
    return {"status": "ok"}

@app.post("/predict")
def predict(req: PredictRequest):
    try:
        img = None
        
        # 1. Handle URL Input (with Super-Human headers)
        if req.image_url:
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
                "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
                "Referer": "https://google.com"
            }
            response = requests.get(req.image_url, stream=True, headers=headers, timeout=10)
            if response.status_code != 200:
                raise Exception(f"External site blocked us with error: {response.status_code}.")
            img = Image.open(response.raw).convert("RGB")
            
        # 2. Handle Base64 Input
        elif req.image:
            b64_data = req.image
            if "," in b64_data:
                b64_data = b64_data.split(",")[1]
            image_bytes = base64.b64decode(b64_data)
            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            
        else:
            raise HTTPException(status_code=400, detail="Must provide 'image' (base64) or 'image_url'.")

        # 3. Execute AI Math
        results = pipe(img, top_k=3)
        
        # 4. Format Output for the STOA Marketplace
        top_3_list = [{"disease": res["label"], "confidence": round(res["score"], 4)} for res in results]
        
        return {
            "prediction": results[0]["label"],
            "confidence": round(results[0]["score"], 4),
            "top_3": top_3_list
        }
        
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to process leaf: {str(e)}")