import logging
from collections import defaultdict
from typing import Dict, Literal

from kiln_ai.adapters.rag.deduplication import (
    deduplicate_chunk_embeddings,
    deduplicate_chunked_documents,
    deduplicate_extractions,
    filter_documents_by_tags,
)
from kiln_ai.adapters.vector_store.vector_store_registry import (
    vector_store_adapter_for_config,
)
from kiln_ai.datamodel.project import Project
from kiln_ai.datamodel.rag import RagConfig
from kiln_ai.datamodel.vector_store import VectorStoreConfig
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


class LogMessage(BaseModel):
    level: Literal["info", "error", "warning"] = Field(
        description="The level of the log message",
    )
    message: str = Field(
        description="The message to display to the user",
    )


class RagProgress(BaseModel):
    total_document_count: int = Field(
        description="The total number of items to process",
        default=0,
    )

    total_document_completed_count: int = Field(
        description="The number of items that have been processed",
        default=0,
    )

    # Progress for indexing is tracked in terms of chunks, not documents. After the initial run
    # the only info we have is how many chunks are in the vector store, and so we need to know
    # the total number of chunks that should be indexed to know if it is completed or not.
    # So we need toset and send that through to the client once we know it (after completing chunking).
    total_chunk_count: int = Field(
        description="The number of chunks that should be indexed for the indexing to be completed.",
        default=0,
    )

    total_chunk_completed_count: int = Field(
        description="The number of chunks that have been indexed",
        default=0,
    )

    total_document_extracted_count: int = Field(
        description="The number of items that have been extracted",
        default=0,
    )

    total_document_extracted_error_count: int = Field(
        description="The number of items that have errored during extraction",
        default=0,
    )

    total_document_chunked_count: int = Field(
        description="The number of items that have been chunked",
        default=0,
    )

    total_document_chunked_error_count: int = Field(
        description="The number of items that have errored during chunking",
        default=0,
    )

    total_document_embedded_count: int = Field(
        description="The number of items that have been embedded",
        default=0,
    )

    total_document_embedded_error_count: int = Field(
        description="The number of items that have errored during embedding",
        default=0,
    )

    total_chunks_indexed_count: int = Field(
        description="The number of chunks that have been indexed",
        default=0,
    )

    total_chunks_indexed_error_count: int = Field(
        description="The number of chunks that have errored during indexing",
        default=0,
    )

    logs: list[LogMessage] | None = Field(
        description="A list of log messages to display to the user",
        default=None,
    )


async def count_records_in_vector_store(
    rag_config: RagConfig,
    vector_store_config: VectorStoreConfig,
) -> int:
    vector_store = await vector_store_adapter_for_config(
        rag_config, vector_store_config
    )
    count = await vector_store.count_records()
    return count


async def count_records_in_vector_store_for_rag_config(
    project: Project,
    rag_config: RagConfig,
) -> int:
    vector_store_config = VectorStoreConfig.from_id_and_parent_path(
        str(rag_config.vector_store_config_id),
        project.path,
    )
    if vector_store_config is None:
        raise ValueError(f"Rag config {rag_config.id} has no vector store config")
    return await count_records_in_vector_store(rag_config, vector_store_config)


async def compute_current_progress_for_rag_configs(
    project: Project,
    rag_configs: list[RagConfig],
) -> Dict[str, RagProgress]:
    # each RAG config represents a unique path: extractor::chunker::embedding
    # different configs can share common prefixes
    # (e.g., extractor-1::chunker-2 for both extractor-1::chunker-2::embedding-3 and extractor-1::chunker-2::embedding-4)
    # we store prefix -> [rag config ids] mappings so at every step
    # we know all the configs that share the same upstream steps
    path_prefixes: dict[str, set[str]] = defaultdict(set)
    for rag_config in rag_configs:
        complete_path: list[str] = [
            str(rag_config.extractor_config_id),
            str(rag_config.chunker_config_id),
            str(rag_config.embedding_config_id),
        ]
        for i in range(len(complete_path)):
            prefix = "::".join(complete_path[: i + 1])
            path_prefixes[prefix].add(str(rag_config.id))

    rag_config_progress_map: dict[str, RagProgress] = {}
    for rag_config in rag_configs:
        all_documents = project.documents(readonly=True)
        if rag_config.tags:
            all_documents = filter_documents_by_tags(all_documents, rag_config.tags)

        rag_config_progress_map[str(rag_config.id)] = RagProgress(
            total_document_count=len(all_documents),
            total_document_completed_count=0,
            total_chunk_count=0,
            total_chunk_completed_count=0,
            total_document_extracted_count=0,
            total_document_chunked_count=0,
            total_document_embedded_count=0,
            total_chunks_indexed_count=await count_records_in_vector_store_for_rag_config(
                project, rag_config
            ),
        )

    # Create a mapping of rag_config_id to its tags for efficient lookup
    rag_config_tags_map = {
        str(rag_config.id): rag_config.tags for rag_config in rag_configs
    }

    for document in project.documents(readonly=True):
        for extraction in deduplicate_extractions(document.extractions(readonly=True)):
            extraction_path_prefix = str(extraction.extractor_config_id)

            # increment the extraction count for every rag config that has this extractor
            # and includes this document based on its tags
            for matching_rag_config_id in path_prefixes[extraction_path_prefix]:
                rag_config_tags = rag_config_tags_map[matching_rag_config_id]
                if not rag_config_tags or (
                    document.tags
                    and any(tag in document.tags for tag in rag_config_tags)
                ):
                    rag_config_progress_map[
                        matching_rag_config_id
                    ].total_document_extracted_count += 1

            for chunked_document in deduplicate_chunked_documents(
                extraction.chunked_documents(readonly=True)
            ):
                # increment the chunked count for every rag config that has this extractor+chunker combo
                # and includes this document based on its tags
                chunking_path_prefix = (
                    f"{extraction_path_prefix}::{chunked_document.chunker_config_id}"
                )
                for matching_rag_config_id in path_prefixes[chunking_path_prefix]:
                    rag_config_tags = rag_config_tags_map[matching_rag_config_id]
                    if not rag_config_tags or (
                        document.tags
                        and any(tag in document.tags for tag in rag_config_tags)
                    ):
                        rag_config_progress_map[
                            matching_rag_config_id
                        ].total_document_chunked_count += 1

                        rag_config_progress_map[
                            matching_rag_config_id
                        ].total_chunk_count += len(chunked_document.chunks)

                for embedding in deduplicate_chunk_embeddings(
                    chunked_document.chunk_embeddings(readonly=True)
                ):
                    # increment the embedding count for every rag config that has this extractor+chunker+embedding combo
                    # and includes this document based on its tags
                    embedding_path_prefix = (
                        f"{chunking_path_prefix}::{embedding.embedding_config_id}"
                    )

                    for matching_rag_config_id in path_prefixes[embedding_path_prefix]:
                        rag_config_tags = rag_config_tags_map[matching_rag_config_id]
                        if not rag_config_tags or (
                            document.tags
                            and any(tag in document.tags for tag in rag_config_tags)
                        ):
                            rag_config_progress_map[
                                matching_rag_config_id
                            ].total_document_embedded_count += 1

    # a document is completed only when all steps are completed, so overall progress is the same
    # as the least complete step
    for _, rag_config_progress in rag_config_progress_map.items():
        rag_config_progress.total_document_completed_count = min(
            rag_config_progress.total_document_extracted_count,
            rag_config_progress.total_document_chunked_count,
            rag_config_progress.total_document_embedded_count,
        )

        rag_config_progress.total_chunk_completed_count = (
            rag_config_progress.total_chunks_indexed_count
        )

    return dict(rag_config_progress_map)


async def compute_current_progress_for_rag_config(
    project: Project,
    rag_config: RagConfig,
) -> RagProgress:
    config_progress = await compute_current_progress_for_rag_configs(
        project, [rag_config]
    )
    if str(rag_config.id) not in config_progress:
        raise ValueError(f"Failed to compute progress for rag config {rag_config.id}")
    return config_progress[str(rag_config.id)]
