import os
import random
import requests
import base64
import time
from io import BytesIO
import numpy as np

from PIL import Image
import msgpack
from laplace import Client
from io import BytesIO


def select_frames(input_frames, num_segments=10):
    indices = np.linspace(start=0, stop=len(input_frames) - 1, num=num_segments).astype(int)
    frames = [input_frames[ind] for ind in indices]
    return frames

client = Client("sd://data.tns.masp_inf2?cluster=default", timeout=100)

video_dir = './v12044gd0000cl5c6rfog65i2eoqcqig'
frames = [(os.path.splitext(item)[0], os.path.join(video_dir, item)) for item in os.listdir(video_dir)]
frames = [item[1] for item in sorted(frames, key=lambda x: int(x[0]))]

out_frames = select_frames(frames)
request = {}
byte_images = []
for image_path in out_frames:
    img = Image.open(image_path)
    byte_io = BytesIO()
    img.save(byte_io, format='PNG')
    byte_images.append(byte_io.getvalue())

# Step 4: Pack the byte data with msgpack
packed_data = msgpack.packb(byte_images)
request['images'] = [packed_data]

# adjust it if you hope to tune those paramters:
temp = 0.01
# prompt = 'Did the image contains a lion? answer yes or no'
request['temperature'] = [str(temp).encode()]
# request['prompt'] = [prompt.encode()]

tic = time.time()
results = client.matx_inference(model_name="data-tns-masp-inf2", input_lists=request)
toc = time.time()
print('time to do the inference')
print(toc - tic)
# results = client.matx_inference(model_name="data-tns-masp-model-b",input_lists=request)
outputs = results.output_bytes_lists
print(outputs['output'][0])