import os
from contextlib import asynccontextmanager

# ─── 1. MANDATORY BACKEND CONFIG (Must be at the very top) ───────────────────
os.environ["KERAS_BACKEND"] = "jax"
os.environ["JAX_PLATFORM_NAME"] = "cpu"
os.environ["HF_HOME"] = "/tmp/huggingface"

import base64
import logging
import numpy as np
import cv2
import requests
import tempfile
from typing import List, Dict, Any, Optional

import keras
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from huggingface_hub import hf_hub_download
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
from inference_sdk import InferenceHTTPClient

# ─── 2. LOGGING & CONFIG ─────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

API_KEY = os.getenv("ROBOFLOW_API_KEY")
WORKSPACE_NAME = "my-workspace-iy9ll"
WORKFLOW_ID = "find-teeth"
IMG_SIZE = 256
THRESHOLD = 0.05
MODEL_REPO = "sairamv1/TeethCariesDetection_CNNAutoencoder"
MODEL_FILENAME = "autoencoder.keras"

MODEL = None
roboflow_client = None

# ─── 3. RETRY-ENABLED MODEL DOWNLOAD ─────────────────────────────────────────
@retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=10),
    retry=retry_if_exception_type((requests.RequestException, OSError))
)
def download_model() -> str:
    """Download model from HuggingFace Hub with retry logic."""
    logger.info(f"🚀 Downloading model from {MODEL_REPO}...")
    local_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=MODEL_FILENAME,
        cache_dir="/tmp/huggingface"
    )
    logger.info(f"✅ Model downloaded to {local_path}")
    return local_path

def init_roboflow_client():
    """Initialize Roboflow client if API key is available."""
    global roboflow_client
    if API_KEY:
        try:
            roboflow_client = InferenceHTTPClient(
                api_url="https://serverless.roboflow.com",
                api_key=API_KEY
            )
            logger.info("✅ Roboflow client initialized")
        except Exception as e:
            logger.error(f"❌ Roboflow client initialization failed: {e}")
            roboflow_client = None
    else:
        logger.warning("⚠️ ROBOFLOW_API_KEY not set")

# ─── 4. LIFESPAN EVENT HANDLER ───────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage model loading at startup and cleanup at shutdown."""
    global MODEL
    try:
        local_model_path = download_model()
        logger.info("🔄 Loading Keras model with JAX backend...")
        MODEL = keras.saving.load_model(local_model_path)
        logger.info(f"✅ Model loaded successfully! Input shape: {MODEL.input_shape}")
        
        # Initialize Roboflow client
        init_roboflow_client()
    except Exception as e:
        logger.error(f"❌ Startup failed: {e}")
        MODEL = None
    
    yield
    logger.info("🛑 Shutting down...")

app = FastAPI(
    title="Teeth Caries Detection API",
    description="Dental anomaly detection using CNN Autoencoder",
    version="1.0.0",
    lifespan=lifespan
)

# ─── 5. ROBOFLOW API CALL ────────────────────────────────────────────────────
def call_roboflow_api(image_path: str) -> Dict:
    """Call Roboflow workflow using inference SDK."""
    if not roboflow_client:
        raise Exception("Roboflow client not initialized")
    
    result = roboflow_client.run_workflow(
        workspace_name=WORKSPACE_NAME,
        workflow_id=WORKFLOW_ID,
        images={"image": image_path},
        use_cache=False
    )
    return result

# ─── 6. ROUTES ───────────────────────────────────────────────────────────────
@app.get("/")
async def root():
    """Health check endpoint."""
    return {
        "status": "active",
        "model_ready": MODEL is not None,
        "backend": os.environ.get("KERAS_BACKEND"),
        "api_configured": API_KEY is not None,
        "roboflow_ready": roboflow_client is not None
    }

@app.get("/health")
async def health():
    """Detailed health check."""
    if MODEL is None:
        return JSONResponse(
            status_code=503,
            content={"status": "unhealthy", "reason": "Model not loaded"}
        )
    return {
        "status": "healthy",
        "model_input_shape": str(MODEL.input_shape),
        "roboflow_status": "ready" if roboflow_client else "not configured"
    }

@app.post("/analyze")
async def analyze(file: UploadFile = File(...)):
    """
    Analyze dental image for caries detection.
    
    Returns anomaly scores for each detected tooth.
    """
    if MODEL is None:
        raise HTTPException(status_code=503, detail="Model not initialized. Please retry later.")
    
    if not roboflow_client:
        raise HTTPException(status_code=500, detail="ROBOFLOW_API_KEY not configured")

    # ─── Read and validate image ─────────────────────────────────────────────
    try:
        contents = await file.read()
        if len(contents) == 0:
            raise HTTPException(status_code=400, detail="Empty file uploaded")
        
        nparr = np.frombuffer(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img is None or img.size == 0:
            raise HTTPException(status_code=400, detail="Invalid image format or corrupted file")
        
        img_h, img_w = img.shape[:2]
        if img_h < 50 or img_w < 50:
            raise HTTPException(status_code=400, detail="Image too small (min 50x50 pixels)")
            
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Image decode error: {e}")
        raise HTTPException(status_code=400, detail=f"Image processing failed: {str(e)}")

    # ─── Call Roboflow detection ─────────────────────────────────────────────
    tmp_path = None
    try:
        # Save to temp file for inference_sdk
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
            tmp.write(contents)
            tmp_path = tmp.name
        
        result_json = call_roboflow_api(tmp_path)
        
        # ─── DEBUG: Log what Roboflow returns ────────────────────────────────
        logger.info(f"🔍 Response Type: {type(result_json)}")
        logger.info(f"🔍 Response Keys: {result_json.keys() if isinstance(result_json, dict) else 'Not a dict'}")
        logger.info(f"🔍 Full Response: {result_json}")
        
    except Exception as e:
        logger.error(f"Roboflow workflow error: {e}")
        raise HTTPException(status_code=502, detail=f"Detection service failed: {str(e)}")
    finally:
        # Clean up temp file
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)

        # ─── Parse detections ────────────────────────────────────────────────────
    predictions = []
    
    try:
        # Handle different response structures
        if isinstance(result_json, list) and len(result_json) > 0:
            # Response is a list
            first_item = result_json[0]
            logger.info(f"🔍 First item type: {type(first_item)}")
            logger.info(f"🔍 First item keys: {first_item.keys() if isinstance(first_item, dict) else 'Not dict'}")
            
            # FIXED: Check for "predictions" directly first!
            if isinstance(first_item, dict):
                if "predictions" in first_item:
                    predictions = first_item["predictions"]
                    logger.info(f"✅ Found predictions directly in first_item")
                elif "outputs" in first_item:
                    outputs = first_item["outputs"]
                    logger.info(f"🔍 Found outputs: {type(outputs)}")
                    if isinstance(outputs, dict):
                        predictions = outputs.get("predictions", [])
                        
        elif isinstance(result_json, dict):
            # Response is a dict
            logger.info(f"🔍 Dict keys: {list(result_json.keys())}")
            
            # Try "predictions" first
            if "predictions" in result_json:
                predictions = result_json["predictions"]
            elif "outputs" in result_json:
                outputs = result_json["outputs"]
                if isinstance(outputs, list) and len(outputs) > 0:
                    output_data = outputs[0]
                elif isinstance(outputs, dict):
                    output_data = outputs
                else:
                    output_data = {}
                predictions = output_data.get("predictions", [])
        
        # Handle case where predictions might be a dict
        if isinstance(predictions, dict) and "predictions" in predictions:
            predictions = predictions["predictions"]
        
        logger.info(f"🔍 Final predictions type: {type(predictions)}")
        logger.info(f"🔍 Final predictions count: {len(predictions) if isinstance(predictions, list) else 0}")
        if isinstance(predictions, list) and len(predictions) > 0:
            logger.info(f"🔍 First prediction keys: {predictions[0].keys()}")
            
    except Exception as e:
        logger.error(f"❌ Parsing error: {e}", exc_info=True)
        predictions = []
    
    if not predictions or (isinstance(predictions, list) and len(predictions) == 0):
        logger.warning("⚠️ No predictions found after parsing")
        return {
            "status": "success",
            "tooth_count": 0,
            "results": [],
            "message": "No teeth detected in image"
        }
    
    logger.info(f"✅ Found {len(predictions)} teeth detections")


    # ─── Process crops ───────────────────────────────────────────────────────
    batch_images = []
    metadata = []

    for i, det in enumerate(predictions[:12]):
        try:
            xc, yc, w, h = det["x"], det["y"], det["width"], det["height"]
            x1 = max(0, int(xc - w/2))
            y1 = max(0, int(yc - h/2))
            x2 = min(img_w, int(xc + w/2))
            y2 = min(img_h, int(yc + h/2))
            
            if x2 <= x1 or y2 <= y1:
                continue
            
            crop = img[y1:y2, x1:x2]
            if crop.size == 0:
                continue

            # Preprocessing pipeline
            gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            gray = clahe.apply(gray)
            
            # Aspect-ratio preserving resize with padding
            h_c, w_c = gray.shape
            scale = IMG_SIZE / max(h_c, w_c)
            nw, nh = int(w_c * scale), int(h_c * scale)
            resized = cv2.resize(gray, (nw, nh), interpolation=cv2.INTER_AREA)
            
            canvas = np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.uint8)
            y_offset = (IMG_SIZE - nh) // 2
            x_offset = (IMG_SIZE - nw) // 2
            canvas[y_offset:y_offset + nh, x_offset:x_offset + nw] = resized
            
            final_img = (canvas.astype(np.float32) / 255.0)[..., np.newaxis]
            batch_images.append(final_img)
            metadata.append({"id": i, "box": [x1, y1, x2, y2]})
            
            logger.info(f"✅ Processed tooth {i}: bbox=[{x1},{y1},{x2},{y2}]")
            
        except Exception as e:
            logger.warning(f"⚠️ Skipping detection {i} due to error: {e}")
            continue

    if not batch_images:
        return {
            "status": "success",
            "tooth_count": 0,
            "results": [],
            "message": "No valid tooth crops extracted"
        }

    logger.info(f"✅ Prepared {len(batch_images)} tooth images for inference")

    # ─── Model inference ─────────────────────────────────────────────────────
    try:
        batch_array = np.array(batch_images)
        logger.info(f"🔄 Running inference on batch shape: {batch_array.shape}")
        reconstructions = MODEL.predict(batch_array, verbose=0)
        errors = np.mean(np.abs(batch_array - reconstructions), axis=(1, 2, 3))
        logger.info(f"✅ Inference complete. Scores: {errors}")
    except Exception as e:
        logger.error(f"❌ Inference error: {e}")
        raise HTTPException(status_code=500, detail="Model inference failed")

    # ─── Format results ──────────────────────────────────────────────────────
    results = []
    for idx, score in enumerate(errors):
        score_val = float(score)
        results.append({
            "tooth_index": metadata[idx]["id"],
            "anomaly_score": round(score_val, 4),
            "is_caries_detected": score_val > THRESHOLD,
            "bounding_box": metadata[idx]["box"],
            "confidence": "high" if score_val > THRESHOLD * 1.5 else "medium" if score_val > THRESHOLD else "low"
        })

    logger.info(f"🎉 Analysis complete: {len(results)} teeth analyzed")
    
    return {
        "status": "success",
        "tooth_count": len(results),
        "detection_threshold": THRESHOLD,
        "results": results
    }

# ─── 7. ERROR HANDLER ────────────────────────────────────────────────────────
@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
    logger.error(f"Unhandled exception: {exc}", exc_info=True)
    return JSONResponse(
        status_code=500,
        content={"detail": "Internal server error occurred"}
    )