import io
import os
import time

import numpy as np
import onnxruntime as ort
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
from torchvision.models import efficientnet_b0

# Configuration
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODEL_DIR = os.path.join(BASE_DIR, "model")
THAI_FOOD_CLASSES = [
    "Green_Curry", 
    "Khao_phat", 
    "Khao_soi", 
    "Massaman_Curry",
    "Pad_Thai", 
    "Phanaeng_Curry", 
    "Phat_kaphrao", 
    "Roti_canai",
    "Tom_kha_gai", 
    "Tom_yum"
]

# Global dictionary to hold model instances in memory
models = {}

def init_worker():
    """Initialize models in the worker process to avoid reloading on each request."""
    global models
    try:
        # 1. Load ONNX
        onnx_path = os.path.join(MODEL_DIR, "efficientnet-b0.onnx")
        if os.path.exists(onnx_path):
            models['onnx'] = ort.InferenceSession(onnx_path)
        else:
            print(f"Warning: File not found {onnx_path}")

        # 2. Load FP16 ONNX (used as the quantized version)
        fp16_path = os.path.join(MODEL_DIR, "efficientnet-b0_fp16.onnx")
        if os.path.exists(fp16_path):
            models['quantized'] = ort.InferenceSession(fp16_path)
        else:
            print(f"Warning: File not found {fp16_path}")

        # 3. Load Original PyTorch with custom weights
        pth_path = os.path.join(MODEL_DIR, "best_thai_food_model.pth")
        models['original'] = efficientnet_b0(pretrained=False)
        models['original'].classifier[1] = nn.Linear(
            models['original'].classifier[1].in_features, 
            len(THAI_FOOD_CLASSES)
        )
        
        if os.path.exists(pth_path):
            models['original'].load_state_dict(torch.load(pth_path, map_location='cpu'))
        else:
            print(f"Warning: File not found {pth_path} (Using random weights)")
            
        models['original'].eval()

    except Exception as e:
        print(f"Error loading models in worker: {e}")

def preprocess_image(image_bytes: bytes):
    """Convert bytes to a normalized tensor suitable for EfficientNet."""
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return transform(image).unsqueeze(0)

def get_model_size_mb(model_type: str) -> float:
    """Helper function to calculate model file size in MB."""
    if model_type == "original":
        path = os.path.join(MODEL_DIR, "best_thai_food_model.pth")
    elif model_type == "onnx":
        path = os.path.join(MODEL_DIR, "efficientnet-b0.onnx")
    elif model_type == "quantized":
        path = os.path.join(MODEL_DIR, "efficientnet-b0_fp16.onnx")
    else:
        return 0.0

    if os.path.exists(path):
        return os.path.getsize(path) / (1024 * 1024)
    return 0.0

def run_inference(image_bytes: bytes, model_type: str):
    """Main inference function executed by the process pool."""
    global models
    
    if model_type not in models:
        raise ValueError(f"Model {model_type} is not loaded.")
    
    input_tensor = preprocess_image(image_bytes)
    start_time = time.perf_counter()
    
    if model_type == "original":
        with torch.no_grad():
            logits = models['original'](input_tensor)
            probs = F.softmax(logits, dim=1)
            confidence, class_idx = torch.max(probs, 1)
            prediction_id = class_idx.item()
            confidence_score = confidence.item()
            
    elif model_type in ["onnx", "quantized"]:
        ort_sess = models[model_type]
        input_data = input_tensor.numpy()
        
        # Dynamically cast input to float16 if required by the model
        input_type = ort_sess.get_inputs()[0].type
        if 'float16' in input_type:
            input_data = input_data.astype(np.float16)
            
        ort_inputs = {ort_sess.get_inputs()[0].name: input_data}
        logits = ort_sess.run(None, ort_inputs)[0]
        
        # Numpy Softmax
        exp_logits = np.exp(logits - np.max(logits))
        probs = exp_logits / exp_logits.sum(axis=1, keepdims=True)
        prediction_id = int(np.argmax(probs[0]))
        confidence_score = float(probs[0][prediction_id])
        
    latency_ms = (time.perf_counter() - start_time) * 1000
    model_size_mb = get_model_size_mb(model_type)

    return {
        "model_type": model_type,
        "prediction_class_id": prediction_id,
        "prediction_class_name": THAI_FOOD_CLASSES[prediction_id],
        "confidence_score": round(confidence_score * 100, 2),
        "latency_ms": round(latency_ms, 2),
        "model_size_mb": round(model_size_mb, 2)
    }