import warnings
import typer
from typing_extensions import Annotated, List
from rag.agents.interface import get_pipeline
import tempfile
import os
from rich import print


# Disable parallelism in the Huggingface tokenizers library to prevent potential deadlocks and ensure consistent behavior.
# This is especially important in environments where multiprocessing is used, as forking after parallelism can lead to issues.
# Note: Disabling parallelism may impact performance, but it ensures safer and more predictable execution.
os.environ['TOKENIZERS_PARALLELISM'] = 'false'


warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)


def run(inputs: Annotated[str, typer.Argument(help="The list of fields to fetch")],
        types: Annotated[str, typer.Argument(help="The list of types of the fields")] = None,
        keywords: Annotated[str, typer.Argument(help="The list of table column keywords")] = None,
        file_path: Annotated[str, typer.Option(help="The file to process")] = None,
        agent: Annotated[str, typer.Option(help="Selected agent")] = "llamaindex",
        index_name: Annotated[str, typer.Option(help="Index to identify embeddings")] = None,
        options: Annotated[List[str], typer.Option(help="Options to pass to the agent")] = None,
        group_by_rows: Annotated[bool, typer.Option(help="Group JSON collection by rows")] = True,
        update_targets: Annotated[bool, typer.Option(help="Update targets")] = True,
        debug: Annotated[bool, typer.Option(help="Enable debug mode")] = False):

    query = 'retrieve ' + inputs
    query_types = types

    query_inputs_arr = [param.strip() for param in inputs.split(',')] if query_types else []
    query_types_arr = [param.strip() for param in query_types.split(',')] if query_types else []
    keywords_arr = [param.strip() for param in keywords.split(',')] if keywords is not None else None

    if not query_types:
        query = inputs

    user_selected_agent = agent  # Modify this as needed

    try:
        rag = get_pipeline(user_selected_agent)
        answer = rag.run_pipeline(user_selected_agent, query_inputs_arr, query_types_arr, keywords_arr, query, file_path,
                                  index_name, options, group_by_rows, update_targets, debug)

        print(f"\nJSON response:\n")
        print(answer)
    except ValueError as e:
        print(f"Caught an exception: {e}")


async def run_from_api_engine(user_selected_agent, query_inputs_arr, query_types_arr, keywords_arr, query, index_name,
                              options_arr, file, group_by_rows, update_targets, debug):
    try:
        rag = get_pipeline(user_selected_agent)

        if file is not None:
            with tempfile.TemporaryDirectory() as temp_dir:
                temp_file_path = os.path.join(temp_dir, file.filename)

                # Save the uploaded file to the temporary directory
                with open(temp_file_path, 'wb') as temp_file:
                    content = await file.read()
                    temp_file.write(content)

                answer = rag.run_pipeline(user_selected_agent, query_inputs_arr, query_types_arr, keywords_arr, query,
                                          temp_file_path, index_name, options_arr, group_by_rows, update_targets,
                                          debug, False)
        else:
            answer = rag.run_pipeline(user_selected_agent, query_inputs_arr, query_types_arr, keywords_arr, query,
                                      None, index_name, options_arr, group_by_rows, update_targets,
                                      debug, False)
    except ValueError as e:
        raise e

    return answer


if __name__ == "__main__":
    typer.run(run)