import hashlib
from fastapi import FastAPI
import socket
import requests
import json
from multiprocessing import Pool
from functools import partial
import os

from worker import Request


def sha256(text):
    h = hashlib.sha256()
    h.update(text.encode("utf-8"))
    return h.hexdigest()


def get_free_port():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(("localhost", 0))
    _, port = sock.getsockname()
    sock.close()
    return port


api = FastAPI()
pool = None


def fetch_results(data, port):
    headers = {
        "Content-Type": "application/json",
    }
    response = requests.post(f"https://localhost:{port}", headers=headers, data=data, verify=False).json()
    return response


@api.post("/")
async def predict_batch(request: Request):
    global pool
    ports = os.environ["PORTS"].split(",")
    print(f"Using responses from ports {ports}")
    if pool is None:
        pool = Pool(len(ports))
    job = partial(fetch_results, json.dumps(dict(request)))
    responses = pool.map_async(job, ports).get()
    return responses
