
import os, json
from tqdm import tqdm

def list_gguf_files():
    files = [f for f in os.listdir() if f.endswith(".gguf")]
    if not files:
        print("No .gguf files found in the current directory.")
        return None
    print("Available .gguf files:")
    for i, f in enumerate(files, 1):
        print(f"{i}. {f}")
    while True:
        try:
            choice = int(input("Select a file by number: "))
            return files[choice - 1]
        except (ValueError, IndexError):
            print("Invalid selection. Try again.")

def get_cutoff_size():
    while True:
        try:
            mb = float(input("Enter cutoff size for each part (in MB): "))
            return int(mb * 1024 * 1024)
        except ValueError:
            print("Please enter a valid number.")

def split_gguf_file(filename, cutoff_size):
    with open(filename, "rb") as f:
        data = f.read()

    total_size = len(data)
    parts = []
    offset = 0
    part_id = 1

    print("\nSplitting file into parts...")
    with tqdm(total=total_size, unit='B', unit_scale=True, desc="Splitting") as pbar:
        while offset < total_size:
            end = min(offset + cutoff_size, total_size)
            part_data = data[offset:end]
            part_name = f"model-{part_id:05d}-of-XXXX.gguf"
            with open(part_name, "wb") as pf:
                pf.write(part_data)
            parts.append(part_name)
            offset = end
            part_id += 1
            pbar.update(len(part_data))

    num_parts = len(parts)
    final_files = []

    print("\nRenaming parts with total count info...")
    for i, old_name in tqdm(list(enumerate(parts)), desc="Renaming", unit="file"):
        new_name = f"model-{i+1:05d}-of-{num_parts:05d}.gguf"
        os.rename(old_name, new_name)
        final_files.append(new_name)

    index = {
        "metadata": {
            "total_size": total_size
        },
        "file_map": {
            f"part_{i+1}": name for i, name in enumerate(final_files)
        }
    }

    with open("model.gguf.index.json", "w") as f:
        json.dump(index, f, indent=2)

    print(f"\nSplit complete! {num_parts} parts created.")
    print("Index file: model.gguf.index.json")

file = list_gguf_files()
if file:
    cutoff = get_cutoff_size()
    split_gguf_file(file, cutoff)