# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_prism

from threading import Lock
from typing import Any, Optional, cast

import torch
from loguru import logger


class HuggingFaceEngine:
    """
    Base class for engines using Hugging Face models with lazy loading.
    Implements the Singleton/Static Lazy Loading pattern.
    """

    # Abstract Attributes to be overridden by subclasses
    MODEL_NAME: str
    PROCESSOR_NAME: Optional[str] = None  # Optional override if processor uses a different repo
    MODEL_CLASS: Any
    PROCESSOR_CLASS: Any

    # Shared State (Types for mypy, actual values initialized in __init_subclass__)
    _model: Any = None
    _processor: Any = None
    _device: Optional[str] = None
    _lock: Lock

    def __init_subclass__(cls, **kwargs: Any) -> None:
        """
        Ensure each subclass has its own independent state (lock, model, processor).
        This prevents different engines from sharing the same lock or model attribute
        inherited from the base class.
        """
        super().__init_subclass__(**kwargs)
        cls._lock = Lock()
        cls._model = None
        cls._processor = None
        cls._device = None

    @classmethod
    def _load_model(cls) -> None:
        """
        Lazy load the model and processor.
        Detects GPU availability and falls back to CPU.
        Thread-safe execution.
        """
        if cls._model is not None:
            return

        with cls._lock:
            # Double-checked locking
            if cls._model is not None:
                return

            logger.info(f"Loading {cls.__name__} model ({cls.MODEL_NAME})...")
            try:
                if torch.cuda.is_available():
                    cls._device = "cuda"
                    logger.info("GPU detected. Using CUDA.")
                else:
                    cls._device = "cpu"
                    logger.warning("No GPU detected. Falling back to CPU. Inference will be slow.")

                # Use cast(Any, ...) to bypass strict type checks on dynamic HF classes
                # equivalent to: cls._processor = cls.PROCESSOR_CLASS.from_pretrained(cls.MODEL_NAME)
                # Use PROCESSOR_NAME if set, otherwise fallback to MODEL_NAME
                processor_name = cls.PROCESSOR_NAME if cls.PROCESSOR_NAME else cls.MODEL_NAME
                cls._processor = cast(Any, cls.PROCESSOR_CLASS).from_pretrained(processor_name)
                cls._model = cast(Any, cls.MODEL_CLASS).from_pretrained(cls.MODEL_NAME).to(cls._device)

                logger.info(f"{cls.__name__} model loaded successfully.")

            except Exception as e:
                logger.exception(f"Failed to load {cls.__name__} model")
                cls._model = None
                cls._processor = None
                raise e
