from PIL import Image
import numpy as np
from fastapi import File, UploadFile
import os
import tensorflow as tf
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Animal names here
ANIMALS = ['Cat', 'Dog', 'Panda']

# Model path - check multiple possible locations for different model formats
model_path = None
model_type = None  # 'savedmodel' or 'keras'

# Check for SavedModel format (try both old and new naming conventions)
savedmodel_paths = [
    "animal-classification/INPUT_model_path/animal-cnn/savedmodel",
    "animal-classification/animal-cnn/savedmodel",
    "/app/animal-classification/INPUT_model_path/animal-cnn/savedmodel",
    "/app/animal-classification/animal-cnn/savedmodel",
    "animal-classification/INPUT_model_path/animal-classification/animal-cnn-savedmodel",
    "animal-classification/animal-cnn-savedmodel",
    "/app/animal-classification/INPUT_model_path/animal-classification/animal-cnn-savedmodel",
    "/app/animal-classification/animal-cnn-savedmodel"
]

# Check for Keras format (.keras file)
keras_paths = [
    "animal-classification/INPUT_model_path/animal-cnn/model.keras",
    "animal-classification/animal-cnn/model.keras",
    "/app/animal-classification/INPUT_model_path/animal-cnn/model.keras",
    "/app/animal-classification/animal-cnn/model.keras"
]

for path in savedmodel_paths:
    if os.path.exists(path):
        model_path = path
        model_type = 'savedmodel'
        break

if not model_path:
    for path in keras_paths:
        if os.path.exists(path):
            model_path = path
            model_type = 'keras'
            break

if not model_path:
    # Fallback: try to find any model in the directory structure
    model_base = "animal-classification"
    print(f"Current working directory: {os.getcwd()}")
    print(f"Files in current directory: {os.listdir('.')}")

    if os.path.exists(model_base):
        print(f"Contents of {model_base}:")
        for root, dirs, files in os.walk(model_base):
            level = root.replace(model_base, '').count(os.sep)
            indent = ' ' * 2 * level
            print(f"{indent}{os.path.basename(root)}/")
            subindent = ' ' * 2 * (level + 1)
            for file in files[:10]:
                print(f"{subindent}{file}")

            # Check for SavedModel directories (multiple naming conventions)
            if 'savedmodel' in dirs:
                model_path = os.path.join(root, 'savedmodel')
                model_type = 'savedmodel'
                break
            if 'animal-cnn-savedmodel' in dirs:
                model_path = os.path.join(root, 'animal-cnn-savedmodel')
                model_type = 'savedmodel'
                break
            # Check if any directory contains saved_model.pb (indicating SavedModel format)
            for dir_name in dirs:
                potential_savedmodel = os.path.join(root, dir_name)
                if os.path.exists(os.path.join(potential_savedmodel, 'saved_model.pb')):
                    model_path = potential_savedmodel
                    model_type = 'savedmodel'
                    break
            if model_path:
                break

            # Check for .keras files
            for file in files:
                if file.endswith('.keras'):
                    model_path = os.path.join(root, file)
                    model_type = 'keras'
                    break
            if model_path:
                break

        if not model_path:
            raise FileNotFoundError(
                f"Could not find any model (SavedModel or .keras) in {model_base}. Directory structure printed above.")
    else:
        raise FileNotFoundError(
            f"Model directory {model_base} not found. Current directory: {os.getcwd()}, Contents: {os.listdir('.')}")

print(f"Loading model from: {model_path}")
print(f"Model type: {model_type}")
print(f"Model path exists: {os.path.exists(model_path)}")

# Load the model based on its type
try:
    if model_type == 'savedmodel':
        loaded_model = tf.saved_model.load(model_path)
        infer = loaded_model.signatures["serving_default"]
        print("SavedModel loaded successfully!")
    else:  # keras
        # Try loading with compile=False to avoid optimizer/loss issues
        try:
            loaded_model = tf.keras.models.load_model(
                model_path, compile=False)
            print("Keras model loaded successfully (compile=False)!")
        except Exception as e1:
            print(f"Failed to load with compile=False: {e1}")
            # Try with safe_mode if available (newer Keras versions)
            try:
                loaded_model = tf.keras.models.load_model(
                    model_path, safe_mode=False)
                print("Keras model loaded successfully (safe_mode=False)!")
            except Exception as e2:
                print(f"Failed to load with safe_mode=False: {e2}")
                raise
        # For keras models, we'll use the model directly, not via signatures
        infer = None
except Exception as e:
    print(f"Error loading model: {e}")
    import traceback
    traceback.print_exc()
    raise


@app.get('/health')
async def health():
    return {"status": "healthy"}


@app.post('/upload/image')
async def uploadImage(img: UploadFile = File(...)):
    # Image inlezen
    original_image = Image.open(img.file)
    resized_image = original_image.resize((64, 64))
    images_to_predict = np.expand_dims(
        np.array(resized_image), axis=0).astype(np.float32)

    # Predict based on model type
    if model_type == 'savedmodel':
        # Tensor maken en infer voor SavedModel
        input_tensor = tf.convert_to_tensor(images_to_predict)
        result = infer(input_tensor)
        predictions = list(result.values())[0].numpy()
    else:  # keras
        # Direct prediction voor Keras model
        predictions = loaded_model.predict(images_to_predict, verbose=0)

    classification = predictions.argmax(axis=1)[0]
    return ANIMALS[classification]