# WarServer - replace HOST with 'localhost' for local testing
import socket
import WarBot

# Kill all warnings
import warnings
warnings.filterwarnings("ignore")

imgWord = 'как выглядит'

# Initialize the base models
model,tokenizer,model_punct,translation_model,translation_tokenizer,imgModel_version = WarBot.initialize()

HOST = '10.0.0.125'
PORT = 5000

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
    # Server sockets operation
    server_socket.bind((HOST, PORT))
    server_socket.listen()
    print(f'Server is listening on port {PORT}')
    while True:
        conn, addr = server_socket.accept()
        with conn:
            print(f'Connected by {addr}')
            data = conn.recv(1024)
            received_string = data.decode()
            #received_string = data.decode('utf-8')

            print(f'Received string from client: {received_string}')

            response = ""

            while not response:
                if received_string.lower().startswith(imgWord): # check if that is a call for an image:
                    received_string = received_string.lower().split(imgWord)[1:][0].strip()  # cut the code word for image
                    # Translate it to english:
                    translated_string = WarBot.translate(received_string, translation_model=translation_model,
                                                         translation_tokenizer=translation_tokenizer)
                    # Generated image url
                    response = WarBot.generate_image(prompt=translated_string, imgModel_version=imgModel_version)
                else:
                    response = WarBot.get_response(received_string, model, tokenizer, model_punct, temperature=0.6)

            response_string = response

            conn.sendall(response_string.encode())
            conn.close()