# --------------------------------------------------------------------------------
# COPY AND PASTE THIS CELL INTO YOUR KAGGLE NOTEBOOK
# --------------------------------------------------------------------------------

import os
import shutil

# --- 1. SETUP PATHS ---
# We look for your specific dataset name in the input directory
expected_name_part = "assyrian-platinum-adapter-auto"
input_root = "/kaggle/input/"

# Find the exact folder name (e.g., "assyrian-translation-adapter-final-v1")
try:
    dataset_folder = [d for d in os.listdir(input_root) if expected_name_part in d.lower()][0]
    dataset_path = os.path.join(input_root, dataset_folder)
    print(f"FOUND DATASET: {dataset_path}")
except IndexError:
    print(f"ERROR: Could not find dataset containing '{expected_name_part}' in {input_root}")
    print("Available inputs:", os.listdir(input_root))
    raise

# --- 2. COPY SCRIPT ---
# Move generate_submission.py to the writable directory
script_src = os.path.join(dataset_path, "generate_submission.py")
if os.path.exists(script_src):
    shutil.copy(script_src, "./generate_submission.py")
    print("Script copied to working directory.")
else:
    print("ERROR: generate_submission.py not found in dataset!")
    print("Contents:", os.listdir(dataset_path))

# --- 3. SETUP ADAPTER ---
# We need the adapter files in "./assyrian_byt5_qlora_v2_refined"
# This logic handles BOTH cases:
# A) You zipped the *folder* (so it exists inside the dataset)
# B) You zipped the *files* (so they are loose in the root)

target_folder = "./assyrian_byt5_qlora_v2_refined"
possible_source_folder = os.path.join(dataset_path, "assyrian_byt5_qlora_v2_refined")

if not os.path.exists(target_folder):
    if os.path.exists(possible_source_folder):
        # Scenario A: The folder exists
        print(f"Copying adapter folder from {possible_source_folder}...")
        shutil.copytree(possible_source_folder, target_folder)
    else:
        # Scenario B: Loose files
        print("Adapter folder not found, assuming loose files. Creating folder...")
        os.makedirs(target_folder, exist_ok=True)
        # Copy all model-related files
        for file in os.listdir(dataset_path):
            if file.endswith(".json") or file.endswith(".safetensors") or file.endswith(".bin"):
                shutil.copy(os.path.join(dataset_path, file), target_folder)
    print("Adapter setup complete.")

# --- 4. INSTALL & RUN ---
import subprocess
print("Installing dependencies and running inference...")
subprocess.check_call(["pip", "install", "-q", "peft", "transformers", "bitsandbytes", "sentencepiece", "pandas"])
subprocess.check_call(["python", "-u", "generate_submission.py"])