from ingestion import IngestAnything, VectorStoreIndex, IngestCode
from typing import Any
from qdrant_client import QdrantClient, AsyncQdrantClient
import weaviate
from weaviate import WeaviateClient
from llama_index.vector_stores.weaviate import WeaviateVectorStore
from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.vector_stores.qdrant import QdrantVectorStore
import pathlib
import sys
import os
from dotenv import load_dotenv
from pymilvus import MilvusClient
from custom_reader_to_test import MarkItDownReader

load_dotenv()

# Readers config
markitdown_reader = MarkItDownReader()


# Abstract client class for collection check and delete operations
class AbstractClient:
    def __init__(self, client: Any):
        self.client = client

    def delete_collection(self, collection_name: str = "Test") -> None:
        if isinstance(self.client, WeaviateClient):
            self.client.collections.delete(collection_name)
        elif isinstance(self.client, MilvusClient):
            self.client.drop_collection(collection_name)
        elif isinstance(self.client, QdrantClient):
            self.client.delete_collection(collection_name)
        return

    def check_collection_exists(self, collection_name: str = "Test") -> bool:
        if isinstance(self.client, WeaviateClient):
            return self.client.collections.exists(collection_name)
        elif isinstance(self.client, MilvusClient):
            r = self.client.list_collections()
            return collection_name in r
        elif isinstance(self.client, QdrantClient):
            return self.client.collection_exists(collection_name)
        return False


# Qdrant Config
client_qdrant = QdrantClient(
    api_key=os.getenv("qdrant_api_key"), url=os.getenv("qdrant_url")
)
aclient_qdrant = AsyncQdrantClient(
    api_key=os.getenv("qdrant_api_key"), url=os.getenv("qdrant_url")
)
vector_store_qdrant = QdrantVectorStore(
    collection_name="Test", client=client_qdrant, aclient=aclient_qdrant
)
abstract_qdrant_client = AbstractClient(client_qdrant)

# Weaviate config
cluster_url = os.getenv("weaviate_cluster_url")
api_key = os.getenv("weaviate_admin_key")
client_weaviate = weaviate.connect_to_weaviate_cloud(
    cluster_url=cluster_url,
    auth_credentials=weaviate.auth.AuthApiKey(api_key),
)
vector_store_weaviate = WeaviateVectorStore(
    weaviate_client=client_weaviate, index_name="Test"
)
abstract_weaviate_client = AbstractClient(client_weaviate)

# Milvus Config
milvus_uri = os.getenv("milvus_uri")
milvus_token = os.getenv("milvus_token")
client_milvus = MilvusClient(uri=milvus_uri, token=milvus_token)
vector_store_milvus = MilvusVectorStore(
    uri=milvus_uri, token=milvus_token, overwrite=True, collection_name="Test", dim=384
)
abstract_milvus_client = AbstractClient(client_milvus)


def test_ingestion():
    test_cases = [
        {
            "chunker": "late",
            "chunk_size": None,
            "chunk_overlap": None,
            "similarity_threshold": None,
            "min_characters_per_chunk": None,
            "min_sentences": None,
            "files_or_dir": "tests/data",
            "tokenizer": "gpt2",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "slumber_genie": None,
            "slumber_model": None,
            "vector_store": vector_store_weaviate,
            "reader": markitdown_reader,
            "client": abstract_qdrant_client,
            "expected": [True, True],
        },
        {
            "chunker": "semantic",
            "chunk_size": None,
            "chunk_overlap": None,
            "similarity_threshold": None,
            "min_characters_per_chunk": None,
            "min_sentences": None,
            "files_or_dir": [
                "tests/data/test.docx",
                "tests/data/test0.png",
                "tests/data/test1.csv",
                "tests/data/test2.json",
                "tests/data/test3.md",
                "tests/data/test4.xml",
                "tests/data/test5.zip",
            ],
            "tokenizer": "gpt2",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "slumber_genie": None,
            "slumber_model": None,
            "vector_store": vector_store_qdrant,
            "reader": markitdown_reader,
            "client": abstract_qdrant_client,
            "expected": [True, True],
        },
        {
            "chunker": "token",
            "chunk_size": None,
            "chunk_overlap": None,
            "similarity_threshold": None,
            "min_characters_per_chunk": None,
            "min_sentences": None,
            "files_or_dir": [
                "tests/data/test.docx",
                "tests/data/test0.png",
                "tests/data/test1.csv",
                "tests/data/test2.json",
                "tests/data/test3.md",
                "tests/data/test4.xml",
                "tests/data/test5.zip",
            ],
            "tokenizer": "gpt2",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "slumber_genie": None,
            "slumber_model": None,
            "vector_store": vector_store_weaviate,
            "reader": markitdown_reader,
            "client": abstract_weaviate_client,
            "expected": [True, True],
        },
        {
            "chunker": "sdpm",
            "chunk_size": None,
            "chunk_overlap": None,
            "similarity_threshold": None,
            "min_characters_per_chunk": None,
            "min_sentences": None,
            "files_or_dir": [
                "tests/data/test.docx",
                "tests/data/test0.png",
                "tests/data/test1.csv",
                "tests/data/test2.json",
                "tests/data/test3.md",
                "tests/data/test4.xml",
                "tests/data/test5.zip",
            ],
            "tokenizer": "gpt2",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "slumber_genie": None,
            "slumber_model": None,
            "vector_store": vector_store_milvus,
            "reader": None,
            "client": abstract_milvus_client,
            "expected": [True, True],
        },
        {
            "chunker": "sentence",
            "chunk_size": None,
            "chunk_overlap": None,
            "similarity_threshold": None,
            "min_characters_per_chunk": None,
            "min_sentences": None,
            "files_or_dir": [
                "tests/data/test.docx",
                "tests/data/test0.png",
                "tests/data/test1.csv",
                "tests/data/test2.json",
                "tests/data/test3.md",
                "tests/data/test4.xml",
                "tests/data/test5.zip",
            ],
            "tokenizer": "gpt2",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "slumber_genie": None,
            "slumber_model": None,
            "vector_store": vector_store_qdrant,
            "reader": None,
            "client": abstract_qdrant_client,
            "expected": [True, True],
        },
        {
            "chunker": "neural",
            "chunk_size": None,
            "chunk_overlap": None,
            "similarity_threshold": None,
            "min_characters_per_chunk": None,
            "min_sentences": None,
            "files_or_dir": [
                "tests/data/test.docx",
                "tests/data/test0.png",
                "tests/data/test1.csv",
                "tests/data/test2.json",
                "tests/data/test3.md",
                "tests/data/test4.xml",
                "tests/data/test5.zip",
            ],
            "tokenizer": "gpt2",
            "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
            "slumber_genie": None,
            "slumber_model": None,
            "vector_store": vector_store_milvus,
            "reader": None,
            "client": abstract_milvus_client,
            "expected": [True, True],
        },
    ]
    for c in test_cases:
        try:
            ingestor = IngestAnything(
                vector_store=c["vector_store"], reader=c["reader"]
            )
            index = ingestor.ingest(
                chunker=c["chunker"],
                chunk_size=c["chunk_size"],
                chunk_overlap=c["chunk_overlap"],
                similarity_threshold=c["similarity_threshold"],
                min_characters_per_chunk=c["min_characters_per_chunk"],
                min_sentences=c["min_sentences"],
                files_or_dir=c["files_or_dir"],
                tokenizer=c["tokenizer"],
                embedding_model=c["embedding_model"],
                slumber_genie=c["slumber_genie"],
                slumber_model=c["slumber_model"],
            )
            outcome = [
                isinstance(index, VectorStoreIndex),
                c["client"].check_collection_exists(),
            ]
        except Exception as e:
            print(e.__str__(), sys.stderr)
            outcome = [None, e.__str__()]
        for f in [
            "tests/data/test.pdf",
            "tests/data/test0.pdf",
            "tests/data/test1.pdf",
            "tests/data/test2.pdf",
            "tests/data/test3.pdf",
            "tests/data/test4.pdf",
            "tests/data/test5.pdf",
        ]:
            if pathlib.Path(f).is_file():
                os.remove(f)
        assert outcome == c["expected"]


def test_code_ingestion():
    test_cases = [
        {
            "id": 1,
            "files": [
                "tests/code/acronym.go",
                "tests/code/animal_magic.go",
                "tests/code/atbash_cipher_test.go",
            ],
            "language": "go",
            "return_type": None,
            "chunk_size": None,
            "tokenizer": "gpt2",
            "include_nodes": True,
            "vector_store": vector_store_qdrant,
            "client": abstract_qdrant_client,
            "expected": [True, True],
        },
        {
            "id": 2,
            "files": [
                "tests/code/acronym.go",
                "tests/code/animal_magic.go",
                "tests/code/atbash_cipher_test.go",
            ],
            "language": "go",
            "return_type": None,
            "chunk_size": None,
            "tokenizer": "gpt2",
            "include_nodes": True,
            "vector_store": vector_store_weaviate,
            "client": abstract_weaviate_client,
            "expected": [True, True],
        },
        {
            "id": 3,
            "files": [
                "tests/code/acronym.go",
                "tests/code/animal_magic.go",
                "tests/code/atbash_cipher_test.go",
            ],
            "language": "go",
            "return_type": None,
            "chunk_size": None,
            "tokenizer": "gpt2",
            "include_nodes": True,
            "vector_store": vector_store_milvus,
            "client": abstract_milvus_client,
            "expected": [True, True],
        },
        {
            "id": 4,
            "files": [
                "tests/code/acrony.go",
                "tests/code/animal_magc.go",
                "tests/code/atbash_cipher_tes.go",
            ],
            "language": "go",
            "return_type": None,
            "chunk_size": None,
            "tokenizer": "gpt2",
            "include_nodes": True,
            "vector_store": vector_store_qdrant,
            "client": abstract_qdrant_client,
            "expected": None,
        },
        {
            "id": 5,
            "files": [
                "tests/code/acronym.go",
                "tests/code/animal_magic.go",
                "tests/code/atbash_cipher_test.go",
            ],
            "language": "pokemon",
            "return_type": None,
            "chunk_size": None,
            "include_nodes": None,
            "tokenizer": "gpt2",
            "vector_store": vector_store_qdrant,
            "client": abstract_qdrant_client,
            "expected": None,
        },
        {
            "id": 6,
            "files": [
                "tests/code/acronym.go",
                "tests/code/animal_magic.go",
                "tests/code/atbash_cipher_test.go",
            ],
            "language": "python",
            "return_type": "tex",
            "chunk_size": None,
            "tokenizer": "gpt2",
            "include_nodes": None,
            "vector_store": vector_store_qdrant,
            "client": abstract_qdrant_client,
            "expected": None,
        },
    ]
    for c in test_cases:
        try:
            ingestor = IngestCode(vector_store=c["vector_store"])
            index = ingestor.ingest(
                files=c["files"],
                embedding_model="sentence-transformers/all-MiniLM-L6-v2",
                language=c["language"],
                return_type=c["return_type"],
                tokenizer=c["tokenizer"],
                chunk_size=c["chunk_size"],
                include_nodes=c["include_nodes"],
            )
            outcome = [
                isinstance(index, VectorStoreIndex),
                c["client"].check_collection_exists(),
            ]
        except Exception:
            outcome = None
        assert outcome == c["expected"], f"Failed on case id {c['id']}"
