import os

import hydra

import aiflows
from aiflows.backends.api_info import ApiInfo
from aiflows.utils.general_helpers import read_yaml_file, quick_load_api_keys

from aiflows import logging
from aiflows.flow_cache import CACHING_PARAMETERS, clear_cache

from aiflows.utils import serving
from aiflows.workers import run_dispatch_worker_thread
from aiflows.messages import FlowMessage
from aiflows.interfaces import KeyInterface
from aiflows.utils.colink_utils import start_colink_server
from aiflows import flow_verse


dependencies = [
    {
        "url": "aiflows/FunSearchFlowModule",
        "revision": os.path.abspath("../")
    }
]
flow_verse.sync_dependencies(dependencies)

logging.set_verbosity_debug()


if __name__ == "__main__":
    
    cl = start_colink_server()
    
    serving.recursive_serve_flow(
        cl=cl,
        flow_class_name="flow_modules.aiflows.FunSearchFlowModule.EvaluatorFlow",
        flow_endpoint="EvaluatorFlow",
    )
        
    run_dispatch_worker_thread(cl)
    
    config_overrides = read_yaml_file(os.path.join(".", "demo.yaml"))

    funsearch_proxy = serving.get_flow_instance(
        cl=cl,
        flow_endpoint="EvaluatorFlow",
        config_overrides=config_overrides,
    )
    data = {
        'artifact': \
            'def solve_function(input) -> str:\n    """Attempt at solving the problem given the input input and returns the predicted output (see the top of the file for problem description)"""\n    return \'YES\'\n'
        }
        
    input_message = funsearch_proxy.package_input_message(data = data)
    
    funsearch_proxy.send_message(input_message)
    
    future = funsearch_proxy.get_reply_future(input_message)
    response = future.get_data()
    print("~~~Response~~~")
    print(response)