import os

import hydra

import aiflows
from aiflows.flow_launchers import FlowLauncher
from aiflows.utils.general_helpers import read_yaml_file
from aiflows.backends.api_info import ApiInfo
from aiflows import logging
from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache

CACHING_PARAMETERS.do_caching = False  # Set to True to enable caching
# clear_cache() # Uncomment this line to clear the cache

logging.set_verbosity_debug()

dependencies = [
    {"url": "aiflows/VectorStoreFlowModule", "revision": os.getcwd()},
]
from aiflows import flow_verse
flow_verse.sync_dependencies(dependencies)

if __name__ == "__main__":
    
    # OpenAI backend
    api_information = [ApiInfo(backend_used="openai",
                              api_key = os.getenv("OPENAI_API_KEY"))]
    # Azure backend
    # api_information = ApiInfo(backend_used = "azure",
    #                           api_base = os.getenv("AZURE_API_BASE"),
    #                           api_key = os.getenv("AZURE_OPENAI_KEY"),
    #                           api_version =  os.getenv("AZURE_API_VERSION") )

    root_dir = "."
    cfg_path = os.path.join(root_dir, "demo.yaml")
    cfg = read_yaml_file(cfg_path)
    
    cfg["vector_store_demo_flow"]["subflows_config"]["vs_db"]["backend"]["api_infos"] = api_information
    cfg["chroma_demo_flow"]["subflows_config"]["chroma_db"]["backend"]["api_infos"] = api_information
    

    # ~~~ Get the data ~~~
    # This can be a list of samples
    data = {"id": 0, "operation": "write", "content": "demo of writing"}  # Add your data here

    # ~~~ Run inference ~~~
    path_to_output_file = None
    # path_to_output_file = "output.jsonl"  # Uncomment this line to save the output to disk
    
    #### CHROMA DEMO ####
    ### DUMBY DEMO OF WRITING  "demo of writing" AND READIN ""  (Nothing)###
    print("DEMO: ChromaDBFlow")

    flow_with_interfaces_chroma = {
        "flow": hydra.utils.instantiate(cfg['chroma_demo_flow'], _recursive_=False, _convert_="partial"),
        "input_interface": (
            None
            if getattr(cfg, "input_interface", None) is None
            else hydra.utils.instantiate(cfg['input_interface'], _recursive_=False)
        ),
        "output_interface": (
            None
            if getattr(cfg, "output_interface", None) is None
            else hydra.utils.instantiate(cfg['output_interface'], _recursive_=False)
        ),
    }
    
    _, outputs = FlowLauncher.launch(
        flow_with_interfaces=flow_with_interfaces_chroma,
        data=data,
        path_to_output_file=path_to_output_file,
    )

    # ~~~ Print the output ~~~
    flow_output_data = outputs[0]
    print(flow_output_data)
    
    #### END CHROM DEMO ####
    
    #### VECTOR STORE DEMO ####
    
    print("DEMO: VECTOR STORE DEMO")
     
    flow_with_interfaces_vstore = {
        "flow": hydra.utils.instantiate(cfg['vector_store_demo_flow'], _recursive_=False, _convert_="partial"),
        "input_interface": (
            None
            if getattr(cfg, "input_interface", None) is None
            else hydra.utils.instantiate(cfg['input_interface'], _recursive_=False)
        ),
        "output_interface": (
            None
            if getattr(cfg, "output_interface", None) is None
            else hydra.utils.instantiate(cfg['output_interface'], _recursive_=False)
        ),
    }
    
    
    _, outputs = FlowLauncher.launch(
        flow_with_interfaces=flow_with_interfaces_vstore,
        data=data,
        path_to_output_file=path_to_output_file,
    )

    # ~~~ Print the output ~~~
    flow_output_data = outputs[0]
    print(flow_output_data)