import asyncio
from uuid import uuid5, NAMESPACE_OID
from typing import Type
from pydantic import BaseModel
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.databases.vector import get_vector_engine, DataPoint
from cognee.modules.data.extraction.extract_categories import extract_categories
from cognee.modules.data.processing.chunk_types.DocumentChunk import DocumentChunk


async def chunk_naive_llm_classifier(data_chunks: list[DocumentChunk], classification_model: Type[BaseModel]):
    if len(data_chunks) == 0:
        return data_chunks

    chunk_classifications = await asyncio.gather(
        *[extract_categories(chunk.text, classification_model) for chunk in data_chunks],
    )

    classification_data_points = []

    for chunk_index, chunk in enumerate(data_chunks):
        chunk_classification = chunk_classifications[chunk_index]
        classification_data_points.append(uuid5(NAMESPACE_OID, chunk_classification.label.type))
        classification_data_points.append(uuid5(NAMESPACE_OID, chunk_classification.label.type))

        for classification_subclass in chunk_classification.label.subclass:
            classification_data_points.append(uuid5(NAMESPACE_OID, classification_subclass.value))

    vector_engine = get_vector_engine()

    class Keyword(BaseModel):
        uuid: str
        text: str
        chunk_id: str
        document_id: str

    collection_name = "classification"

    if await vector_engine.has_collection(collection_name):
        existing_data_points = await vector_engine.retrieve(
            collection_name,
            list(set(classification_data_points)),
        ) if len(classification_data_points) > 0 else []

        existing_points_map = {point.id: True for point in existing_data_points}
    else:
        existing_points_map = {}
        await vector_engine.create_collection(collection_name, payload_schema=Keyword)

    data_points = []
    nodes = []
    edges = []

    for (chunk_index, data_chunk) in enumerate(data_chunks):
        chunk_classification = chunk_classifications[chunk_index]
        classification_type_label = chunk_classification.label.type
        classification_type_id = uuid5(NAMESPACE_OID, classification_type_label)

        if classification_type_id not in existing_points_map:
            data_points.append(
                DataPoint[Keyword](
                    id=str(classification_type_id),
                    payload=Keyword.parse_obj({
                        "uuid": str(classification_type_id),
                        "text": classification_type_label,
                        "chunk_id": str(data_chunk.chunk_id),
                        "document_id": str(data_chunk.document_id),
                    }),
                    embed_field="text",
                )
            )

            nodes.append((
                str(classification_type_id),
                dict(
                    id=str(classification_type_id),
                    name=classification_type_label,
                    type=classification_type_label,
                )
            ))
            existing_points_map[classification_type_id] = True

        edges.append((
            str(data_chunk.chunk_id),
            str(classification_type_id),
            "is_media_type",
            dict(
                relationship_name="is_media_type",
                source_node_id=str(data_chunk.chunk_id),
                target_node_id=str(classification_type_id),
            ),
        ))

        for classification_subclass in chunk_classification.label.subclass:
            classification_subtype_label = classification_subclass.value
            classification_subtype_id = uuid5(NAMESPACE_OID, classification_subtype_label)

            if classification_subtype_id not in existing_points_map:
                data_points.append(
                    DataPoint[Keyword](
                        id=str(classification_subtype_id),
                        payload=Keyword.parse_obj({
                            "uuid": str(classification_subtype_id),
                            "text": classification_subtype_label,
                            "chunk_id": str(data_chunk.chunk_id),
                            "document_id": str(data_chunk.document_id),
                        }),
                        embed_field="text",
                    )
                )

                nodes.append((
                    str(classification_subtype_id),
                    dict(
                        id=str(classification_subtype_id),
                        name=classification_subtype_label,
                        type=classification_subtype_label,
                    )
                ))
                edges.append((
                    str(classification_subtype_id),
                    str(classification_type_id),
                    "is_subtype_of",
                    dict(
                        relationship_name="contains",
                        source_node_id=str(classification_type_id),
                        target_node_id=str(classification_subtype_id),
                    ),
                ))

                existing_points_map[classification_subtype_id] = True

            edges.append((
                str(data_chunk.chunk_id),
                str(classification_subtype_id),
                "is_classified_as",
                dict(
                    relationship_name="is_classified_as",
                    source_node_id=str(data_chunk.chunk_id),
                    target_node_id=str(classification_subtype_id),
                ),
            ))

    if len(nodes) > 0 or len(edges) > 0:
        await vector_engine.create_data_points(collection_name, data_points)

        graph_engine = await get_graph_engine()

        await graph_engine.add_nodes(nodes)
        await graph_engine.add_edges(edges)

    return data_chunks
