import os
import json
import uuid
import threading
from typing import Dict, Optional, List, Type, Any, Literal, cast
from datetime import datetime
import redis
from SimpleLLMFunc import OpenAICompatible
from context.schemas import Message
from context.context import ContextBackend, RedisFileContextBackend
from config.config import get_config
from SimpleLLMFunc.logger import push_warning, app_log


class ContextManager:
    """
    General-purpose context manager that supports different backend implementations.

    Main responsibilities:
    1. Manage creation and lifecycle of ContextBackend instances
    2. Provide advanced convenience interfaces
    3. Handle batch operations and cleanup tasks
    4. Support pluggable backend implementations
    """

    _instance = None
    _lock: threading.Lock = threading.Lock()

    def __new__(cls, backend_class: Type[ContextBackend] = RedisFileContextBackend):
        """Singleton pattern implementation."""
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super(ContextManager, cls).__new__(cls)
                    cls._instance.backend_class = backend_class
        return cls._instance

    def __init__(self, backend_class: Type[ContextBackend]):
        """
        Initialize the context manager.

        Args:
            backend_class: Backend implementation class, defaulting to RedisFileBackend
        """
        # Prevent duplicate initialization.
        if hasattr(self, "_initialized"):
            return

        self.backend_class = backend_class
        self.config = get_config()
        self.context_dir = self.config.CONTEXT_DIR
        self._active_contexts: Dict[str, ContextBackend] = {}

        # Ensure the directory exists.
        os.makedirs(self.context_dir, exist_ok=True)

        self._initialized = True

    def _redis_client(self) -> redis.Redis:
        return redis.Redis(
            host=self.config.REDIS_HOST,
            port=int(self.config.REDIS_PORT),
            db=int(self.config.REDIS_DB),
            decode_responses=True,
        )

    def _list_context_ids_from_redis(self) -> set[str]:
        context_ids: set[str] = set()
        try:
            client = self._redis_client()
            raw_keys = cast(Any, client.keys("context:*:*"))
            for key in cast(List[str], raw_keys):
                parts = key.split(":", 2)
                if len(parts) >= 3 and parts[0] == "context" and parts[1]:
                    context_ids.add(parts[1])
        except Exception as e:
            print(f"Warning: Failed to list context ids from Redis: {e}")
        return context_ids

    def _delete_context_redis_keys(self, context_id: str) -> bool:
        try:
            client = self._redis_client()
            raw_keys = cast(Any, client.keys(f"context:{context_id}:*"))
            keys = cast(List[str], raw_keys)
            if not keys:
                return False
            deleted = cast(Any, client.delete(*keys))
            return int(deleted) > 0
        except Exception as e:
            print(f"Warning: Failed to delete Redis context keys for {context_id}: {e}")
            return False

    def create_context(
        self,
        context_id: Optional[str] = None,
        llm_interface: Optional[
            OpenAICompatible
        ] = get_config().CONTEXT_SUMMARY_INTERFACE,
        max_history_length: int = get_config().CONTEXT_MAX_HISTORY_LENGTH,
        auto_summarize_trigger: int = get_config().CONTEXT_AUTO_SUMMARIZE_TRIGGER,
        **backend_kwargs,
    ) -> ContextBackend:
        """
        Create a new context object.

        Args:
            context_id: Context ID; generated automatically if None
            llm_interface: LLM interface
            max_history_length: Maximum history length
            auto_summarize_trigger: Automatic summary trigger threshold
            **backend_kwargs: Extra parameters passed to the backend

        Returns:
            ContextBackend: Created context object
        """
        with self._lock:
            if context_id is None:
                context_id = str(uuid.uuid4())

            # Check whether it already exists.
            if context_id in self._active_contexts:
                app_log(
                    f"Context {context_id} already exists, and is in active contexts. Returning the existing context."
                )
                return self._active_contexts[context_id]

            # Generate the file path if the backend needs one.
            if "file_path" not in backend_kwargs:
                context_file = os.path.join(self.context_dir, f"ctx_{context_id}.json")
                backend_kwargs["file_path"] = context_file
                push_warning(f"Context file path: {context_file}")

            # Create the context object.
            context = self.backend_class(
                context_id=context_id,
                llm_interface=llm_interface,
                max_history_length=max_history_length,
                auto_summarize_trigger=auto_summarize_trigger,
                **backend_kwargs,
            )

            # Add it to the active context list.
            self._active_contexts[context_id] = context

            return context

    def get_context(self, context_id: str) -> Optional[ContextBackend]:
        """
        Get the context object with the specified ID.

        Args:
            context_id: Context ID

        Returns:
            ContextBackend: Context object, or None if it does not exist
        """
        with self._lock:
            # First check active contexts.
            if context_id in self._active_contexts:
                return self._active_contexts[context_id]

            # Try to load from file if the backend supports it.
            context_file = os.path.join(self.context_dir, f"ctx_{context_id}.json")
            if os.path.exists(context_file):
                try:
                    context = self.backend_class(
                        context_id=context_id,
                        llm_interface=self.config.CONTEXT_SUMMARY_INTERFACE,  # Can be configured later.
                        max_history_length=self.config.CONTEXT_MAX_HISTORY_LENGTH,
                        auto_summarize_trigger=self.config.CONTEXT_AUTO_SUMMARIZE_TRIGGER,
                        file_path=context_file,
                    )
                    self._active_contexts[context_id] = context
                    return context
                except Exception as e:
                    print(f"Warning: Failed to load context {context_id}: {e}")

            return None

    def delete_context(self, context_id: str) -> bool:
        """
        Delete the context object with the specified ID.

        Args:
            context_id: Context ID

        Returns:
            bool: Whether deletion succeeded
        """
        with self._lock:
            success = False

            # Remove it from active contexts.
            if context_id in self._active_contexts:
                del self._active_contexts[context_id]
                success = True

            # Delete context keys from Redis.
            if self._delete_context_redis_keys(context_id):
                success = True

            # Delete the file if it exists.
            context_file = os.path.join(self.context_dir, f"ctx_{context_id}.json")
            if os.path.exists(context_file):
                try:
                    os.remove(context_file)
                    success = True
                except Exception as e:
                    print(f"Warning: Failed to delete context file {context_file}: {e}")

            return success

    def list_context_ids(self) -> List[str]:
        """List all known context IDs, including Redis and the file system."""
        context_ids = set(self._active_contexts.keys())
        context_ids.update(self._list_context_ids_from_redis())

        try:
            for filename in os.listdir(self.context_dir):
                if filename.startswith("ctx_") and filename.endswith(".json"):
                    context_ids.add(filename[4:-5])
        except Exception as e:
            print(f"Warning: Failed to scan context dir for ids: {e}")

        return sorted(context_ids)

    def list_contexts(self) -> List[Dict[str, Any]]:
        """
        List all available contexts.

        Returns:
            List[Dict]: Context information list
        """
        contexts = []

        # Scan context files in the file system.
        try:
            for filename in os.listdir(self.context_dir):
                if filename.startswith("ctx_") and filename.endswith(".json"):
                    context_id = filename[4:-5]  # Remove the "ctx_" prefix and ".json" suffix.

                    context_info = {
                        "context_id": context_id,
                        "file_path": os.path.join(self.context_dir, filename),
                        "is_active": context_id in self._active_contexts,
                    }

                    # Try to read basic information.
                    try:
                        file_path = context_info["file_path"]
                        if isinstance(file_path, str):
                            with open(file_path, "r", encoding="utf-8") as f:
                                data = json.load(f)
                                metadata = data.get("metadata", {})
                                context_info.update(
                                    {
                                        "start_time": metadata.get("start_time"),
                                        "last_activity": metadata.get("last_activity"),
                                        "total_messages": metadata.get(
                                            "total_messages", 0
                                        ),
                                    }
                                )
                    except Exception:
                        pass  # Ignore read errors.

                    contexts.append(context_info)

        except Exception as e:
            print(f"Warning: Failed to list contexts: {e}")

        return contexts

    async def save_context(self, context_id: str) -> bool:
        """
        Manually save the specified context to file.

        Args:
            context_id: Context ID

        Returns:
            bool: Whether saving succeeded
        """
        with self._lock:
            context = self._active_contexts.get(context_id)

        if context is None:
            return False

        try:
            return await context.persist()
        except Exception as e:
            print(f"Warning: Failed to save context {context_id}: {e}")
            return False

    async def save_all_contexts(self) -> int:
        """
        Save all active contexts to files.

        Returns:
            int: Number of contexts successfully saved
        """
        saved_count = 0
        with self._lock:
            context_ids = list(self._active_contexts.keys())

        for context_id in context_ids:
            if await self.save_context(context_id):
                saved_count += 1

        return saved_count

    async def cleanup_inactive_contexts(self, max_inactive_time: int = 3600) -> int:
        """
        Clean up contexts that have been inactive for a long time.

        Args:
            max_inactive_time: Maximum inactive time in seconds

        Returns:
            int: Number of cleaned contexts
        """
        cleaned_count = 0
        current_time = datetime.now()
        contexts_to_persist: list[tuple[str, ContextBackend]] = []

        with self._lock:
            contexts_to_remove = []

            for context_id, context in self._active_contexts.items():
                try:
                    metadata = context.get_metadata()
                    last_activity_str = metadata.get("last_activity")
                    if last_activity_str:
                        last_activity = datetime.fromisoformat(last_activity_str)
                        inactive_time = (current_time - last_activity).total_seconds()

                        if inactive_time > max_inactive_time:
                            # Persist the context before removing it.
                            contexts_to_persist.append((context_id, context))
                            contexts_to_remove.append(context_id)
                            cleaned_count += 1
                except Exception as e:
                    print(
                        f"Warning: Error checking activity for context {context_id}: {e}"
                    )

        for context_id, context in contexts_to_persist:
            try:
                await context.persist()
            except Exception as e:
                print(f"Warning: Failed to persist inactive context {context_id}: {e}")

        with self._lock:
            for context_id in contexts_to_remove:
                del self._active_contexts[context_id]

        return cleaned_count

    # ===== Convenience interfaces =====

    async def add_message(
        self,
        context_id: str,
        message: Message,
    ) -> bool:
        """
        Convenience method for adding a message.

        Args:
            context_id: Context ID
            role: Message role
            content: Message content
            **message_kwargs: Other message parameters

        Returns:
            bool: Whether adding succeeded
        """
        backend = self.get_context(context_id)
        if not backend:
            return False

        try:
            await backend.store_message(message)
            return True
        except Exception as e:
            print(f"Warning: Failed to add message: {e}")
            return False

    def get_history(self, context_id: str, limit: Optional[int] = None) -> List:
        """
        Get conversation history.

        Args:
            context_id: Context ID
            limit: Limit on the number of returned messages

        Returns:
            List: Message history
        """
        backend = self.get_context(context_id)
        if not backend:
            return []

        return backend.retrieve_messages(limit)

    async def summarize_context(self, context_id: str) -> Optional[str]:
        """
        Summarize the context.

        Args:
            context_id: Context ID

        Returns:
            Optional[str]: Summary content
        """
        backend = self.get_context(context_id)
        if not backend:
            return None

        return await backend.auto_summarize()


# Global instance.
_global_context_manager: Optional[ContextManager] = None


def get_context_manager() -> ContextManager:
    """Get the global ContextManager instance."""
    global _global_context_manager
    if _global_context_manager is None:
        # Get Redis configuration from config.
        config = get_config()

        # redis config
        redis_host = config.REDIS_HOST
        redis_port = int(config.REDIS_PORT)
        redis_db = int(config.REDIS_DB)

        # Create a custom backend class with preconfigured Redis parameters.
        class ConfiguredRedisFileBackend(RedisFileContextBackend):
            def __init__(
                self,
                context_id: str,
                llm_interface: Optional[
                    OpenAICompatible
                ] = config.CONTEXT_SUMMARY_INTERFACE,
                max_history_length: int = config.CONTEXT_MAX_HISTORY_LENGTH,
                auto_summarize_trigger: int = config.CONTEXT_AUTO_SUMMARIZE_TRIGGER,
                redis_host: str = redis_host,
                redis_port: int = redis_port,
                redis_db: int = redis_db,
                file_path: str = "",
            ):
                super().__init__(
                    context_id=context_id,
                    llm_interface=llm_interface,
                    max_history_length=max_history_length,
                    auto_summarize_trigger=auto_summarize_trigger,
                    redis_host=redis_host,
                    redis_port=redis_port,
                    redis_db=redis_db,
                    file_path=file_path,
                )
                self.file_path = file_path
                self.context_id = context_id
                self.llm_interface = llm_interface
                self.max_history_length = max_history_length
                self.auto_summarize_trigger = auto_summarize_trigger
                self.redis_host = redis_host
                self.redis_port = redis_port
                self.redis_db = redis_db

        # Create ContextManager using the configured backend class.
        _global_context_manager = ContextManager(
            backend_class=ConfiguredRedisFileBackend
        )
    return _global_context_manager