import os
import shutil
import logging
import tensorflow as tf
from tensorflow.keras.layers import Layer
from huggingface_hub import snapshot_download
from config import Config

# Model config
REPO_ID = Config.IMAGE_CLASSIFIER_REPO_ID
MODEL_DIR = Config.IMAGE_CLASSIFIER_MODEL_DIR
WEIGHTS_PATH = os.path.join(MODEL_DIR, Config.IMAGE_CLASSIFIER_WEIGHTS_FILE)
HF_TOKEN = Config.HF_TOKEN

# Device info (for logging)
gpus = tf.config.list_physical_devices("GPU")
device = "cuda" if gpus else "cpu"

# Global model reference
_model_img = None

# Custom layer used in the model
class Cast(Layer):
    def call(self, inputs):
        return tf.cast(inputs, tf.float32)

def warmup():
    global _model_img
    download_model_repo()
    _model_img = load_model()
    logging.info("Image model is ready.")

def download_model_repo():
    if os.path.exists(MODEL_DIR) and os.path.isdir(MODEL_DIR):
        logging.info("Image model already exists, skipping download.")
        return
    snapshot_path = snapshot_download(repo_id=REPO_ID, token=HF_TOKEN)
    os.makedirs(MODEL_DIR, exist_ok=True)
    shutil.copytree(snapshot_path, MODEL_DIR, dirs_exist_ok=True)

def load_model():
    global _model_img
    if _model_img is not None:
        return _model_img

    print(f"{'GPU detected' if device == 'cuda' else 'No GPU detected'}, loading model on {device.upper()}.")

    _model_img = tf.keras.models.load_model(
        WEIGHTS_PATH, custom_objects={"Cast": Cast}
    )
    print("Model input shape:", _model_img.input_shape)
    return _model_img

def get_model():
    global _model_img
    if _model_img is None:
        download_model_repo()
        _model_img = load_model()
    return _model_img