import asyncio, time
from playwright.async_api import async_playwright

CELLS = [
    "!pip install -q unsloth datasets trl peft accelerate bitsandbytes",
    "!wget -q 'https://huggingface.co/datasets/yace222/weval-brain-dataset/resolve/main/train_chatml.jsonl' -O /content/dataset.jsonl\n!wc -l /content/dataset.jsonl",
    "from unsloth import FastLanguageModel\nimport torch\nmodel, tokenizer = FastLanguageModel.from_pretrained(model_name='unsloth/Qwen2.5-7B-Instruct-bnb-4bit', max_seq_length=2048, load_in_4bit=True)\nmodel = FastLanguageModel.get_peft_model(model, r=16, target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'], lora_alpha=16, lora_dropout=0, bias='none', use_gradient_checkpointing='unsloth')\nprint('Model loaded!')",
    "from datasets import load_dataset\ndataset = load_dataset('json', data_files='/content/dataset.jsonl', split='train')\ndef format_chat(example):\n    text = tokenizer.apply_chat_template(example['messages'], tokenize=False, add_generation_prompt=False)\n    return {'text': text}\ndataset = dataset.map(format_chat)\nprint(f'Formatted: {len(dataset)} examples')",
    "from trl import SFTTrainer\nfrom transformers import TrainingArguments\ntrainer = SFTTrainer(model=model, tokenizer=tokenizer, train_dataset=dataset, dataset_text_field='text', max_seq_length=2048, args=TrainingArguments(per_device_train_batch_size=2, gradient_accumulation_steps=4, warmup_steps=5, num_train_epochs=3, learning_rate=2e-4, fp16=not torch.cuda.is_bf16_supported(), bf16=torch.cuda.is_bf16_supported(), logging_steps=1, output_dir='outputs', optim='adamw_8bit', seed=42))\nstats = trainer.train()\nprint(f'Done! Loss: {stats.training_loss:.4f}')",
    "model.save_pretrained_gguf('weval-brain-gguf', tokenizer, quantization_method='q4_k_m')\nprint('GGUF Q4_K_M saved!')\nimport os\nfor f in os.listdir('weval-brain-gguf'):\n    print(f'{f}: {os.path.getsize(os.path.join(\"weval-brain-gguf\",f))//1024//1024}MB')",
]

async def run():
    async with async_playwright() as p:
        ctx = await p.chromium.launch_persistent_context(
            "/tmp/chrome-weval", headless=True,
            args=["--no-sandbox","--disable-blink-features=AutomationControlled"],
            viewport={"width":1920,"height":1080},
        )
        page = ctx.pages[0] if ctx.pages else await ctx.new_page()
        await page.add_init_script("Object.defineProperty(navigator,'webdriver',{get:()=>undefined})")
        
        # Open blank Colab
        print(f"[{time.strftime('%H:%M:%S')}] Opening Colab...")
        await page.goto("https://colab.research.google.com/#create=true", timeout=30000)
        await page.wait_for_timeout(10000)
        
        # Try setting GPU via menu
        print(f"[{time.strftime('%H:%M:%S')}] Setting GPU runtime...")
        # Use Command Palette: Ctrl+Shift+P
        await page.keyboard.press("Control+Shift+p")
        await page.wait_for_timeout(2000)
        await page.keyboard.type("change runtime", delay=80)
        await page.wait_for_timeout(2000)
        await page.keyboard.press("Enter")
        await page.wait_for_timeout(3000)
        await page.screenshot(path="/tmp/colab-gpu-dialog.png")
        
        # Try selecting T4 GPU
        try:
            # Look for GPU option in dialog
            t4_option = page.locator("text=T4")
            if await t4_option.count() > 0:
                await t4_option.first.click()
                await page.wait_for_timeout(1000)
                print("T4 GPU selected!")
            
            # Click Save
            save_btn = page.locator("text=Save, button:has-text('Save')")
            if await save_btn.count() > 0:
                await save_btn.first.click()
                await page.wait_for_timeout(2000)
                print("GPU saved!")
        except Exception as e:
            print(f"GPU setup: {e}")
        
        await page.keyboard.press("Escape")
        await page.wait_for_timeout(1000)
        
        # Paste all cells
        print(f"[{time.strftime('%H:%M:%S')}] Pasting {len(CELLS)} cells...")
        
        # Click on first cell
        cells = page.locator(".cell, .codecell, [class*='cell']")
        first_cell = page.locator("textarea, .CodeMirror, .monaco-editor, [contenteditable]").first
        
        for i, code in enumerate(CELLS):
            if i > 0:
                # Add new cell: Ctrl+M then B
                await page.keyboard.press("Control+m")
                await page.wait_for_timeout(300)
                await page.keyboard.press("b")
                await page.wait_for_timeout(1000)
            
            # Type code into cell
            await page.keyboard.type(code, delay=2)
            await page.wait_for_timeout(300)
            print(f"  Cell {i+1}/{len(CELLS)} pasted ({len(code)} chars)")
        
        await page.screenshot(path="/tmp/colab-all-cells.png")
        
        # RUN ALL: Ctrl+F9
        print(f"[{time.strftime('%H:%M:%S')}] Running all cells (Ctrl+F9)...")
        await page.keyboard.press("Control+F9")
        await page.wait_for_timeout(10000)
        await page.screenshot(path="/tmp/colab-executing.png")
        
        print(f"[{time.strftime('%H:%M:%S')}] TRAINING LAUNCHED!")
        print("Cells are executing. Training will take ~20 min on T4 GPU.")
        print("Check: /tmp/colab-executing.png")
        
        # Copy screenshots
        import subprocess
        for f in ['colab-gpu-dialog.png','colab-all-cells.png','colab-executing.png']:
            subprocess.run(['cp', f'/tmp/{f}', f'/var/www/html/api/notebooks/{f}'], capture_output=True)
        
        await ctx.close()

asyncio.run(run())