import random
from typing import Literal
import time
import sys

def RandomColour(w=5, h=5, mode: Literal['print', 'return'] = 'print'):
    full = []
    for ih in range(h):
        row = []
        for iw in range(w):
            r, g, b = random.randint(0,255), random.randint(0,255), random.randint(0,255)
            row.append(f"\033[48;2;{r};{g};{b}m  \033[0m")
        full.append(''.join(row))
    if mode == 'print':
        print('\n'.join(full))
    elif mode == 'return':
        return '\n'.join(full)

def DiscoDanceFloor(w=5, h=5, frames: Literal['inf']|int = 'inf', fps=5):
    if fps <= 0:
        print("Bro... fps=0 means 0 frames per second. You want a STILL IMAGE?! 🤣")
        return
    if fps > 60:
        print("Bro... 60 FPS is the limit. Your eyes can't see beyond that anyway! 😂")
        fps = 60
    first = RandomColour(w=w, h=h, mode='return')
    print(first)
    
    def next_frame():
        for i in range(h):
            sys.stdout.write(f"\033[2K")
            sys.stdout.write(f"\033[1A")
        sys.stdout.write(f"\033[2K")
        frame = RandomColour(w=w, h=h, mode='return')
        sys.stdout.write(frame + "\n")
        sys.stdout.flush()
        time.sleep(float(1 / fps))

    if frames == 'inf':
        while True:
            next_frame()
    if isinstance(frames, int):
        for i in range(frames):
            next_frame()

RandomColour()
print("\n\n\n")
DiscoDanceFloor(w=10, h=10)