import os
import io
import torch
import torch.nn.functional as F
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from PIL import Image
import mlflow.pytorch
from transformers import AutoTokenizer, AutoModel
from health_multimodal.image.model.pretrained import get_biovil_t_image_encoder
from health_multimodal.image.data.transforms import create_chest_xray_transform_for_inference

app = FastAPI(title="BioVil Cross-Attention+MLP Inference API")

# Global instances for your 3 models and required processors
device = None
tokenizer = None
text_model = None
image_model = None
image_transform = None
cross_att_classifier = None

# STARTUP COMPONENT
@app.on_event("startup")
def load_all_models_and_assets():
    global device, tokenizer, text_model, image_model, image_transform, cross_att_classifier
    try:
        # Setup device configuration
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {device}")

        # Load the comprehensive BioViL-T repo for Text
        model_id = "microsoft/BiomedVLP-BioViL-T"
        
        # Specialized CXR-BERT tokenizer and model
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        text_model = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device)
        text_model.eval()

        # Instantiate the BioViL-T Image Engine
        image_model = get_biovil_t_image_encoder().to(device)
        image_transform = create_chest_xray_transform_for_inference(resize=512, center_crop_size=448)
        image_model.eval()

        # Connect to Hugging Face MLflow instance and pull the Cross-Attention Classifier
        mlflow.set_tracking_uri(os.environ.get("APP_URI"))
        model_uri = "models:/biovil_cross_attention_mlp/latest"
        
        cross_att_classifier = mlflow.pytorch.load_model(model_uri, map_location=torch.device('cpu'))
        cross_att_classifier.to(device).eval()
        
        print("All 3 models and processors loaded into memory successfully!")
    except Exception as e:
        print(f"❌ Startup Error: {str(e)}")
        raise e

# PREPROCESSING PIPELINES
def get_text_embeddings(report_text):
    inputs = tokenizer(
        report_text, 
        padding="max_length", 
        truncation=True, 
        max_length=512, 
        return_tensors="pt"
    ).to(device)
    
    with torch.no_grad():
        outputs = text_model(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            return_dict=True
        )
    return outputs.last_hidden_state

def get_image_embeddings_from_pil(pil_image):
    # Adjusted to accept the PIL image object directly from memory
    raw_image = pil_image.convert("L")
    processed_tensor = image_transform(raw_image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        image_outputs = image_model(processed_tensor)
    return image_outputs.projected_patch_embeddings 

# THE PREDICT ENDPOINT
@app.post("/predict")
async def predict(
    text_input: str = Form(...), 
    image_file: UploadFile = File(...)
):
    # Guard against queries hitting the server before models are fully loaded
    if None in (cross_att_classifier, text_model, image_model):
        raise HTTPException(status_code=503, detail="Models are initializing. Try again shortly.")
        
    try:
        # Read incoming file stream directly into memory as a PIL Image
        image_bytes = await image_file.read()
        pil_image = Image.open(io.BytesIO(image_bytes))

        # Use custom preprocessing pipelines
        sequence_outputs = get_text_embeddings(text_input)      
        patch_img_emb = get_image_embeddings_from_pil(pil_image) 

        # Run inputs through registered Cross-Attention Classifier
        with torch.no_grad():
            outputs = cross_att_classifier(patch_img_emb, sequence_outputs[:, :256, :]).squeeze(1)
            probability = torch.sigmoid(outputs).item()
            prediction = int(probability >= 0.5)

        # Return JSON response back to Streamlit
        return {
            "status": "success",
            "prediction": prediction,
            "probability": round(probability, 4)
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")