import grpc
from concurrent import futures
import time
import executor_pb2
import executor_pb2_grpc
from agent import cmd
import threading
from store import GlobalStore

VALID_TOKEN = "your-valid-token"

class AuthInterceptor(grpc.ServerInterceptor):

    def __init__(self, valid_token):
        self.valid_token = valid_token

    def intercept_service(self, continuation, handler_call_details):
        # Extract the metadata from the request
        metadata = dict(handler_call_details.invocation_metadata)

        # Check for the token in the metadata
        token = metadata.get('token')
        if token != self.valid_token:
            def denied_handler(request, context):
                context.abort(grpc.StatusCode.UNAUTHENTICATED, 'Invalid token')
            return grpc.unary_unary_rpc_method_handler(denied_handler)

        # If the token is valid, proceed with the request
        return continuation(handler_call_details)

class ExecutorService(executor_pb2_grpc.ExecutorServiceServicer):
    def __init__(self):
        self._store = GlobalStore()

    def CreateExecutor(self, request, context):
        from queue import Queue
        response_queue = Queue()

        # Start the act method in a separate thread
        act_thread = threading.Thread(target=cmd, args=(request.query,response_queue))
        act_thread.start()

        while act_thread.is_alive() or not response_queue.empty():
            if not response_queue.empty():
                message = response_queue.get()
                yield message
            else:
                time.sleep(0.1)  # Avoid busy waiting

    def ExecutorInput(self, request, context):
        id = request.id
        user_input = request.input
        res_id =  self._store.add_value(id, user_input)
        return executor_pb2.ExecutorInputResponse(id=res_id, success=True)

def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10), interceptors=[AuthInterceptor(VALID_TOKEN)])
    executor_pb2_grpc.add_ExecutorServiceServicer_to_server(ExecutorService(), server)
    server.add_insecure_port('[::]:50051')
    server.start()
    try:
        while True:
            time.sleep(86400)
    except KeyboardInterrupt:
        server.stop(0)

if __name__ == '__main__':
    serve()