"""
Crop Classifier REST API  v3.1
================================
Status: Stability & Debug Update
Changes: Supabase resilience, FileResponse, detailed JSON errors.
"""

from fastapi import FastAPI, File, UploadFile, HTTPException, Header, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from pydantic import BaseModel, EmailStr
import numpy as np
import json, os, io, time, logging, base64, uuid, secrets
from datetime import datetime, timezone
from PIL import Image
import tensorflow as tf
from tensorflow.keras.applications.efficientnet import preprocess_input
import requests as req_lib
import httpx

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# ── Supabase config ─────────────────────────────────────────────────────────
SUPABASE_URL = os.environ.get("SUPABASE_URL", "https://ykvatttsnpjrwqfhhysu.supabase.co")
SUPABASE_KEY = os.environ.get("SUPABASE_KEY", 
    "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InlrdmF0dHRzbnBqcndxZmhoeXN1Iiwicm9sZSI6ImFub24iLCJpYXQiOjE3NzA0OTk5NjQsImV4cCI6MjA4NjA3NTk2NH0.5Njnh8NBEcPDddHjwv3CoUpCcAHu-ALNUQHQVdAdq-Y"
)
SB_HEADERS = {"apikey": SUPABASE_KEY, "Authorization": f"Bearer {SUPABASE_KEY}", "Content-Type": "application/json"}
SB_TABLE   = f"{SUPABASE_URL}/rest/v1/crop_api_keys"

# ── In-memory key cache ──────────────────────────────────────────────────────
_key_cache: dict = {}        # {api_key: {name, email, id}}
_cache_ts: float = 0.0

async def refresh_key_cache():
    global _key_cache, _cache_ts
    try:
        async with httpx.AsyncClient() as client:
            r = await client.get(f"{SB_TABLE}?is_active=eq.true&select=api_key,name,email,id", 
                                 headers=SB_HEADERS, timeout=10)
            if r.status_code == 200:
                _key_cache = {row["api_key"]: row for row in r.json()}
                _cache_ts  = time.time()
                logger.info(f"Key cache refreshed: {len(_key_cache)} keys")
            else:
                logger.warning(f"Key cache failed: Supabase returned {r.status_code}")
    except Exception as e:
        logger.error(f"Critical error refreshing key cache: {e}")

# ── App ──────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="🌾 Crop Classifier API",
    version="3.1.0",
)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, 
                   allow_methods=["*"], allow_headers=["*"])

# ── Startup ──────────────────────────────────────────────────────────────────
MODEL_PATH = os.environ.get("MODEL_PATH", "best_v6.keras")
JSON_PATH  = os.environ.get("JSON_PATH",  "class_names.json")

model = None
class_names = []

@app.on_event("startup")
async def startup():
    global model, class_names
    try:
        logger.info(f"Loading model: {MODEL_PATH}")
        model = tf.keras.models.load_model(MODEL_PATH)
        with open(JSON_PATH) as f:
            class_names = json.load(f)["class_names"]
        logger.info(f"Model loaded successfully. {len(class_names)} classes.")
        await refresh_key_cache()
    except Exception as e:
        logger.error(f"STARTUP FAILED: {e}")
        # Dont raise here - allow server to start so /health works for debugging

# ── LLaMA ────────────────────────────────────────────────────────────────────
NVIDIA_API_KEY = os.environ.get("NVIDIA_API_KEY", "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")
LLAMA_URL = "https://integrate.api.nvidia.com/v1/chat/completions"

# ── Helpers ──────────────────────────────────────────────────────────────────
def confidence_label(pct: float) -> str:
    return "High" if pct >= 70 else "Medium" if pct >= 40 else "Low"

def preprocess_image(image_bytes: bytes) -> np.ndarray:
    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        image = image.resize((224, 224))
        arr = np.expand_dims(np.array(image, dtype=np.float32), axis=0)
        return preprocess_input(arr)
    except Exception as e:
        logger.warning(f"Image preprocessing failed: {e}")
        return None

def compress_image(image_bytes: bytes) -> bytes:
    try:
        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        img.thumbnail((768, 768))
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=85, optimize=True)
        return buf.getvalue()
    except Exception:
        return image_bytes

def call_llama_vision(image_bytes: bytes, top3_preds: list) -> dict:
    try:
        img_b64 = base64.b64encode(compress_image(image_bytes)).decode("utf-8")
        predictions_str = ", ".join(f"{p['crop']} ({p['confidence_percent']}%)" for p in top3_preds)
        prompt = (
            "You are an expert agricultural scientist. Analyze the crop in this image.\n"
            f"Possible crops: {predictions_str}\n\n"
            "Format your response as EXACTLY these fields, one per line:\n"
            "**Crop Name:** [Correct common name]\n"
            "**Scientific Name:** [Latin name]\n"
            "**Characteristics:** [Visual features]\n"
            "**Quality:** [Premium, Excellent, Very Good, Good, Fair, or Bad]\n"
            "**Market Grade:** [Grade A, Grade B, Grade C, or Ungraded]\n"
            "**Prediction Accuracy:** [Correct, Partially Correct, or Incorrect]\n"
            "**Storage Tip:** [1 sentence]\n"
            "**Explanation:** [2 sentences]"
        )
        payload = {
            "model": "meta/llama-3.2-90b-vision-instruct",
            "messages": [{"role": "user", "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}}
            ]}],
            "max_tokens": 600, "temperature": 0.2, "stream": False
        }
        headers = {"Authorization": f"Bearer {NVIDIA_API_KEY}", "Accept": "application/json"}
        resp = req_lib.post(LLAMA_URL, headers=headers, json=payload, timeout=60)
        resp.raise_for_status()
        raw_text = resp.json()["choices"][0]["message"]["content"]

        fields = {"crop_name": None, "scientific_name": None, "characteristics": None,
                  "quality": None, "market_grade": None, "prediction_accuracy": None,
                  "storage_tip": None, "explanation": None}
        for line in raw_text.splitlines():
            clean = line.replace("**", "").replace("- ","").replace("* ","").strip()
            if ":" in clean:
                k, v = clean.split(":", 1)
                k = k.lower().strip()
                if "crop name" in k: fields["crop_name"] = v.strip()
                elif "scientific name" in k: fields["scientific_name"] = v.strip()
                elif "characteristics" in k: fields["characteristics"] = v.strip()
                elif "quality" in k: fields["quality"] = v.strip()
                elif "market grade" in k: fields["market_grade"] = v.strip()
                elif "prediction accuracy" in k: fields["prediction_accuracy"] = v.strip()
                elif "storage tip" in k: fields["storage_tip"] = v.strip()
                elif "explanation" in k: fields["explanation"] = v.strip()
        fields["raw"] = raw_text
        return fields
    except Exception as e:
        return {"error": str(e), "raw": None}

def build_final_answer(top3: list, ai: dict | None) -> dict:
    ai_ok = ai and ai.get("crop_name") and not ai.get("error")
    return {
        "crop_name":        ai.get("crop_name")        if ai_ok else top3[0]["crop"],
        "quality":          ai.get("quality")          if ai_ok else "Unavailable",
        "market_grade":     ai.get("market_grade")     if ai_ok else "Unavailable",
        "characteristics":  ai.get("characteristics")  if ai_ok else "Unavailable",
        "explanation":      ai.get("explanation")      if ai_ok else "Unavailable",
        "storage_tip":      ai.get("storage_tip")      if ai_ok else "Unavailable",
        "confidence_label": confidence_label(top3[0]["confidence_percent"]),
    }

async def increment_usage(key_id: str):
    try:
        async with httpx.AsyncClient() as client:
            await client.post(f"{SUPABASE_URL}/rest/v1/rpc/increment_usage",
                              headers=SB_HEADERS, json={"row_id": key_id}, timeout=3)
    except: pass

async def validate_api_key(x_api_key: str = Header(..., alias="x-api-key")):
    global _cache_ts
    if time.time() - _cache_ts > 60: await refresh_key_cache()
    info = _key_cache.get(x_api_key)
    if not info:
        raise HTTPException(status_code=401, detail={"error": "Invalid API key", "portal": "/portal"})
    return info

# ── Registration ─────────────────────────────────────────────────────────────
class RegisterRequest(BaseModel):
    name: str
    email: str

@app.post("/register", tags=["Admin"])
async def register(body: RegisterRequest):
    try:
        new_key = f"crop-{secrets.token_urlsafe(20)}"
        async with httpx.AsyncClient() as client:
            # Check existing
            r = await client.get(f"{SB_TABLE}?email=eq.{body.email}&select=api_key,name", headers=SB_HEADERS)
            if r.status_code == 200 and r.json():
                return {"success": True, "api_key": r.json()[0]["api_key"], "is_new": False}
            
            # Create new
            r = await client.post(SB_TABLE, headers=SB_HEADERS, 
                                  json={"api_key": new_key, "name": body.name, "email": body.email})
            if r.status_code not in (200, 201):
                raise Exception(f"Supabase error: {r.text}")
        
        await refresh_key_cache()
        return {"success": True, "api_key": new_key, "is_new": True}
    except Exception as e:
        logger.error(f"Registration failed: {e}")
        raise HTTPException(500, detail={"error": "Failed to create key", "message": str(e)})

# ── Routes ───────────────────────────────────────────────────────────────────
@app.get("/portal", response_class=HTMLResponse)
async def portal_page():
    if not os.path.exists("portal.html"):
        return HTMLResponse("<h1>Portal page missing</h1>", status_code=404)
    return FileResponse("portal.html")

@app.get("/health")
async def health():
    return {"status": "ok", "model_ready": model is not None, "classes": len(class_names)}

@app.get("/crops")
async def list_crops():
    return {"crops": [n.replace("_"," ").title() for n in class_names]}

@app.post("/predict")
async def predict(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    key_info: dict = Depends(validate_api_key),
):
    if model is None:
        raise HTTPException(503, detail="Model still loading or failed to load. Check /health.")
    
    img_bytes = await file.read()
    if not img_bytes: raise HTTPException(422, detail="Empty image file.")
    
    processed = preprocess_image(img_bytes)
    if processed is None:
        raise HTTPException(422, detail="Invalid image file. Could not decode.")

    try:
        # Inference
        res = model.predict(processed, verbose=0)[0]
        top3_idx = np.argsort(res)[-3:][::-1]
        top3 = [{"rank": i+1, "crop": class_names[idx].replace("_"," ").title(),
                 "confidence_percent": round(float(res[idx])*100, 2),
                 "confidence_label": confidence_label(round(float(res[idx])*100, 2))}
                for i, idx in enumerate(top3_idx)]
        
        # LLaMA
        ai = call_llama_vision(img_bytes, top3)
        background_tasks.add_task(increment_usage, key_info["id"])
        
        return {
            "success": True,
            "final_answer": build_final_answer(top3, ai),
            "model_prediction": {"top3": top3},
            "ai_expert": ai,
            "request_id": str(uuid.uuid4())
        }
    except Exception as e:
        logger.error(f"Prediction failed: {e}")
        raise HTTPException(500, detail={"error": "Prediction error", "message": str(e)})

@app.post("/predict/fast")
async def predict_fast(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    key_info: dict = Depends(validate_api_key),
):
    if model is None: raise HTTPException(503, detail="Model loading.")
    img_bytes = await file.read()
    processed = preprocess_image(img_bytes)
    if processed is None: raise HTTPException(422, detail="Bad image.")
    
    res = model.predict(processed, verbose=0)[0]
    top3_idx = np.argsort(res)[-3:][::-1]
    top3 = [{"rank": i+1, "crop": class_names[idx].replace("_"," ").title(),
             "confidence_percent": round(float(res[idx])*100, 2),
             "confidence_label": confidence_label(round(float(res[idx])*100, 2))}
            for i, idx in enumerate(top3_idx)]
    
    background_tasks.add_task(increment_usage, key_info["id"])
    return {"success": True, "top_prediction": top3[0]["crop"], "top3": top3}