from typing import List, Optional, Union

from langchain_openai import AzureOpenAIEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client.models import Distance, PointStruct, VectorParams
from ai_workspace.database import MongoDB, Qdrant

from ai_workspace.schemas import WorkspaceSchema, DocumentSchema
from ai_workspace.exceptions import WorkspaceError, handle_exceptions
from .base import BaseWorkspace
from ai_workspace.utils.logger import setup_logger
from typing import ClassVar
from logging import Logger
from ai_workspace.packages.shared import IDGenerator
from ai_workspace.packages.document import Document
from ai_workspace.packages.embed import EmbedDocument


class Workspace(BaseWorkspace):
    """A class to manage vector document storage and retrieval using Qdrant and Azure OpenAI.

    This class provides functionality to store, embed, and retrieve documents using
    Qdrant vector database and Azure OpenAI embeddings.
    """

    mongodb: MongoDB
    workspace_collection: str = "workspaces"
    qdrant: Qdrant
    embedding: AzureOpenAIEmbeddings

    logger: ClassVar[Logger] = setup_logger(__name__)
    id_generator: ClassVar[IDGenerator] = IDGenerator()
    vector_config: VectorParams = VectorParams(
        size=3072,  # Size of the embedding vector
        distance=Distance.COSINE,  # Using cosine distance for similarity
    )

    @property
    def document_client(self) -> Document:
        """Get the Document client for managing documents in the workspace."""
        return Document(mongodb=self.mongodb)

    @property
    def embed_client(self) -> EmbedDocument:
        """Get the EmbedDocument client for managing embeddings in the workspace."""
        return EmbedDocument(embedding=self.embedding)

    @property
    def collection(self):
        """Get the collection client for the workspace collection."""
        return self.mongodb.get_db().get_collection(self.workspace_collection)

    @property
    def qdrant_client(self):
        """Get the Qdrant client."""
        return self.qdrant.get_client()

    def get_vector_store(self, workspace_id: str) -> QdrantVectorStore:
        """Get the vector store for a specific workspace.

        Args:
            workspace_id (str): The ID of the workspace

        Returns:
            QdrantVectorStore: The vector store for the specified workspace
        """
        return QdrantVectorStore(
            qdrant_client=self.qdrant_client,
            collection_name=workspace_id,
            embedding=self.embed_client.embedding,
        )

    @handle_exceptions
    def get_instructions(self, workspace_id: str) -> str:
        """Get instructions from database schema.

        Args:
            workspace_id (str): The ID of the workspace

        Returns:
            str: The instructions for the workspace

        Raises:
            WorkspaceError: If workspace is not found or has no instructions
        """
        workspace = self.collection.find_one({"workspace_id": workspace_id})
        if not workspace:
            raise WorkspaceError(
                f"Workspace {workspace_id} not found or has no instructions"
            )

        if "instructions" not in workspace:
            raise WorkspaceError(f"Workspace {workspace_id} has no instructions")

        return workspace["instructions"]

    @handle_exceptions
    def create(self, workspace_data: WorkspaceSchema) -> Optional[str]:
        """Add a new workspace to MongoDB.

        Args:
            workspace_data (WorkspaceSchema): Data for the workspace to add

        Returns:
            str: ID of the newly added workspace
        """
        workspace_data.workspace_id = (
            str(next(self.id_generator))
            if not workspace_data.workspace_id
            else workspace_data.workspace_id
        )
        workspace_dict = workspace_data.model_dump(exclude_unset=True)

        self.collection.insert_one(workspace_dict)

        # Create a Qdrant collection for the workspace
        self.qdrant_client.create_collection(
            collection_name=workspace_data.workspace_id,
            vectors_config=self.vector_config,
        )
        return workspace_data.workspace_id

    def add_session(self, workspace_id: str, session_id: str) -> bool:
        """Add a session to a workspace.

        Args:
            workspace_id (str): ID of the workspace to add the session to
            session_id (str): ID of the session to add
        Returns:
            bool: True if the session was added successfully, False otherwise
        """
        try:
            self.collection.find_one_and_update(
                {"workspace_id": workspace_id},
                {"$addToSet": {"chat_sessions": session_id}},
            )
            return True
        except Exception as e:
            self.logger.error(
                f"Cannot add session {session_id} to workspace {workspace_id}: {e}"
            )
            return False

    def delete_workspace(self, workspace_id: str) -> bool:
        """Delete a workspace from MongoDB.

        Args:
            workspace_id (str): ID of the workspace to delete
            This will also delete the corresponding Qdrant collection.

        Returns:
            bool: True if deletion was successful, False otherwise
        """
        try:
            result = self.collection.delete_one({"workspace_id": workspace_id})

            if result.deleted_count > 0:
                self.qdrant_client.delete_collection(workspace_id)
                return True
            else:
                self.logger.warning(f"No workspace found with ID {workspace_id}")
                return False

        except Exception as e:
            self.logger.error(f"ERROR in delete_workspace: {e}")
            return False

    def update_instructions(self, workspace_id: str, instructions: str) -> bool:
        """Update instructions for a workspace.

        Args:
            workspace_id (str): ID of the workspace to update
            instructions (str): New instructions for the workspace

        Returns:
            bool: True if update was successful, False otherwise
        """
        try:
            result = self.collection.update_one(
                {"workspace_id": workspace_id},
                {"$set": {"instructions": instructions}},
            )

            if result.modified_count > 0:
                return True
            else:
                self.logger.warning(f"No workspace found with ID {workspace_id}")
                return False

        except Exception as e:
            self.logger.error(f"ERROR in update_instructions: {e}")
            return False

    def delete_instructions(self, workspace_id: str) -> bool:
        """Delete instructions for a workspace.

        Args:
            workspace_id (str): ID of the workspace to delete instructions for

        Returns:
            bool: True if deletion was successful, False otherwise
        """
        try:
            result = self.collection.update_one(
                {"workspace_id": workspace_id},
                {"$unset": {"instructions": ""}},
            )

            if result.modified_count > 0:
                return True
            else:
                self.logger.warning(f"No workspace found with ID {workspace_id}")
                return False

        except Exception as e:
            self.logger.error(f"ERROR in delete_instructions: {e}")
            return False

    def find_workspace_by_id(self, workspace_id: str) -> Optional[WorkspaceSchema]:
        """Find a workspace by its ID.

        Args:
            workspace_id (str): ID of the workspace to find

        Returns:
            Optional[WorkspaceSchema]: The found workspace or None if not found
        """
        workspace = self.collection.find_one({"workspace_id": workspace_id})
        if not workspace:
            return None
        return WorkspaceSchema.model_validate(workspace)

    def find_workspace_by_user_id(
        self, user_id: str
    ) -> Optional[List[WorkspaceSchema]]:
        """Find a workspace by its user ID.

        Args:
            user_id (str): ID of the user to find the workspace for

        Returns:
            Optional[WorkspaceSchema]: The found workspace or None if not found
        """
        workspace = self.collection.find({"user_id": user_id})
        if not workspace:
            return None
        return [WorkspaceSchema.model_validate(ws) for ws in workspace]

    @handle_exceptions
    def add_knowledge(
        self,
        document: DocumentSchema,
        raw_texts: Optional[List[str]] = None,
        metadata: Optional[dict | List[dict]] = None,
        to_workspace: bool = True,
    ) -> str:
        """Add a document to the workspace.
        Saves the document to MongoDB and returns its ID.
        Upload the embedded document to the Qdrant collection.
        If `to_workspace` is True, the document will be added to the workspace's collection else it will be added to session collection.
        Args:
            document (DocumentSchema): The document to upload
            raw_texts (Optional[List[str]]): List of raw texts to embed and store
            metadata (Optional[dict | List[dict]]): Metadata for the document.
                If a single dict is provided, it will be applied to all texts.
                If a list of dicts is provided, each text will get its corresponding metadata.

        Returns:
            str: The ID of the uploaded document
        """
        document_id = self.document_client.upload_document(document)
        if not to_workspace and not document.session_id:
            raise WorkspaceError(
                "Cannot add document to session without a session_id in the document"
            )
        if not to_workspace and document.session_id:
            to_collection = document.session_id
        else:
            to_collection = document.workspace_id

        if metadata and isinstance(metadata, dict):
            metadata = [metadata]
        if raw_texts:
            _vectors = self.embed_client.embed_texts(raw_texts)
            self.qdrant_client.upsert(
                collection_name=to_collection,
                points=[
                    PointStruct(
                        id=next(self.id_generator),
                        vector=_vector,
                        payload={
                            "page_content": raw_texts[i],
                            "metadata": {
                                "file_name": document.file_name,
                                "file_mime": document.file_mime,
                                "file_suffix": document.file_suffix,
                                "file_url": document.file_url,
                                "document_id": document_id,
                                "created_at": document.created_at,
                                "updated_at": document.updated_at,
                                **(metadata[i] if metadata else {}),
                            },
                        },
                    )
                    for i, _vector in enumerate(_vectors)
                ],
            )
        return document_id
