import base64
from datetime import datetime
import io
import pygame, os
import logging
import time


def capture(pygame_screen, output_dir):
    if output_dir == "":
        pygame.image.save(pygame_screen, "tmp.png")
        filepath = "tmp.png"
    else:
        filepath = os.path.join(output_dir, "screen_" + datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f") + ".png")
        logging.info("filepath: " + filepath)
        print("filepath: " + filepath)
        # 获取当前屏幕的大小
        
        screen_width, screen_height = pygame_screen.get_size()
        image_size_height = 360
        
        image_size = (int(screen_width * image_size_height / screen_height), image_size_height)
        print("image_size: ", image_size)
        
        # 创建一个新表面，并将当前屏幕内容缩放到该表面
        scaled_surface = pygame.transform.scale(pygame_screen, image_size)
        
        if os.path.exists(output_dir):
            # 使用pygame.image.save将缩放后的内容保存到文件
            pygame.image.save(scaled_surface, filepath)
        # pygame.image.save(pygame_screen, filepath)
    
    if os.path.exists(filepath):
        with open(filepath, "rb") as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
    else:
        byte_io = io.BytesIO()
        pygame.image.save(pygame_screen, byte_io) 
        byte_io.seek(0)  # 将文件指针移到开始
        encoded_string = base64.b64encode(byte_io.read()).decode('utf-8')
        
    time.sleep(0.02)
    
    return filepath, encoded_string

def seed_everything(seed):
    try:
        import random
        random.seed(seed)
        print("random seed set to ", seed)
    except:
        print("random seed failed")
    
    try:
        import numpy as np
        np.random.seed(seed)
        print("numpy seed set to ", seed)
    except:
        print("numpy seed failed")
        
    try:
        import torch
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        print("torch seed set to ", seed)
    except:
        print("torch seed failed")