import grpc
import executor_pb2
import executor_pb2_grpc
from colorama import Fore, Style, init

VALID_TOKEN = "your-valid-token"

def run():
    with grpc.insecure_channel('react.pingcap.net:81') as channel:
        metadata = [('token', VALID_TOKEN)]
        stub = executor_pb2_grpc.ExecutorServiceStub(channel)

        while True:
            query = input("Enter a question (or type 'quit' to exit): ")

            if query.lower() == 'quit':
                break

            request = executor_pb2.CreateExecutorRequest(query=query)
            response_iterator = stub.CreateExecutor(request, metadata=metadata)
            for response in response_iterator:
                if response.type == executor_pb2.ACTIONSTART:
                    print(f"{Fore.YELLOW}Start action{Style.RESET_ALL}")
                elif response.type == executor_pb2.ACTIONEND:
                    print(f"{Fore.CYAN}End action{Style.RESET_ALL}")
                elif response.type == executor_pb2.INPUT:
                    user_input = input(f"{Fore.BLUE}Provide input for AI: {Style.RESET_ALL}")
                    input_req = executor_pb2.ExecutorInputRequest(id=response.input_id, input=user_input)
                    input_rep = stub.ExecutorInput(input_req, metadata=metadata)
                    continue
                else:
                    print(f"{Fore.GREEN}Final answer:{Style.RESET_ALL}  {response.answer}\n")

                if response.answer == "":
                    message = response.message.replace("\\n", "\n")
                    print(f"{message}\n")


if __name__ == '__main__':
    run()