from agent_types import FunctionAgent, ReActAgent
from agent import (
    IngestAgent,
    IngestAnythingFunctionAgent,
    IngestAnythingReActAgent,
    IngestCodeFunctionAgent,
    IngestCodeReActAgent,
)
from llama_index.core import Settings
from llama_index.llms.mistralai import MistralAI
import os
from dotenv import load_dotenv
import weaviate
from llama_index.vector_stores.weaviate import WeaviateVectorStore

load_dotenv()

llm = MistralAI(api_key=os.getenv("mistralai_api_key"), model="mistral-small-latest")
Settings.llm = llm

# 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"
)


def addition_tool(a: int, b: int):
    """Tool useful to sum two integer numbers"""
    return a + b


async def multiplication_tool(a: int, b: int):
    """Tool useful to multiply two integer numbers"""
    return a * b


# files to ingest
anything_files = [
    "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",
]
code_files = [
    "tests/code/acronym.go",
    "tests/code/animal_magic.go",
    "tests/code/atbash_cipher_test.go",
]
result_files = [
    "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",
]


def test_everything():
    ## INITIALIZATION TEST ##

    agent_1 = IngestAgent().create_agent(
        vector_database=vector_store_weaviate,
        ingestion_type="anything",
        agent_type="function_calling",
        llm=llm,
        tools=[addition_tool, multiplication_tool],
    )
    agent_2 = IngestAgent().create_agent(
        vector_database=vector_store_weaviate,
        ingestion_type="anything",
        agent_type="react",
        llm=llm,
    )
    agent_3 = IngestAgent().create_agent(
        vector_database=vector_store_weaviate,
        ingestion_type="code",
        agent_type="function_calling",
        llm=llm,
    )
    agent_4 = IngestAgent().create_agent(
        vector_database=vector_store_weaviate,
        ingestion_type="code",
        agent_type="react",
        llm=llm,
        tools=[addition_tool, multiplication_tool],
    )
    assert [
        isinstance(agent_1, IngestAnythingFunctionAgent),
        isinstance(agent_2, IngestAnythingReActAgent),
        isinstance(agent_3, IngestCodeFunctionAgent),
        isinstance(agent_4, IngestCodeReActAgent),
    ] == [True, True, True, True]

    ## AGENTS TEST ##

    # anything agents
    agent_1.ingest(
        files_or_dir=anything_files,
        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
        chunker="token",
        tokenizer="gpt2",
    )
    function_agent_1 = agent_1.get_agent()
    agent_2.ingest(
        files_or_dir=anything_files,
        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
        chunker="token",
        tokenizer="gpt2",
    )
    react_agent_1 = agent_2.get_agent()

    # code agents
    agent_3.ingest(
        files=code_files,
        language="go",
        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
        tokenizer="gpt2",
    )
    function_agent_2 = agent_3.get_agent()
    agent_4.ingest(
        files=code_files,
        language="go",
        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
        tokenizer="gpt2",
    )
    react_agent_2 = agent_4.get_agent()
    assert [
        isinstance(function_agent_1, FunctionAgent),
        isinstance(react_agent_1, ReActAgent),
        isinstance(function_agent_2, FunctionAgent),
        isinstance(react_agent_2, ReActAgent),
    ] == [True, True, True, True]
    for fl in result_files:
        os.remove(fl)
