import json
import re

from haiku.rag.chunker import chunker
from haiku.rag.embeddings import get_embedder
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.repositories.base import BaseRepository


class ChunkRepository(BaseRepository[Chunk]):
    """Repository for Chunk database operations."""

    def __init__(self, store):
        super().__init__(store)
        self.embedder = get_embedder()

    async def create(self, entity: Chunk, commit: bool = True) -> Chunk:
        """Create a chunk in the database."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()
        cursor.execute(
            """
            INSERT INTO chunks (document_id, content, metadata)
            VALUES (:document_id, :content, :metadata)
            """,
            {
                "document_id": entity.document_id,
                "content": entity.content,
                "metadata": json.dumps(entity.metadata),
            },
        )

        entity.id = cursor.lastrowid

        # Generate and store embedding
        embedding = await self.embedder.embed(entity.content)
        serialized_embedding = self.store.serialize_embedding(embedding)
        cursor.execute(
            """
            INSERT INTO chunk_embeddings (chunk_id, embedding)
            VALUES (:chunk_id, :embedding)
            """,
            {"chunk_id": entity.id, "embedding": serialized_embedding},
        )

        # Insert into FTS5 table for full-text search
        cursor.execute(
            """
            INSERT INTO chunks_fts(rowid, content)
            VALUES (:rowid, :content)
            """,
            {"rowid": entity.id, "content": entity.content},
        )

        if commit:
            self.store._connection.commit()
        return entity

    async def get_by_id(self, entity_id: int) -> Chunk | None:
        """Get a chunk by its ID."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()
        cursor.execute(
            """
            SELECT id, document_id, content, metadata
            FROM chunks WHERE id = :id
            """,
            {"id": entity_id},
        )

        row = cursor.fetchone()
        if row is None:
            return None

        chunk_id, document_id, content, metadata_json = row
        metadata = json.loads(metadata_json) if metadata_json else {}

        return Chunk(
            id=chunk_id, document_id=document_id, content=content, metadata=metadata
        )

    async def update(self, entity: Chunk) -> Chunk:
        """Update an existing chunk."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")
        if entity.id is None:
            raise ValueError("Chunk ID is required for update")

        cursor = self.store._connection.cursor()
        cursor.execute(
            """
            UPDATE chunks
            SET document_id = :document_id, content = :content, metadata = :metadata
            WHERE id = :id
            """,
            {
                "document_id": entity.document_id,
                "content": entity.content,
                "metadata": json.dumps(entity.metadata),
                "id": entity.id,
            },
        )

        # Regenerate and update embedding
        embedding = await self.embedder.embed(entity.content)
        serialized_embedding = self.store.serialize_embedding(embedding)
        cursor.execute(
            """
            UPDATE chunk_embeddings
            SET embedding = :embedding
            WHERE chunk_id = :chunk_id
            """,
            {"embedding": serialized_embedding, "chunk_id": entity.id},
        )

        # Update FTS5 table
        cursor.execute(
            """
            UPDATE chunks_fts
            SET content = :content
            WHERE rowid = :rowid
            """,
            {"content": entity.content, "rowid": entity.id},
        )

        self.store._connection.commit()
        return entity

    async def delete(self, entity_id: int, commit: bool = True) -> bool:
        """Delete a chunk by its ID."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()

        # Delete from FTS5 table first
        cursor.execute(
            "DELETE FROM chunks_fts WHERE rowid = :rowid", {"rowid": entity_id}
        )

        # Delete the embedding
        cursor.execute(
            "DELETE FROM chunk_embeddings WHERE chunk_id = :chunk_id",
            {"chunk_id": entity_id},
        )

        # Delete the chunk
        cursor.execute("DELETE FROM chunks WHERE id = :id", {"id": entity_id})

        deleted = cursor.rowcount > 0
        if commit:
            self.store._connection.commit()
        return deleted

    async def list_all(
        self, limit: int | None = None, offset: int | None = None
    ) -> list[Chunk]:
        """List all chunks with optional pagination."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()
        query = "SELECT id, document_id, content, metadata FROM chunks ORDER BY document_id, id"
        params = {}

        if limit is not None:
            query += " LIMIT :limit"
            params["limit"] = limit

        if offset is not None:
            query += " OFFSET :offset"
            params["offset"] = offset

        cursor.execute(query, params)
        rows = cursor.fetchall()

        return [
            Chunk(
                id=chunk_id,
                document_id=document_id,
                content=content,
                metadata=json.loads(metadata_json) if metadata_json else {},
            )
            for chunk_id, document_id, content, metadata_json in rows
        ]

    async def create_chunks_for_document(
        self, document_id: int, content: str, commit: bool = True
    ) -> list[Chunk]:
        """Create chunks and embeddings for a document."""
        # Chunk the document content
        chunk_texts = await chunker.chunk(content)
        created_chunks = []

        # Create chunks with embeddings using the create method
        for order, chunk_text in enumerate(chunk_texts):
            # Create chunk with order in metadata
            chunk = Chunk(
                document_id=document_id, content=chunk_text, metadata={"order": order}
            )

            created_chunk = await self.create(chunk, commit=commit)
            created_chunks.append(created_chunk)

        return created_chunks

    async def delete_by_document_id(
        self, document_id: int, commit: bool = True
    ) -> bool:
        """Delete all chunks for a document."""
        chunks = await self.get_by_document_id(document_id)

        deleted_any = False
        for chunk in chunks:
            if chunk.id is not None:
                deleted = await self.delete(chunk.id, commit=False)
                deleted_any = deleted_any or deleted

        if commit and deleted_any and self.store._connection:
            self.store._connection.commit()
        return deleted_any

    async def search_chunks(
        self, query: str, limit: int = 5
    ) -> list[tuple[Chunk, float]]:
        """Search for relevant chunks using vector similarity."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()

        # Generate embedding for the query
        query_embedding = await self.embedder.embed(query)
        serialized_query_embedding = self.store.serialize_embedding(query_embedding)

        # Search for similar chunks using sqlite-vec
        cursor.execute(
            """
            SELECT c.id, c.document_id, c.content, c.metadata, distance, d.uri, d.metadata as document_metadata
            FROM chunk_embeddings
            JOIN chunks c ON c.id = chunk_embeddings.chunk_id
            JOIN documents d ON c.document_id = d.id
            WHERE embedding MATCH :embedding AND k = :k
            ORDER BY distance
            """,
            {"embedding": serialized_query_embedding, "k": limit},
        )

        results = cursor.fetchall()
        return [
            (
                Chunk(
                    id=chunk_id,
                    document_id=document_id,
                    content=content,
                    metadata=json.loads(metadata_json) if metadata_json else {},
                    document_uri=document_uri,
                    document_meta=json.loads(document_metadata_json)
                    if document_metadata_json
                    else {},
                ),
                1.0 / (1.0 + distance),
            )
            for chunk_id, document_id, content, metadata_json, distance, document_uri, document_metadata_json in results
        ]

    async def search_chunks_fts(
        self, query: str, limit: int = 5
    ) -> list[tuple[Chunk, float]]:
        """Search for chunks using FTS5 full-text search."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()

        # Clean the query for FTS5 - extract keywords for better matching
        # Remove special characters and split into words
        words = re.findall(r"\b\w+\b", query.lower())
        # Join with OR to find chunks containing any of the keywords
        fts_query = " OR ".join(words) if words else query

        # Search using FTS5
        cursor.execute(
            """
            SELECT c.id, c.document_id, c.content, c.metadata, rank, d.uri, d.metadata as document_metadata
            FROM chunks_fts
            JOIN chunks c ON c.id = chunks_fts.rowid
            JOIN documents d ON c.document_id = d.id
            WHERE chunks_fts MATCH :query
            ORDER BY rank
            LIMIT :limit
            """,
            {"query": fts_query, "limit": limit},
        )

        results = cursor.fetchall()

        return [
            (
                Chunk(
                    id=chunk_id,
                    document_id=document_id,
                    content=content,
                    metadata=json.loads(metadata_json) if metadata_json else {},
                    document_uri=document_uri,
                    document_meta=json.loads(document_metadata_json)
                    if document_metadata_json
                    else {},
                ),
                -rank,
            )
            for chunk_id, document_id, content, metadata_json, rank, document_uri, document_metadata_json in results
            # FTS5 rank is negative BM25 score
        ]

    async def search_chunks_hybrid(
        self, query: str, limit: int = 5, k: int = 60
    ) -> list[tuple[Chunk, float]]:
        """Hybrid search using Reciprocal Rank Fusion (RRF) combining vector similarity and FTS5 full-text search."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()

        # Generate embedding for the query
        query_embedding = await self.embedder.embed(query)
        serialized_query_embedding = self.store.serialize_embedding(query_embedding)

        # Clean the query for FTS5 - extract keywords for better matching
        # Remove special characters and split into words
        words = re.findall(r"\b\w+\b", query.lower())
        # Join with OR to find chunks containing any of the keywords
        fts_query = " OR ".join(words) if words else query
        # Perform hybrid search using RRF (Reciprocal Rank Fusion)
        cursor.execute(
            """
            WITH vector_search AS (
                SELECT
                    c.id,
                    c.document_id,
                    c.content,
                    c.metadata,
                    ROW_NUMBER() OVER (ORDER BY ce.distance) as vector_rank
                FROM chunk_embeddings ce
                JOIN chunks c ON c.id = ce.chunk_id
                WHERE ce.embedding MATCH :embedding AND k = :k_vector
                ORDER BY ce.distance
            ),
            fts_search AS (
                SELECT
                    c.id,
                    c.document_id,
                    c.content,
                    c.metadata,
                    ROW_NUMBER() OVER (ORDER BY chunks_fts.rank) as fts_rank
                FROM chunks_fts
                JOIN chunks c ON c.id = chunks_fts.rowid
                WHERE chunks_fts MATCH :fts_query
                ORDER BY chunks_fts.rank
            ),
            all_chunks AS (
                SELECT id, document_id, content, metadata FROM vector_search
                UNION
                SELECT id, document_id, content, metadata FROM fts_search
            ),
            rrf_scores AS (
                SELECT
                    a.id,
                    a.document_id,
                    a.content,
                    a.metadata,
                    COALESCE(1.0 / (:k + v.vector_rank), 0) + COALESCE(1.0 / (:k + f.fts_rank), 0) as rrf_score
                FROM all_chunks a
                LEFT JOIN vector_search v ON a.id = v.id
                LEFT JOIN fts_search f ON a.id = f.id
            )
            SELECT r.id, r.document_id, r.content, r.metadata, r.rrf_score, d.uri, d.metadata as document_metadata
            FROM rrf_scores r
            JOIN documents d ON r.document_id = d.id
            ORDER BY r.rrf_score DESC
            LIMIT :limit
            """,
            {
                "embedding": serialized_query_embedding,
                "k_vector": limit * 3,
                "fts_query": fts_query,
                "k": k,
                "limit": limit,
            },
        )

        results = cursor.fetchall()
        return [
            (
                Chunk(
                    id=chunk_id,
                    document_id=document_id,
                    content=content,
                    metadata=json.loads(metadata_json) if metadata_json else {},
                    document_uri=document_uri,
                    document_meta=json.loads(document_metadata_json)
                    if document_metadata_json
                    else {},
                ),
                rrf_score,
            )
            for chunk_id, document_id, content, metadata_json, rrf_score, document_uri, document_metadata_json in results
        ]

    async def get_by_document_id(self, document_id: int) -> list[Chunk]:
        """Get all chunks for a specific document."""
        if self.store._connection is None:
            raise ValueError("Store connection is not available")

        cursor = self.store._connection.cursor()
        cursor.execute(
            """
            SELECT c.id, c.document_id, c.content, c.metadata, d.uri, d.metadata as document_metadata
            FROM chunks c
            JOIN documents d ON c.document_id = d.id
            WHERE c.document_id = :document_id
            ORDER BY JSON_EXTRACT(c.metadata, '$.order')
            """,
            {"document_id": document_id},
        )

        rows = cursor.fetchall()
        return [
            Chunk(
                id=chunk_id,
                document_id=document_id,
                content=content,
                metadata=json.loads(metadata_json) if metadata_json else {},
                document_uri=document_uri,
                document_meta=json.loads(document_metadata_json)
                if document_metadata_json
                else {},
            )
            for chunk_id, document_id, content, metadata_json, document_uri, document_metadata_json in rows
        ]
