import os
from haystack import Pipeline
from haystack.components.writers import DocumentWriter
from haystack_integrations.components.embedders.mistral.document_embedder import (
    MistralDocumentEmbedder,
)
from neo4j_haystack import Neo4jDocumentStore


from .file_upload import knowledgebase_file_to_docs
from .models import KnowledgeBase

MISTRAL_API_KEY = os.environ["MISTRAL_API_KEY"]
NEO4J_URI = os.environ["NEO4J_URI"]
NEO4J_USER = os.environ["NEO4J_USERNAME"]
NEO4J_PASS = os.environ["NEO4J_PASSWORD"]


def publish_to_neo4j(kb_instance: KnowledgeBase):
    """
    Process a KnowledgeBase file and store embeddings in Neo4j.
    """
    file_path = kb_instance.document.path
    print(f"[Publish Service] Processing file: {file_path}")

    # Step 1 – Prepare final docs
    final_docs = knowledgebase_file_to_docs(file_path)
    if not final_docs:
        raise ValueError("No valid documents extracted from file")

    # Step 2 – Connect to Neo4j
    document_store = Neo4jDocumentStore(
        url=NEO4J_URI,
        username=NEO4J_USER,
        password=NEO4J_PASS,
        database="fc875647",  # you may want to move this into settings
        index="document-embeddings",
        embedding_field="embedding",
        embedding_dim=1024,
        node_label="Document",
    )
    try:
        document_store.create_index_if_not_exists()
    except Exception as e:
        print(f"[Publish Service] Index creation warning: {e}")

    # Step 3 – Init embedder + writer
    embedder = MistralDocumentEmbedder()
    writer = DocumentWriter(document_store=document_store)

    indexing_pipeline = Pipeline()
    indexing_pipeline.add_component("embedder", embedder)
    indexing_pipeline.add_component("writer", writer)
    indexing_pipeline.connect("embedder", "writer")

    # Step 4 – Run pipeline in batches
    batch_size = 3
    for i in range(0, len(final_docs), batch_size):
        documents_batch = final_docs[i : i + batch_size]
        print(f"[Publish Service] Processing batch {i//batch_size + 1}")
        result = indexing_pipeline.run(
            data={"embedder": {"documents": documents_batch}}
        )
        print(
            f"[Publish Service] Documents written: {result['writer']['documents_written']}"
        )

    # Step 5 – Update KnowledgeBase instance
    kb_instance.is_published = True
    kb_instance.save(update_fields=["is_published"])

    print(f"[Publish Service] Finished publishing {kb_instance.id}")
    return True
