from typing import Dict, List, Optional
import logging

logger = logging.getLogger(__name__)


class CacheManager:
    _cache: Dict[str, Dict] = {}

    @classmethod
    def set_prompt_details(cls, prompt_id: str, details: Dict) -> None:
        """Store prompt details in cache with version info"""
        logger.info(f"Caching prompt details for {prompt_id}: {details}")
        cls._cache[prompt_id] = details

    @classmethod
    def get_prompt_details(cls, prompt_id: str) -> Optional[Dict]:
        """Get prompt details from cache"""
        details = cls._cache.get(prompt_id)
        logger.info(f"Retrieved prompt details for {prompt_id}: {details}")
        return details

    @classmethod
    def has_prompt_details(cls, prompt_id: str) -> bool:
        """Check if prompt details exist in cache"""
        return prompt_id in cls._cache

    @classmethod
    def print_cache_contents(cls) -> None:
        """Print all contents of the cache for debugging"""
        print("\n=== Cache Contents ===")
        if not cls._cache:
            print("Cache is empty")
        for key, value in cls._cache.items():
            print(f"\nKey: {key}")
            print(f"Value: {value}")


class InteractionCacheManager:
    _interaction_cache: Dict[str, Dict] = (
        {}
    )  # Structure: {prompt_id: {session_id: {version: {messages: [], lastResponseAt: None}}}}
    _user_messages: Dict[str, Dict] = {}
    _prompt_version_messages: Dict[str, List] = {}

    @classmethod
    def save_user_messages(cls, prompt_id: str, messages: Dict) -> None:
        """Save user messages with system message and variables"""
        logger.info(f"Saving user messages for prompt_id: {prompt_id}")
        logger.debug(f"Messages content: {messages}")
        cls._user_messages[prompt_id] = messages

    @classmethod
    def save_interaction(
        cls,
        session_id: str,
        interaction_data: Dict,
        prompt_id: str,
        version: str,
    ) -> None:
        """
        Save interaction with hierarchical tracking
        Structure: prompt_id -> session_id -> version -> messages
        """
        logger.info(
            f"Saving interaction - Prompt ID: {prompt_id}, Session: {session_id}, Version: {version}"
        )

        # Initialize prompt_id if not exists
        if prompt_id not in cls._interaction_cache:
            cls._interaction_cache[prompt_id] = {}
            logger.debug(f"Created new prompt cache for {prompt_id}")

        # Initialize session if not exists
        if session_id not in cls._interaction_cache[prompt_id]:
            cls._interaction_cache[prompt_id][session_id] = {}
            logger.debug(f"Created new session cache for {session_id}")

        # Initialize version if not exists
        if version not in cls._interaction_cache[prompt_id][session_id]:
            cls._interaction_cache[prompt_id][session_id][version] = {
                "messages": [],
                "lastResponseAt": None,
            }
            logger.debug(f"Created new version cache for {version}")

        # Add new messages
        cls._interaction_cache[prompt_id][session_id][version]["messages"].extend(
            interaction_data["messages"]
        )
        cls._interaction_cache[prompt_id][session_id][version]["lastResponseAt"] = (
            interaction_data["lastResponseAt"]
        )

        logger.info(
            f"Saved interaction with {len(interaction_data['messages'])} messages"
        )
        logger.debug(
            f"Current cache state for {prompt_id}/{session_id}/{version}: "
            f"{cls._interaction_cache[prompt_id][session_id][version]}"
        )

    @classmethod
    def get_interaction_history(
        cls, prompt_id: str, session_id: str, version: str
    ) -> Optional[Dict]:
        """Get complete interaction history for a specific prompt, session and version"""
        try:
            history = (
                cls._interaction_cache.get(prompt_id, {})
                .get(session_id, {})
                .get(version, {})
            )
            logger.info(
                f"Retrieved interaction history - Prompt ID: {prompt_id}, "
                f"Session: {session_id}, Version: {version}"
            )
            logger.debug(f"History content: {history}")
            return history
        except Exception as e:
            logger.error(f"Error retrieving interaction history: {str(e)}")
            return None

    @classmethod
    def get_all_prompt_interactions(cls, prompt_id: str) -> Dict:
        """Get all interactions for a specific prompt across all sessions and versions"""
        interactions = cls._interaction_cache.get(prompt_id, {})
        logger.info(f"Retrieved all interactions for prompt ID: {prompt_id}")
        logger.debug(f"Interactions content: {interactions}")
        return interactions

    @classmethod
    def clear_cache(cls) -> None:
        """Clear all caches"""
        cls._interaction_cache.clear()
        cls._user_messages.clear()
        cls._prompt_version_messages.clear()
        logger.info("All caches cleared")

    @classmethod
    def print_cache_contents(cls) -> None:
        """Print all contents of different caches for debugging"""
        print("\n=== Interaction Cache Contents ===")
        if not cls._interaction_cache:
            print("Interaction cache is empty")
        for prompt_id, sessions in cls._interaction_cache.items():
            print(f"\nPrompt ID: {prompt_id}")
            for session_id, versions in sessions.items():
                print(f"  Session ID: {session_id}")
                for version, data in versions.items():
                    print(f"    Version: {version}")
                    print(f"    Messages: {data['messages']}")
                    print(f"    Last Response At: {data['lastResponseAt']}")

        print("\n=== User Messages Cache Contents ===")
        if not cls._user_messages:
            print("User messages cache is empty")
        for key, value in cls._user_messages.items():
            print(f"\nPrompt ID: {key}")
            print(f"Messages: {value}")
