from .models import Document, GraphChunk, ChunkRelations, QUERY


class AsyncNeo4jRepository:
    def __init__(self, neo4j):
        self.neo4j = neo4j

    async def add_document_with_chunk(
        self, doc: Document, chunk: GraphChunk, relations: ChunkRelations
    ):
        async with self.neo4j.session() as session:
            await session.run(
                QUERY,
                milvus_id=chunk.milvus_id,
                doc_id=doc.doc_id,
                dataset_id=doc.dataset_id,
                chunk_id=chunk.chunk_id,
                topic=chunk.topic,
                domain=chunk.domain,
                text_type=chunk.text_type,
                task_type=chunk.task_type,
                entities=relations.entities,
                concepts=relations.concepts,
                keywords=relations.keywords,
                domain_name=relations.domain,
                topic_name=relations.topic,
            )
        print(f"✅ {doc.doc_id} - {chunk.chunk_id} 已写入 Neo4j (async)")