from collections import OrderedDict
from functools import wraps

from common.logger import logger

# Global cache for prompt embeddings
# Key is (pipeline_type, args, kwargs)
GLOBAL_PROMPT_CACHE = OrderedDict()
MAX_PROMPT_CACHE_SIZE = 64


def clear_global_prompt_cache():
    """Clear the global prompt embeddings cache."""
    GLOBAL_PROMPT_CACHE.clear()
    logger.debug("Global prompt cache cleared")


def enable_prompt_caching(pipeline):
    """
    Generic wrapper to cache the results of encode_prompt on any diffusers pipeline.
    Uses a global cache shared across pipeline instances to avoid redundant text encoding
    even when pipelines are reloaded.
    """
    if not hasattr(pipeline, "encode_prompt"):
        logger.warning("Pipeline does not have encode_prompt method; cannot enable prompt caching")
        return pipeline

    if hasattr(pipeline, "_prompt_cache_enabled"):
        return pipeline  # Already enabled

    original_encode_prompt = pipeline.encode_prompt
    pipeline_identity = pipeline.__class__.__name__

    def make_hashable(obj):
        if isinstance(obj, (list, tuple)):
            return tuple(make_hashable(i) for i in obj)
        if isinstance(obj, dict):
            return tuple(sorted((k, make_hashable(v)) for k, v in obj.items()))
        return obj

    @wraps(original_encode_prompt)
    def wrapped_encode_prompt(*args, **kwargs):
        try:
            # Create a cache key from identity and hashable representation of all arguments
            # Identity ensures we don't use Flux embeddings for a Wan model, etc.
            cache_key = (pipeline_identity, make_hashable(args), make_hashable(kwargs))
        except (TypeError, ValueError):
            # Fallback: if something isn't hashable, just compute normally
            return original_encode_prompt(*args, **kwargs)

        if cache_key in GLOBAL_PROMPT_CACHE:
            # Move to end (Most Recently Used)
            GLOBAL_PROMPT_CACHE.move_to_end(cache_key)
            logger.info("Using cached prompt embeddings")
            return GLOBAL_PROMPT_CACHE[cache_key]

        # Compute new results (e.g., prompt_embeds, negative_prompt_embeds)
        result = original_encode_prompt(*args, **kwargs)

        # Store in global cache
        GLOBAL_PROMPT_CACHE[cache_key] = result
        if len(GLOBAL_PROMPT_CACHE) > MAX_PROMPT_CACHE_SIZE:
            GLOBAL_PROMPT_CACHE.popitem(last=False)  # Remove Least Recently Used

        return result

    # Monkey patch the instance method
    pipeline.encode_prompt = wrapped_encode_prompt
    pipeline._prompt_cache_enabled = True
    return pipeline