#!/usr/bin/env python

import os

# Display header
print("==========================================================")
print(" ██████  ██████  ████████ ██   ██  █████  ██      ██      ")
print("██       ██   ██    ██    ██   ██ ██   ██ ██      ██      ")
print("██   ███ ██████     ██    ███████ ███████ ██      ██      ")
print("██    ██ ██         ██         ██ ██   ██ ██      ██      ")
print(" ██████  ██         ██         ██ ██   ██ ███████ ███████ ")
print(" └─> https://github.com/nomic-ai/gpt4all")

# Function to detect macOS architecture and set the binary filename
def detect_mac_arch():
    mac_arch = os.uname().machine
    if mac_arch == "arm64":
        os_type = "M1 Mac/OSX"
        binary_filename = "gpt4all-lora-quantized-OSX-m1"
    elif mac_arch == "x86_64":
        os_type = "Intel Mac/OSX"
        binary_filename = "gpt4all-lora-quantized-OSX-intel"
    else:
        print("Unknown macOS architecture")
        exit(1)
    return os_type, binary_filename

# Detect operating system and set the binary filename
if os.name == "posix":
    if os.uname().sysname == "Darwin":
        os_type, binary_filename = detect_mac_arch()
    elif os.uname().sysname == "Linux":
        if "Microsoft" in open("/proc/version").read():
            os_type = "Windows (WSL)"
            binary_filename = "gpt4all-lora-quantized-win64.exe"
        else:
            os_type = "Linux"
            binary_filename = "gpt4all-lora-quantized-linux-x86"
    else:
        print("Unknown operating system")
        exit(1)
    print("================================")
    print(f"== You are using {os_type}.")

elif os.name == "nt":
    os_type = "Windows (Cygwin/MSYS/MINGW)"
    binary_filename = "gpt4all-lora-quantized-win64.exe"
    print("================================")
    print(f"== You are using {os_type}.")

else:
    print("Unknown operating system")
    exit(1)

# Change to the chat directory
os.chdir("chat")

# List .bin files and prompt user to select one
bin_files = [file for file in os.listdir() if file.endswith(".bin")]
print("== Available .bin files:")
for i, file in enumerate(bin_files):
    print(f"   [{i 1}] {file}")
    
# Function to get user input and validate it
def get_valid_user_input():
    input_valid = False
    while not input_valid:
        print("==> Please enter a number:")
        user_selection = input().strip()
        if user_selection.isdigit() and 1 <= int(user_selection) <= len(bin_files):
            input_valid = True
        else:
            print(f"Invalid input. Please enter a number between 1 and {len(bin_files)}")
    return int(user_selection)

user_selection = get_valid_user_input()
selected_bin_file = bin_files[user_selection-1]

# Run the selected .bin file with the appropriate command
os.system(f"./{binary_filename} -m {selected_bin_file}")