import gradio as gr

import os
from IPython import get_ipython
from IPython.display import display, Markdown

COLAB = False

if COLAB:
  from google.colab.output import clear as clear_output
else:
  from IPython.display import clear_output

#@title ## 🚩 Start Here

#@markdown ### 1️⃣ Setup
#@markdown This cell will load some requirements and create the necessary folders in your Google Drive. <p>
#@markdown Your project name can't contain spaces but it can contain a single / to make a subfolder in your dataset.
project_name = "LillieBlueDress" #@param {type:"string"}
project_name = project_name.strip()
#@markdown The folder structure doesn't matter and is purely for comfort. Make sure to always pick the same one. I like organizing by project.
folder_structure = "Organize by project (MyDrive/Loras/project_name/dataset)" #@param ["Organize by category (MyDrive/lora_training/datasets/project_name)", "Organize by project (MyDrive/Loras/project_name/dataset)"]

if not project_name or any(c in project_name for c in " .()\"'\\") or project_name.count("/") > 1:
  print("Please write a valid project_name, project name cannot have.()\"'\\.")
else:
  if COLAB and not os.path.exists('/content/drive'):
    from google.colab import drive
    print("📂 Connecting to Google Drive...")
    drive.mount('/content/drive')

  project_base = project_name if "/" not in project_name else project_name[:project_name.rfind("/")]
  project_subfolder = project_name if "/" not in project_name else project_name[project_name.rfind("/")+1:]

  root_dir = "/content" if COLAB else "~/Loras"
  deps_dir = os.path.join(root_dir, "deps")
  
  if "/Loras" in folder_structure:
    main_dir      = os.path.join(root_dir, "drive/MyDrive/Loras") if COLAB else root_dir
    config_folder = os.path.join(main_dir, project_base)
    images_folder = os.path.join(main_dir, project_base, "dataset")
    if "/" in project_name:
      images_folder = os.path.join(images_folder, project_subfolder)
  else:
    main_dir      = os.path.join(root_dir, "drive/MyDrive/lora_training") if COLAB else root_dir
    config_folder = os.path.join(main_dir, "config", project_name)
    images_folder = os.path.join(main_dir, "datasets", project_name)

  for dir in [main_dir, deps_dir, images_folder, config_folder]:
    os.makedirs(dir, exist_ok=True)

  print(f"✅ Project {project_name} is ready!")
  step1_installed_flag = True