import os
import io
import cv2
import logging
import numpy as np

# ==========================================
# 0. PRODUCTION LOGGING & ENVIRONMENT SETUP
# ==========================================
# Set TensorFlow variables BEFORE importing TF to suppress warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

import tensorflow as tf
from tensorflow import keras
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware

# Configure production logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("AiroDx-API")

# ==========================================
# 1. CONFIGURATION & PATHS
# ==========================================
# Use Environment Variables for paths (Allows Docker to change them without changing code)
MODEL_PATH = os.getenv("MODEL_PATH", "./model/remedis_multiview_final.keras")
REMEDIS_BASE_PATH = os.getenv("REMEDIS_BASE_PATH", "./model/cxr-50x1-remedis-m")

IMAGE_SIZE = (448, 448)
TARGET_DISEASES = [
    'Atelectasis', 'Cardiomegaly', 'Consolidation', 'Emphysema', 'Hernia', 
    'Infiltrates', 'Mass', 'Nodule', 'Pleural Effusion', 'Pleural Thickening', 
    'Pneumonia', 'Pneumothorax', 'Pulmonary Edema', 'Pulmonary Fibrosis'
]

# ==========================================
# 2. REDEFINE CUSTOM LAYER
# ==========================================
@tf.keras.utils.register_keras_serializable()
class REMEDISBackbone(keras.layers.Layer):
    def __init__(self, model_path=None, **kwargs):
        # Accept legacy serialized model_path, but force the local deployment path.
        super().__init__(**kwargs)
        self.model_path = REMEDIS_BASE_PATH
        self._remedis_model = None
        
    def build(self, input_shape):
        if self._remedis_model is None:
            if not os.path.exists(self.model_path):
                logger.error(f"Base REMEDIS model not found at {self.model_path}")
                raise FileNotFoundError(f"Base REMEDIS model not found at {self.model_path}")
            self._remedis_model = tf.saved_model.load(self.model_path)
        super().build(input_shape)
    
    def call(self, inputs, training=None):
        return self._remedis_model(inputs)
    
    def compute_output_shape(self, input_shape):
        return (input_shape[0], 14, 14, 2048)
    
    def get_config(self):
        config = super().get_config()
        return config
    
    @classmethod
    def from_config(cls, config):
        # Remove stale serialized path (e.g., Kaggle) before layer construction.
        config.pop("model_path", None)
        return cls(**config)

# ==========================================
# 3. INITIALIZE APP & LOAD MODEL
# ==========================================
app = FastAPI(
    title="AiroDx Multi-View Diagnostic API", 
    description="Multimodal diagnostic support using Siamese REMEDIS network.",
    version="1.0.0"
)

# CORS setup for production (Can be restricted via ENV variable)
ALLOWED_ORIGINS = os.getenv("ALLOWED_ORIGINS", "*").split(",")
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS, 
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

logger.info(f"Loading weights from {MODEL_PATH}...")
try:
    model = keras.models.load_model(
        MODEL_PATH, 
        custom_objects={"REMEDISBackbone": REMEDISBackbone}
    )
    logger.info("Model loaded successfully into memory!")
except Exception as e:
    logger.critical(f"Failed to load model: {e}")
    model = None

# ==========================================
# 4. HEALTH CHECK ENDPOINT (For Cloud Hosting)
# ==========================================
@app.get("/health")
async def health_check():
    """Endpoint used by load balancers and Docker to ensure API is alive."""
    if model is None:
        raise HTTPException(status_code=503, detail="Model failed to load.")
    return {"status": "healthy", "model_ready": True}

# ==========================================
# 5. PREPROCESSING FUNCTION
# ==========================================
async def process_image(file: UploadFile) -> np.ndarray:
    try:
        # Read the file bytes
        contents = await file.read()
        nparr = np.frombuffer(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if img is None:
            raise ValueError(f"Invalid image file: {file.filename}")
            
        # Exact same preprocessing used in training
        img = cv2.resize(img, IMAGE_SIZE)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype('float32') / 255.0
        
        # Add batch dimension: (448, 448, 3) -> (1, 448, 448, 3)
        return np.expand_dims(img, axis=0)
    except Exception as e:
        logger.error(f"Image processing error: {str(e)}")
        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")

# ==========================================
# 6. PREDICTION ENDPOINT
# ==========================================
@app.post("/predict")
async def predict(
    frontal_image: UploadFile = File(..., description="The PA/AP Frontal X-ray"),
    lateral_image: UploadFile = File(..., description="The Lateral X-ray")
):
    logger.info(f"Received prediction request for files: {frontal_image.filename}, {lateral_image.filename}")
    
    # 1. Process both images
    img_pa = await process_image(frontal_image)
    img_l = await process_image(lateral_image)
    
    # 2. Run inference
    predictions = model.predict({
        "pa_image": img_pa,
        "l_image": img_l
    })
    
    # 3. Format results
    probabilities = predictions[0].tolist() 
    
    results = []
    for disease, prob in zip(TARGET_DISEASES, probabilities):
        results.append({
            "disease": disease,
            "probability": round(prob, 4),
            "flagged": bool(prob > 0.5) # You can adjust this threshold later based on your F1 tuning
        })
    
    # Sort results by probability (highest first)
    results.sort(key=lambda x: x["probability"], reverse=True)
    
    logger.info("Prediction processed successfully.")
    
    return {
        "status": "success",
        "predictions": results
    }