from datetime import datetime, timezone
import json
import uuid
from enum import Enum
from typing import (
    IO,
    Any,
    Iterator,
    List,
    Literal,
    Tuple,
    Type,
    TypeVar,
    Optional,
    overload,
)
from .types import IDatabase
from .locks import LockManager
from .vectors import NumpyVectorIndex as VectorIndex
from .manager import ManagerBase, synced, emits

import numpy as np

from pydantic import BaseModel, Field

# --- Fuzzy Search Helper Functions ---


def _levenshtein_distance(s1: str, s2: str) -> int:
    """Calculates the Levenshtein distance between two strings."""
    if len(s1) < len(s2):
        return _levenshtein_distance(s2, s1)
    if len(s2) == 0:
        return len(s1)

    previous_row = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)
            current_row.append(min(insertions, deletions, substitutions))
        previous_row = current_row
    return previous_row[-1]


def _get_trigrams(text: str) -> set[str]:
    """Generates a set of 3-character trigrams from a string."""
    if not text or len(text) < 3:
        return set()
    return {text[i : i + 3] for i in range(len(text) - 2)}


def _sliding_window_levenshtein(query: str, content: str, fuzziness: int) -> int:
    """
    Finds the best Levenshtein match for a query within a larger text
    by comparing it against relevant substrings.
    """
    query_tokens = query.lower().split()
    content_tokens = content.lower().split()
    query_len = len(query_tokens)
    if query_len == 0:
        return 0

    min_dist = float("inf")
    query_norm = " ".join(query_tokens)

    # The window size can be slightly smaller or larger than the query length
    # to account for missing or extra words in a fuzzy match.
    for window_size in range(max(1, query_len - fuzziness), query_len + fuzziness + 1):
        if window_size > len(content_tokens):
            continue
        for i in range(len(content_tokens) - window_size + 1):
            window_text = " ".join(content_tokens[i : i + window_size])
            dist = _levenshtein_distance(query_norm, window_text)
            if dist < min_dist:
                min_dist = dist

    return int(min_dist)


class Document[B](BaseModel):
    """A data class representing a single item in a collection."""

    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    embedding: list[float] | None = None
    body: B | None = None

    def serialize_embedding(self) -> bytes | None:
        """Serializes the embedding to bytes for storage."""
        if self.embedding is None:
            return None

        array = np.array(self.embedding, dtype=np.float32)
        return array.tobytes()

    @classmethod
    def deserialize_embedding(cls, data: bytes | None):
        """Deserializes bytes back into the embedding list."""
        if data is None:
            return None

        array = np.frombuffer(data, dtype=np.float32)
        return array.tolist()

    def _flatten_metadata(self, metadata: dict, prefix: str = "") -> dict[str, Any]:
        """Flattens a nested dictionary for indexing."""
        flat_dict = {}
        for key, value in metadata.items():
            new_key = f"{prefix}.{key}" if prefix else key
            if isinstance(value, dict):
                flat_dict.update(self._flatten_metadata(value, new_key))
            else:
                flat_dict[new_key] = value
        return flat_dict

    def serialize_body(self) -> dict:
        """Serializes the body to a dictionary for storage."""
        if self.body is None:
            return {}

        if isinstance(self.body, dict):
            return self.body
        elif isinstance(self.body, BaseModel):
            return self.body.model_dump(mode="json")
        else:
            return {"value": self.body}

    def flattened_body(self) -> dict[str, Any]:
        """Returns the flattened body metadata."""
        return self._flatten_metadata(self.serialize_body())


class CollectionManager[B: BaseModel](ManagerBase[B]):
    """
    A wrapper for multi-modal collection operations, including document storage,
    FTS, fuzzy search, graph traversal, and persistent vector search.
    """

    def __init__(self, name: str, db: IDatabase, model: Type[B] | None = None):
        super().__init__(name, db, model)

        # All vector-related operations are now delegated to the VectorIndex class.
        self._vector_index = VectorIndex(name, db)

        self._compact_lock = LockManager(
            db, f"__lock__internal__collection__{name}__compact", timeout=0
        )

    def _needs_compactation(self):
        return self._vector_index.delta_size >= 100

    @synced
    def compact(self):
        """
        Triggers a compaction of the vector index for all processes.

        If a compaction is already running for this collection, this method returns
        immediately without starting a new one.
        """
        if self._compact_lock.acquire(block=False):
            cursor = self._db.connection.cursor()
            self._vector_index.compact(cursor)
            self._compact_lock.release()

    @emits("index", payload=lambda document, *args, **kwargs: dict(id=document.id))
    @synced
    def index(
        self,
        document: Document[B],
        *,
        fts: bool | list[str] = True,
        fuzzy: bool = False,
    ):
        """
        Indexes a Document, including vector, FTS, and fuzzy search data.
        The entire operation is performed in a single atomic transaction.
        """
        if not isinstance(document, Document):
            raise TypeError("Item to index must be a Document object.")

        cursor = self._db.connection.cursor()

        # Step 1: Core Document and Vector Storage
        cursor.execute(
            "INSERT OR REPLACE INTO beaver_collections (collection, item_id, item_vector, metadata) VALUES (?, ?, ?, ?)",
            (
                self._name,
                document.id,
                document.serialize_embedding(),
                json.dumps(document.serialize_body()),
            ),
        )

        # Step 2: Delegate to the VectorIndex if an embedding exists.
        if document.embedding is not None:
            self._vector_index.index(document.embedding, document.id, cursor)

        # Step 3: FTS and Fuzzy Indexing
        cursor.execute(
            "DELETE FROM beaver_fts_index WHERE collection = ? AND item_id = ?",
            (self._name, document.id),
        )
        cursor.execute(
            "DELETE FROM beaver_trigrams WHERE collection = ? AND item_id = ?",
            (self._name, document.id),
        )

        flat_metadata = document.flattened_body()
        fields_to_index: dict[str, str] = {}
        if isinstance(fts, list):
            fields_to_index = {
                k: v
                for k, v in flat_metadata.items()
                if k in fts and isinstance(v, str)
            }
        elif fts:
            fields_to_index = {
                k: v for k, v in flat_metadata.items() if isinstance(v, str)
            }

        if fields_to_index:
            fts_data = [
                (self._name, document.id, path, content)
                for path, content in fields_to_index.items()
            ]
            cursor.executemany(
                "INSERT INTO beaver_fts_index (collection, item_id, field_path, field_content) VALUES (?, ?, ?, ?)",
                fts_data,
            )
            if fuzzy:
                trigram_data = []
                for path, content in fields_to_index.items():
                    for trigram in _get_trigrams(content.lower()):
                        trigram_data.append((self._name, document.id, path, trigram))
                if trigram_data:
                    cursor.executemany(
                        "INSERT INTO beaver_trigrams (collection, item_id, field_path, trigram) VALUES (?, ?, ?, ?)",
                        trigram_data,
                    )

        # After the transaction commits, check if auto-compaction is needed.
        if self._needs_compactation():
            self.compact()

    @emits("drop", payload=lambda document, *args, **kwargs: dict(id=document.id))
    @synced
    def drop(self, document: Document[B]):
        """Removes a document and all its associated data from the collection."""
        if not isinstance(document, Document):
            raise TypeError("Item to drop must be a Document object.")

        cursor = self._db.connection.cursor()
        cursor.execute(
            "DELETE FROM beaver_collections WHERE collection = ? AND item_id = ?",
            (self._name, document.id),
        )
        cursor.execute(
            "DELETE FROM beaver_fts_index WHERE collection = ? AND item_id = ?",
            (self._name, document.id),
        )
        cursor.execute(
            "DELETE FROM beaver_trigrams WHERE collection = ? AND item_id = ?",
            (self._name, document.id),
        )
        cursor.execute(
            "DELETE FROM beaver_edges WHERE collection = ? AND (source_item_id = ? OR target_item_id = ?)",
            (self._name, document.id, document.id),
        )
        self._vector_index.drop(document.id, cursor)

        # Check for auto-compaction after a drop as well.
        if self._needs_compactation():
            self.compact()

    def __iter__(self) -> Iterator[Document[B]]:
        """Returns an iterator over all documents in the collection."""
        cursor = self._db.connection.cursor()
        cursor.execute(
            "SELECT item_id, item_vector, metadata FROM beaver_collections WHERE collection = ?",
            (self._name,),
        )
        for row in cursor:
            embedding = (
                np.frombuffer(row["item_vector"], dtype=np.float32).tolist()
                if row["item_vector"]
                else None
            )
            yield Document(
                id=row["item_id"],
                embedding=embedding,
                body=self._deserialize(row["metadata"]),
            )
        cursor.close()

    @synced
    def search(
        self, vector: list[float], top_k: int = 10
    ) -> List[Tuple[Document[B], float]]:
        """Performs a fast, in-memory nearest neighbor search."""
        if not isinstance(vector, list):
            raise TypeError("Search vector must be a list of floats.")

        search_results = self._vector_index.search(vector, top_k=top_k)
        if not search_results:
            return []

        result_ids = [item[0] for item in search_results]
        distance_map = {item[0]: item[1] for item in search_results}

        placeholders = ",".join("?" for _ in result_ids)
        sql = f"SELECT item_id, item_vector, metadata FROM beaver_collections WHERE collection = ? AND item_id IN ({placeholders})"

        cursor = self._db.connection.cursor()
        rows = cursor.execute(sql, (self._name, *result_ids)).fetchall()

        doc_map = {
            row["item_id"]: Document(
                id=row["item_id"],
                embedding=Document.deserialize_embedding(row["item_vector"]),
                body=self._deserialize(row["metadata"]),
            )
            for row in rows
        }

        final_results = []

        for doc_id in result_ids:
            if doc_id in doc_map:
                doc = doc_map[doc_id]
                distance = distance_map[doc_id]
                final_results.append((doc, float(distance)))

        return final_results

    @synced
    def match(
        self,
        query: str,
        *,
        on: str | list[str] | None = None,
        top_k: int = 10,
        fuzziness: int = 0,
    ) -> list[tuple[Document[B], float]]:
        """
        Performs a full-text or fuzzy search on indexed string fields.
        """
        if isinstance(on, str):
            on = [on]

        if fuzziness == 0:
            return self._perform_fts_search(query, on, top_k)
        else:
            return self._perform_fuzzy_search(query, on, top_k, fuzziness)

    def _perform_fts_search(
        self, query: str, on: list[str] | None, top_k: int
    ) -> list[tuple[Document[B], float]]:
        """Performs a standard FTS search."""
        cursor = self._db.connection.cursor()
        sql_query = """
            SELECT t1.item_id, t1.item_vector, t1.metadata, fts.rank
            FROM beaver_collections AS t1 JOIN (
                SELECT DISTINCT item_id, rank FROM beaver_fts_index
                WHERE beaver_fts_index MATCH ? {} ORDER BY rank LIMIT ?
            ) AS fts ON t1.item_id = fts.item_id
            WHERE t1.collection = ? ORDER BY fts.rank
        """
        params: list[Any] = [query]
        field_filter_sql = ""
        if on:
            placeholders = ",".join("?" for _ in on)
            field_filter_sql = f"AND field_path IN ({placeholders})"
            params.extend(on)

        params.extend([top_k, self._name])
        rows = cursor.execute(
            sql_query.format(field_filter_sql), tuple(params)
        ).fetchall()
        results = []
        for row in rows:
            doc = Document(
                id=row["item_id"],
                embedding=Document.deserialize_embedding(row["item_vector"]),
                body=self._deserialize(row["metadata"]),
            )
            results.append((doc, row["rank"]))
        return results

    def _get_trigram_candidates(self, query: str, on: list[str] | None) -> set[str]:
        """
        Gets document IDs that meet a trigram similarity threshold with the query.
        """
        query_trigrams = _get_trigrams(query.lower())
        if not query_trigrams:
            return set()

        similarity_threshold = int(len(query_trigrams) * 0.3)
        if similarity_threshold == 0:
            return set()

        cursor = self._db.connection.cursor()
        sql = """
            SELECT item_id FROM beaver_trigrams
            WHERE collection = ? AND trigram IN ({}) {}
            GROUP BY item_id
            HAVING COUNT(DISTINCT trigram) >= ?
        """
        params: list[Any] = [self._name]
        trigram_placeholders = ",".join("?" for _ in query_trigrams)
        params.extend(query_trigrams)

        field_filter_sql = ""
        if on:
            field_placeholders = ",".join("?" for _ in on)
            field_filter_sql = f"AND field_path IN ({field_placeholders})"
            params.extend(on)

        params.append(similarity_threshold)
        cursor.execute(
            sql.format(trigram_placeholders, field_filter_sql), tuple(params)
        )
        return {row["item_id"] for row in cursor.fetchall()}

    def _perform_fuzzy_search(
        self, query: str, on: list[str] | None, top_k: int, fuzziness: int
    ) -> list[tuple[Document[B], float]]:
        """Performs a 3-stage fuzzy search: gather, score, and sort."""
        fts_results = self._perform_fts_search(query, on, top_k)
        fts_candidate_ids = {doc.id for doc, _ in fts_results}
        trigram_candidate_ids = self._get_trigram_candidates(query, on)
        candidate_ids = fts_candidate_ids.union(trigram_candidate_ids)
        if not candidate_ids:
            return []

        cursor = self._db.connection.cursor()
        id_placeholders = ",".join("?" for _ in candidate_ids)
        sql_text = f"SELECT item_id, field_path, field_content FROM beaver_fts_index WHERE collection = ? AND item_id IN ({id_placeholders})"
        params_text: list[Any] = [self._name]
        params_text.extend(candidate_ids)
        if on:
            sql_text += f" AND field_path IN ({','.join('?' for _ in on)})"
            params_text.extend(on)

        cursor.execute(sql_text, tuple(params_text))
        candidate_texts: dict[str, dict[str, str]] = {}
        for row in cursor.fetchall():
            item_id = row["item_id"]
            if item_id not in candidate_texts:
                candidate_texts[item_id] = {}
            candidate_texts[item_id][row["field_path"]] = row["field_content"]

        scored_candidates = []
        fts_rank_map = {doc.id: rank for doc, rank in fts_results}

        for item_id in candidate_ids:
            if item_id not in candidate_texts:
                continue
            min_dist = float("inf")
            for content in candidate_texts[item_id].values():
                dist = _sliding_window_levenshtein(query, content, fuzziness)
                if dist < min_dist:
                    min_dist = dist
            if min_dist <= fuzziness:
                scored_candidates.append(
                    {
                        "id": item_id,
                        "distance": min_dist,
                        "fts_rank": fts_rank_map.get(item_id, 0),
                    }
                )

        scored_candidates.sort(key=lambda x: (x["distance"], x["fts_rank"]))
        top_ids = [c["id"] for c in scored_candidates[:top_k]]
        if not top_ids:
            return []

        id_placeholders = ",".join("?" for _ in top_ids)
        sql_docs = f"SELECT item_id, item_vector, metadata FROM beaver_collections WHERE collection = ? AND item_id IN ({id_placeholders})"
        cursor.execute(sql_docs, (self._name, *top_ids))
        doc_map = {
            row["item_id"]: Document(
                id=row["item_id"],
                embedding=Document.deserialize_embedding(row["item_vector"]),
                body=self._deserialize(row["metadata"]),
            )
            for row in cursor.fetchall()
        }

        final_results = []
        distance_map = {c["id"]: c["distance"] for c in scored_candidates}
        for doc_id in top_ids:
            if doc_id in doc_map:
                final_results.append((doc_map[doc_id], float(distance_map[doc_id])))
        return final_results

    @emits(
        "connect",
        payload=lambda src, tgt, *args, **kwargs: dict(src_id=src.id, tgt_id=tgt.id),
    )
    @synced
    def connect(self, source: Document, target: Document, label: str):
        """Creates a directed edge between two documents."""
        if not isinstance(source, Document) or not isinstance(target, Document):
            raise TypeError("Source and target must be Document objects.")

        self._db.connection.execute(
            """INSERT OR REPLACE INTO beaver_edges
            (collection, source_item_id, target_item_id, label, metadata)
            VALUES (?, ?, ?, ?, ?)
            """,
            (
                self._name,
                source.id,
                target.id,
                label,
                None,
            ),
        )

    def neighbors(
        self, doc: Document[B], label: str | None = None
    ) -> list[Document[B]]:
        """Retrieves the neighboring documents connected to a given document."""
        sql = """
        SELECT DISTINCT t1.item_id, t1.item_vector, t1.metadata
        FROM beaver_collections AS t1
        JOIN beaver_edges AS t2
        ON t1.item_id = t2.target_item_id AND t1.collection = t2.collection
        WHERE t2.collection = ? AND t2.source_item_id = ?
        """
        params = [self._name, doc.id]
        if label:
            sql += " AND t2.label = ?"
            params.append(label)

        rows = self._db.connection.cursor().execute(sql, tuple(params)).fetchall()
        return [
            Document(
                id=row["item_id"],
                embedding=Document.deserialize_embedding(row["item_vector"]),
                body=self._deserialize(row["metadata"]),
            )
            for row in rows
        ]

    def walk(
        self,
        source: Document[B],
        labels: List[str],
        depth: int,
        *,
        direction: Literal["outgoing", "incoming"] = "outgoing",
    ) -> List[Document[B]]:
        """Performs a graph traversal (BFS) from a starting document."""
        if not isinstance(source, Document):
            raise TypeError("The starting point must be a Document object.")

        if depth <= 0:
            return []

        source_col, target_col = (
            ("source_item_id", "target_item_id")
            if direction == "outgoing"
            else ("target_item_id", "source_item_id")
        )
        sql = f"""
        WITH RECURSIVE walk_bfs(item_id, current_depth) AS (
            SELECT ?, 0
            UNION ALL
            SELECT edges.{target_col}, bfs.current_depth + 1
            FROM beaver_edges AS edges JOIN walk_bfs AS bfs ON edges.{source_col} = bfs.item_id
            WHERE edges.collection = ? AND bfs.current_depth < ? AND edges.label IN ({','.join('?' for _ in labels)})
        )
        SELECT DISTINCT t1.item_id, t1.item_vector, t1.metadata
        FROM beaver_collections AS t1 JOIN walk_bfs AS bfs ON t1.item_id = bfs.item_id
        WHERE t1.collection = ? AND bfs.current_depth > 0
        """
        params = [source.id, self._name, depth] + labels + [self._name]

        rows = self._db.connection.cursor().execute(sql, tuple(params)).fetchall()
        return [
            Document(
                id=row["item_id"],
                embedding=Document.deserialize_embedding(row["item_vector"]),
                body=self._deserialize(row["metadata"]),
            )
            for row in rows
        ]

    def __len__(self) -> int:
        """Returns the number of documents in the collection."""
        cursor = self._db.connection.cursor()
        cursor.execute(
            "SELECT COUNT(*) FROM beaver_collections WHERE collection = ?",
            (self._name,),
        )
        count = cursor.fetchone()[0]
        cursor.close()
        return count

    def _get_dump_object(self) -> dict:
        """Builds the JSON-compatible dump object."""

        # The __iter__ method yields Document instances.
        # doc.to_dict(metadata_only=False) correctly serializes
        # the Document (including embedding) to a plain dictionary.
        items_list = [doc.model_dump(mode="json") for doc in self]

        metadata = {
            "type": "Collection",
            "name": self._name,
            "count": len(items_list),
            "dump_date": datetime.now(timezone.utc).isoformat(),
        }

        return {"metadata": metadata, "items": items_list}

    @overload
    def dump(self) -> dict:
        pass

    @overload
    def dump(self, fp: IO[str]) -> None:
        pass

    def dump(self, fp: IO[str] | None = None) -> dict | None:
        """
        Dumps the entire contents of the collection to a JSON-compatible
        Python object or a file-like object.

        Args:
            fp: A file-like object opened in text mode (e.g., with 'w').
                If provided, the JSON dump will be written to this file.
                If None (default), the dump will be returned as a dictionary.

        Returns:
            A dictionary containing the dump if fp is None.
            None if fp is provided.
        """
        dump_object = self._get_dump_object()

        if fp:
            json.dump(dump_object, fp, indent=2)
            return None

        return dump_object

    @synced
    def clear(self):
        """
        Atomically removes all documents and associated data (vectors,
        FTS indexes, graph edges) from this collection.

        This operation is protected by the manager's inter-process lock
        and will reset the vector index.
        """
        cursor = self._db.connection.cursor()

        # 1. Delete from all related tables
        tables_to_clear = [
            "beaver_collections",
            "beaver_fts_index",
            "beaver_trigrams",
            "beaver_edges",
        ]
        for table in tables_to_clear:
            cursor.execute(f"DELETE FROM {table} WHERE collection = ?", (self._name,))

        # 2. Clear vector-specific tables
        cursor.execute(
            "DELETE FROM beaver_vector_change_log WHERE collection_name = ?",
            (self._name,),
        )
        cursor.execute(
            "DELETE FROM beaver_collection_versions WHERE collection_name = ?",
            (self._name,),
        )

        # 3. Reset the in-memory vector index
        # Re-initializing the VectorIndex is the cleanest way
        # to reset its internal state.
        self._vector_index = VectorIndex(self._name, self._db)


def rerank[B: BaseModel](
    *results: list[Document[B]], weights: list[float] | None = None, k: int = 60
) -> list[Document[B]]:
    """
    Reranks documents from multiple search result lists using Reverse Rank Fusion (RRF).
    """
    if not results:
        return []

    if weights is None:
        weights = [1.0] * len(results)

    if len(results) != len(weights):
        raise ValueError("The number of result lists must match the number of weights.")

    rrf_scores: dict[str, float] = {}
    doc_store: dict[str, Document[B]] = {}

    for result_list, weight in zip(results, weights):
        for rank, doc in enumerate(result_list):
            doc_id = doc.id
            if doc_id not in doc_store:
                doc_store[doc_id] = doc
            score = weight * (1 / (k + rank))
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0.0) + score

    sorted_doc_ids = sorted(
        rrf_scores.keys(), key=lambda k: rrf_scores[k], reverse=True
    )
    return [doc_store[doc_id] for doc_id in sorted_doc_ids]
