# 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 pathlib import Path
from typing import Any, List, Union, cast

import open_clip  # NEW: We use open_clip instead of transformers
import torch
from loguru import logger
from PIL import Image

from coreason_prism.utils.engine import HuggingFaceEngine


class Embedder(HuggingFaceEngine):
    """The Embedder (Multi-Modal Vectors).

    Generates embeddings for non-text objects using BioCLIP.
    Projects molecular graphs or histology images into a vector space.
    """

    # We use 'hf-hub:' prefix to tell open_clip to look at Hugging Face
    MODEL_NAME = "hf-hub:imageomics/bioclip"

    # These are kept for compatibility with the base class structure,
    # though we override the loading logic below.
    PROCESSOR_NAME = None
    MODEL_CLASS = None
    PROCESSOR_CLASS = None

    # New shared state for the tokenizer (specific to open_clip)
    _tokenizer: Any = None

    @classmethod
    def _load_model(cls) -> None:
        """
        Lazy load the model and processor using open_clip.
        """
        if cls._model is not None:
            return

        with cls._lock:
            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.")

                # Load via open_clip
                model, _, preprocess = open_clip.create_model_and_transforms(cls.MODEL_NAME)
                tokenizer = open_clip.get_tokenizer(cls.MODEL_NAME)

                cls._model = model.to(cls._device)
                cls._processor = preprocess  # This is the image transform
                cls._tokenizer = tokenizer

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

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

    @classmethod
    def embed_image(cls, image_path: Union[str, Path]) -> List[float]:
        """Generate an embedding for a given image using OpenCLIP."""
        path = Path(image_path)
        if not path.exists():
            logger.error(f"Image file not found: {path}")
            raise FileNotFoundError(f"Image file not found: {path}")

        cls._load_model()
        assert cls._model is not None
        assert cls._processor is not None
        assert cls._device is not None

        try:
            logger.info(f"Embedding image: {path}")
            image = Image.open(path).convert("RGB")

            # 1. Preprocess image
            # processor is the open_clip transform function here
            preprocessed_image = cls._processor(image).unsqueeze(0).to(cls._device)

            # 2. Encode
            with torch.no_grad():
                image_features = cls._model.encode_image(preprocessed_image)
                # Normalize features (standard practice for CLIP models)
                image_features /= image_features.norm(dim=-1, keepdim=True)

            embedding = image_features.squeeze().tolist()
            if isinstance(embedding, float):
                return [embedding]
            return cast(List[float], embedding)

        except Exception as e:
            logger.exception(f"Error embedding image {path}")
            raise e

    @classmethod
    def embed_text(cls, text: str) -> List[float]:
        """Generate an embedding for a given text using OpenCLIP."""
        if not text:
            logger.error("Empty text provided for embedding")
            raise ValueError("Text cannot be empty")

        cls._load_model()
        assert cls._model is not None
        assert cls._tokenizer is not None
        assert cls._device is not None

        try:
            logger.info("Embedding text input")

            # 1. Tokenize
            text_tokens = cls._tokenizer([text]).to(cls._device)

            # 2. Encode
            with torch.no_grad():
                text_features = cls._model.encode_text(text_tokens)
                # Normalize features
                text_features /= text_features.norm(dim=-1, keepdim=True)

            embedding = text_features.squeeze().tolist()
            if isinstance(embedding, float):
                return [embedding]
            return cast(List[float], embedding)

        except Exception as e:
            logger.exception("Error embedding text")
            raise e
