from typing import Dict, List, Any
import tensorflow as tf
import numpy as np
from PIL import Image
import io
import base64
import os


class EndpointHandler():
    def __init__(self, path=""):
        """
        Initialize the handler by loading the model and class labels.

        Args:
            path: Path to the model directory (provided by Hugging Face)
        """
        # Load the trained model
        model_path = os.path.join(path, "food_cassifier.h5")
        self.model = tf.keras.models.load_model(model_path)

        # Load class names from classes.txt
        classes_path = os.path.join(path, "classes.txt")
        with open(classes_path, 'r') as f:
            self.class_names = [line.strip() for line in f.readlines()]

        print(f"Model loaded successfully with {len(self.class_names)} classes")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Process the inference request.

        Args:
            data: Dictionary containing the input data
                  Expected format: {"inputs": "<base64_encoded_image>"}
                  or {"inputs": {"image": "<base64_encoded_image>"}}

        Returns:
            List of predictions with class labels and confidence scores
        """
        # Extract the input image
        inputs = data.get("inputs", "")

        # Handle different input formats
        if isinstance(inputs, dict):
            image_data = inputs.get("image", "")
        else:
            image_data = inputs

        # Decode base64 image
        try:
            # Remove data URL prefix if present
            if isinstance(image_data, str) and "base64," in image_data:
                image_data = image_data.split("base64,")[1]

            # Decode base64
            image_bytes = base64.b64decode(image_data)
            image = Image.open(io.BytesIO(image_bytes))
        except Exception as e:
            return [{"error": f"Failed to decode image: {str(e)}"}]

        # Preprocess the image
        try:
            preprocessed_image = self.preprocess_image(image)
        except Exception as e:
            return [{"error": f"Failed to preprocess image: {str(e)}"}]

        # Make prediction
        try:
            prediction = self.model.predict(preprocessed_image)
        except Exception as e:
            return [{"error": f"Failed to make prediction: {str(e)}"}]

        # Postprocess the prediction
        results = self.postprocess_prediction(prediction)

        return results

    def preprocess_image(self, image: Image.Image) -> np.ndarray:
        """
        Preprocess the input image for the model.

        Args:
            image: PIL Image object

        Returns:
            Preprocessed numpy array of shape (1, 224, 224, 3)
        """
        # Convert to RGB if needed (handles RGBA, grayscale, etc.)
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Convert to numpy array
        img_array = tf.keras.preprocessing.image.img_to_array(image)

        # Resize to (224, 224)
        img_array = tf.image.resize(img_array, (224, 224))

        # Expand dimensions to create batch: (224, 224, 3) -> (1, 224, 224, 3)
        img_array = np.expand_dims(img_array, axis=0)

        # Preprocess using MobileNetV2 preprocessing
        img_array = tf.keras.applications.mobilenet_v2.preprocess_input(img_array)

        return img_array

    def postprocess_prediction(self, prediction: np.ndarray) -> List[Dict[str, Any]]:
        """
        Postprocess the model prediction to return human-readable results.

        Args:
            prediction: Model output array of shape (1, 101)

        Returns:
            List containing prediction results with top predictions
        """
        # Get the predicted class index
        predicted_idx = int(np.argmax(prediction[0]))
        confidence = float(prediction[0][predicted_idx])
        predicted_label = self.class_names[predicted_idx]

        # Get top 5 predictions
        top_5_indices = np.argsort(prediction[0])[-5:][::-1]
        top_5_predictions = [
            {
                "label": self.class_names[int(idx)],
                "score": float(prediction[0][idx])
            }
            for idx in top_5_indices
        ]

        # Return the results
        return [
            {
                "predicted_class": predicted_label,
                "predicted_index": predicted_idx,
                "confidence": confidence,
                "top_5_predictions": top_5_predictions
            }
        ]