import os
import json
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import logging

import dotenv
dotenv.load_dotenv()

from sqlalchemy import Column, String, Text, DateTime, ForeignKey, create_engine, Index
from sqlalchemy.dialects.sqlite import JSON
from sqlalchemy.orm import relationship, sessionmaker
from sqlalchemy.exc import SQLAlchemyError

from qdrant_client import QdrantClient, models
from sqlalchemy.orm import declarative_base

from dalistore.memory_bank import MemoryBank

# Configure logging
logging.basicConfig(level=logging.ERROR)
logger = logging.getLogger(__name__)

Base = declarative_base()


class MemoryUnit(Base):
    __tablename__ = "memory_units"
    unique_id = Column(String, primary_key=True)
    timestamp = Column(DateTime, nullable=False)

    # Inputs
    type = Column(String, nullable=False)
    description = Column(Text)
    metadata_ = Column("metadata", JSON)
    state_before = Column(Text)

    # Outputs
    status = Column(String)
    human_comment = Column(Text)
    ai_comment = Column(Text)
    state_after = Column(Text)

    # Relationships
    parent_id = Column(String, ForeignKey("memory_units.unique_id"))
    parent = relationship(
        "MemoryUnit",
        remote_side=[unique_id],
        backref="children",
        foreign_keys=[parent_id],
    )

    # Indexes
    __table_args__ = (
        Index("ix_memory_units_timestamp", "timestamp"),
        Index("ix_memory_units_type", "type"),
        Index("ix_memory_units_status", "status"),
        Index("ix_memory_units_parent_id", "parent_id"),
    )


class TextMemoryBank(MemoryBank):
    """
    TextMemoryBank handles storage and retrieval of memory units containing textual data.
    It uses SQLAlchemy ORM for structured data storage and Qdrant for vector embeddings to enable similarity searches.
    """

    def __init__(self, remote: bool = False, home_dir: str = None, collection_name: str = "text_memories", db_name: str = "text_memories"):
        super().__init__(remote, home_dir, collection_name, db_name)
        self.model = self.model_manager.load_text_model("all-MiniLM-L6-v2")
        self.engine = None
        self.Session = None
        self.initialize()

    def initialize(self):
        """Initialize the SQLAlchemy engine and session."""
        try:
            os.makedirs(self.home_dir, exist_ok=True)
            self._initialize_sqlalchemy()
            self._initialize_qdrant()
            self.initialized = True
        except Exception as e:
            logger.error(f"Error initializing TextMemoryBank: {e}")
            raise e

    def store(self, data: Dict[str, Any]) -> str:
        """
        Store the memory unit data in the database and embeddings in Qdrant.

        Parameters:
        - data: Dictionary containing the memory unit data.

        Returns:
        - The unique_id of the stored memory unit.
        """
        self._ensure_initialized()
        if "unique_id" not in data:
            unique_id = self._generate_unique_id()
        else:
            unique_id = data["unique_id"]

        try:
            # Create a new SQLAlchemy session
            with self.Session() as session:
                memory_unit = MemoryUnit(
                    unique_id=unique_id,
                    timestamp=data.get("timestamp", datetime.utcnow()),
                    type=data["type"],
                    description=data.get("description"),
                    status=data.get("status"),
                    parent_id=data.get("parent_id"),
                    metadata_=data.get("metadata"),
                    state_before=data.get("state_before"),
                    state_after=data.get("state_after"),
                    human_comment=data.get("human_comment"),
                    ai_comment=data.get("ai_comment"),
                )
                session.add(memory_unit)
                session.commit()
            return unique_id
        except SQLAlchemyError as e:
            logger.error(f"Error storing memory unit: {e}")
            raise e

    def retrieve_by_id(
        self, unique_id: str, depth: int = 1
    ) -> Optional[Dict[str, Any]]:
        """
        Retrieve the memory unit by unique ID from the database.

        Parameters:
        - unique_id: The unique identifier of the memory unit.
        - depth: The depth of child units to retrieve.

        Returns:
        - A dictionary representing the memory unit, or None if not found.
        """
        return self._fetch_unit(unit_id=unique_id, current_depth=0, depth=depth)
    
    def retrieve_by_filter(self, filter_data: Dict[str, Any], max_results: int = 10) -> List[str]:
        """
        Retrieve memory units by filter.
        """
        self._ensure_initialized()
        try:
            with self.Session() as session:
                # Start building the query
                query = session.query(MemoryUnit)
                
                # Apply filters from filter_data
                for key, value in filter_data.items():
                    if hasattr(MemoryUnit, key):
                        query = query.filter(getattr(MemoryUnit, key) == value)
                
                # Order by timestamp descending and limit results
                results = query.order_by(MemoryUnit.timestamp.desc()).limit(max_results).all()
                
                # Return list of tuples with id and score 1.0 (since this is exact match)
                return [result.unique_id for result in results]
                
        except SQLAlchemyError as e:
            logger.error(f"Error retrieving memory units by filter: {e}")
            raise e

    def embed(self, unique_id: str, data: Dict[str, Any]) -> None:
        """
        Generate embeddings for the memory unit and store them in Qdrant.

        Parameters:
        - unique_id: Unique identifier for the memory unit.
        - data: Dictionary containing the memory unit data.
        """
        self._ensure_initialized()
        try:
            vectors = {}
            input_embeddings_list = []
            output_embeddings_list = []
            all_embeddings_list = []
            input_fields = ["type", "description", "metadata", "state_before"]
            output_fields = ["status", "human_comment", "ai_comment", "state_after"]
            all_fields = input_fields + output_fields

            for field in all_fields:
                content = data.get(field)
                if content:
                    if field == "metadata":
                        # Convert metadata dict to string if necessary
                        if isinstance(content, dict):
                            content = json.dumps(content)
                    embedding = self._embed_text(content)
                    vectors[field] = embedding.tolist()
                    if field in input_fields:
                        input_embeddings_list.append(embedding)
                    else:
                        output_embeddings_list.append(embedding)
                    all_embeddings_list.append(embedding)

            # Calculate the averaged embeddings
            if input_embeddings_list:
                input_embedding = np.mean(input_embeddings_list, axis=0)
                vectors["inputs"] = input_embedding.tolist()
            if output_embeddings_list:
                output_embedding = np.mean(output_embeddings_list, axis=0)
                vectors["outputs"] = output_embedding.tolist()
            if all_embeddings_list:
                omnibus_embedding = np.mean(all_embeddings_list, axis=0)
                vectors["all"] = omnibus_embedding.tolist()

            # Upsert embeddings into Qdrant with named vectors
            point = models.PointStruct(
                id=unique_id,
                vector=vectors,  # Pass the dictionary of named vectors
                payload={"unique_id": unique_id, "type": data.get("type")},
            )

            self.qdrant_client.upsert(
                collection_name=self.collection_name, points=[point]
            )

        except Exception as e:
            logger.error(f"Error storing embeddings in Qdrant: {e}")
            raise e

    def retrieve_similar(
        self,
        query_data: Dict[str, Any],
        query_weights: Dict[str, float] = {},
        filter_data: Dict[str, Any] = {},
        max_results: int = 10,
        score_threshold: float = 0.5,
    ) -> List[Tuple[str, float]]:
        """
        Retrieve IDs of similar memory units based on the provided query, along with their scores.

        Parameters:
        - query_data: A dictionary containing the query to use for similarity search.
        - query_weights: A dictionary containing the weights to use for similarity search.
        - filter_data: A dictionary containing the data to filter the results by.
        - max_results: The maximum number of results to return.
        - score_threshold: The minimum score threshold for the results.

        Returns:
        - A list of tuples representing similar memory units and their scores.
        """
        self._ensure_initialized()
        # collect all filter_data into conditions
        conditions = []
        for field, value in filter_data.items():
            conditions.append(
                models.FieldCondition(key=field, match=models.MatchValue(value=value))
            )
        if conditions:
            filter = models.Filter(must=conditions)
        else:
            filter = None

        # make independent query for each field in query_data with max_results * len(query_data)
        results = {}
        factor = len(query_data)
        for field, value in query_data.items():
            query_embedding = self._embed_text(value)
            query_results = self.qdrant_client.search(
                collection_name=self.collection_name,
                query_vector=(field, query_embedding),
                limit=max_results * factor,
                score_threshold=score_threshold,
                query_filter=filter,
            )
            for hit in query_results:
                id = hit.payload["unique_id"]
                score = hit.score
                if id not in results:
                    results[id] = {}
                results[id][field] = score

        # calculate the total weight of the query
        total_weight = 0
        for field in query_data:
            total_weight += query_weights[field] if field in query_weights else 1

        # combine results with a simple weighted average
        combined_results = []
        for id, scores in results.items():
            if len(scores) == len(query_data):
                combined_score = 0
                for field, score in scores.items():
                    combined_score += score * (
                        query_weights[field] if field in query_weights else 1
                    )
                combined_results.append((id, combined_score / total_weight))

        # return the top N sorted combined results
        filtered_results = [
            (id, score) for id, score in combined_results if score > score_threshold
        ]
        sorted_results = sorted(filtered_results, key=lambda x: x[1], reverse=True)

        return sorted_results[:max_results]

    def delete(self, unique_id: str, delete_from_qdrant: bool = True) -> None:
        """
        Delete the memory unit with the specified unique ID from both the database and Qdrant.

        Parameters:
        - unique_id: The unique identifier of the memory unit to delete.
        - delete_from_qdrant: Whether to delete the embeddings from Qdrant.
        """
        self._ensure_initialized()
        try:
            with self.Session() as session:
                memory_unit = (
                    session.query(MemoryUnit)
                    .filter_by(unique_id=unique_id)
                    .one_or_none()
                )
                if memory_unit:
                    session.delete(memory_unit)
                    session.commit()
                else:
                    logger.warning(
                        f"Memory unit with ID {unique_id} not found in database."
                    )
            if delete_from_qdrant:
                self.qdrant_client.delete(
                    collection_name=self.collection_name, points_selector=[unique_id]
                )
        except SQLAlchemyError as e:
            logger.error(f"Error deleting memory unit '{unique_id}': {e}")
            raise e

    def edit(self, unique_id: str, new_data: Dict[str, Any]) -> None:
        """
        Edit the memory unit with the specified unique ID using the provided data.

        Parameters:
        - unique_id: The unique identifier of the memory unit to edit.
        - data: A dictionary containing the updated data.
        """
        self._ensure_initialized()
        session = self.Session()
        try:
            memory_unit = (
                session.query(MemoryUnit).filter_by(unique_id=unique_id).one_or_none()
            )
            if memory_unit is None:
                raise ValueError(f"Memory unit with unique_id {unique_id} not found.")

            # Map input fields to ORM fields
            field_mapping = {
                "description": "description",
                "status": "status",
                "human_comment": "human_comment",
                "ai_comment": "ai_comment",
                "metadata": "metadata_",
                "type": "type",
            }

            for key, value in new_data.items():
                if key in field_mapping:
                    setattr(memory_unit, field_mapping[key], value)
                elif hasattr(memory_unit, key):
                    setattr(memory_unit, key, value)
                else:
                    logger.warning(f"Unknown field '{key}' in edit data.")

            # Update embeddings if necessary
            fields_to_embed = [
                "description",
                "human_comment",
                "ai_comment",
                "metadata",
                "status",
            ]
            embedding_vectors = {}
            for field in fields_to_embed:
                if field in new_data:
                    embedding_vectors[field] = self._embed_text(new_data[field])

            if embedding_vectors:
                # Update embeddings in Qdrant
                self.qdrant_client.upsert(
                    collection_name=self.collection_name,
                    points=[
                        models.PointStruct(
                            id=unique_id, vector=embedding_vectors, payload={}
                        )
                    ],
                )

            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"Error editing memory unit with ID {unique_id}: {e}")
            raise
        finally:
            session.close()

    def close(self):
        """Close the database connection and other resources."""
        if self.engine:
            self.engine.dispose()
        if self.qdrant_client:
            self.qdrant_client.close()

    def _fetch_unit(
        self,
        unit_id: str,
        current_depth: int,
        depth: int,
        visited: Optional[set] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Fetches a memory unit recursively up to the specified depth.

        Parameters:
        - unit_id: The unique identifier of the memory unit.
        - current_depth: The current depth of recursion.
        - depth: The maximum depth to recurse.
        - visited: A set of visited unit_ids to prevent cycles.

        Returns:
        - A dictionary representing the memory unit, or None if not found.

        Raises:
        - Exception: If an error occurs during database operations.
        """
        self._ensure_initialized()

        if visited is None:
            visited = set()

        if unit_id in visited:
            logger.warning(
                f"Cycle detected while fetching unit {unit_id}. Skipping to prevent infinite recursion."
            )
            return None

        visited.add(unit_id)

        try:
            with self.Session() as session:
                # Fetch the unit
                memory_unit = (
                    session.query(MemoryUnit).filter_by(unique_id=unit_id).one_or_none()
                )
                if not memory_unit:
                    logger.warning(f"Memory unit with ID {unit_id} not found.")
                    return None

                # Convert SQLAlchemy object to dictionary
                unit = memory_unit.__dict__.copy()
                unit.pop("_sa_instance_state", None)

                # If using 'metadata_' attribute, map it back to 'metadata'
                if "metadata_" in unit:
                    unit["metadata"] = unit.pop("metadata_", None)

                # Fetch media blobs if supported
                if self.supports_media_blobs:
                    unit["media_blobs"] = []
                    for media_blob in memory_unit.media_blobs:
                        blob_dict = media_blob.__dict__.copy()
                        blob_dict.pop("_sa_instance_state", None)
                        unit["media_blobs"].append(blob_dict)

                # Fetch descendants if depth allows
                if current_depth < depth:
                    unit["subunits"] = []
                    child_units = (
                        session.query(MemoryUnit).filter_by(parent_id=unit_id).all()
                    )
                    for child_unit in child_units:
                        try:
                            child_data = self._fetch_unit(
                                child_unit.unique_id, current_depth + 1, depth, visited
                            )
                            if child_data:
                                unit["subunits"].append(child_data)
                        except Exception as e:
                            logger.error(
                                f"Error fetching child unit with ID {child_unit.unique_id}: {e}"
                            )
                            continue  # Skip this child and continue with others
                return unit
        except Exception as e:
            logger.error(f"Error fetching unit with ID {unit_id}: {e}")
            raise  # Re-raise the exception after logging

    def _generate_unique_id(self) -> str:
        return str(uuid.uuid4())

    def _embed_text(self, text: str) -> np.ndarray:
        return self.model.encode([text])[0]

    def _initialize_qdrant(self):
        # Initialize qdrant
        if self.qdrant_client is None:
            if not self.remote:
                self.qdrant_client = QdrantClient(path=self.qdrant_path)
            else:
                self.qdrant_client = QdrantClient(url=os.getenv("QDRANT_URL"), api_key=os.getenv("QDRANT_API_KEY"))
        # Get embedding size dynamically
        embedding_size = self.model_manager.get_text_embedding_size()

        # Define vector configurations for each named vector
        vector_configs = {
            # Inputs
            "type": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "description": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "metadata": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "state_before": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            # Outputs
            "status": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "human_comment": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "ai_comment": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "state_after": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            # Combined embeddings
            "inputs": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "outputs": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
            "all": models.VectorParams(
                size=embedding_size, distance=models.Distance.COSINE
            ),
        }

        if not self.qdrant_client.collection_exists(self.collection_name):
            # Create collection with multiple vectors
            self.qdrant_client.create_collection(
                collection_name=self.collection_name, vectors_config=vector_configs
            )

    def _initialize_sqlalchemy(self):
        """
        Initialize the SQLAlchemy engine and session factory.
        """
        if self.remote:
            username = os.getenv("POSTGRES_USER")
            password = os.getenv("POSTGRES_PASSWORD")
            host = os.getenv("POSTGRES_HOST")  
            port = os.getenv("POSTGRES_PORT")
            database = os.getenv("POSTGRES_DB")
            url = f"postgresql://{username}:{password}@{host}:{port}/{database}"
            self.engine = create_engine(url)
        else:
            self.engine = create_engine(
                f"sqlite:///{self.db_path}", connect_args={"check_same_thread": False}
            )

        def set_test_table_name(base, suffix):
            for cls in base.__subclasses__():
                cls.__tablename__ = f"{cls.__tablename__}_{suffix}"

        set_test_table_name(Base, self.db_name)

        self.Session = sessionmaker(bind=self.engine)
        Base.metadata.create_all(self.engine)

        if not self.remote:
            logger.debug(f"Database path: {self.db_path}")
            logger.debug(f"Database file permissions: {oct(os.stat(self.db_path).st_mode)}")
