#!/usr/bin/env bash
# Phase-0 go/no-go for Qwen3.6-27B on the GPU host:
#   1) config + tokenizer (does transformers know qwen3_5?)
#   2) full bf16 load + a thinking-mode generation (does it load + forward?)
#   3) a tiny LoRA training step (does the LoRA training path run?)
#
# Writes machine-readable results to reports/cybergym/qwen36_27b_base/phase0_env.md
# and per-stage logs under /workspace/tmp. Exits non-zero if the load stage fails.
set -uo pipefail

MODEL="${MODEL:-Qwen/Qwen3.6-27B}"
export HF_HOME=${HF_HOME:-/workspace/hf-cache}
export HF_HUB_CACHE=${HF_HUB_CACHE:-/workspace/hf-cache/hub}
export PYTHONUNBUFFERED=1
LOG=/workspace/tmp/phase0
mkdir -p "$LOG" reports/cybergym/qwen36_27b_base
REPORT=reports/cybergym/qwen36_27b_base/phase0_env.md

echo "# Phase-0 Environment Report" > "$REPORT"
echo "" >> "$REPORT"
echo "- Model: \`$MODEL\`" >> "$REPORT"

echo "===== [1/3] config check =====" | tee "$LOG/config.log"
python training/scripts/phase0_qwen_smoke.py --model "$MODEL" --mode config 2>&1 | tee -a "$LOG/config.log"
CFG_RC=${PIPESTATUS[0]}
echo "" >> "$REPORT"; echo "## Config check (rc=$CFG_RC)" >> "$REPORT"
echo '```json' >> "$REPORT"; tail -n 40 "$LOG/config.log" >> "$REPORT"; echo '```' >> "$REPORT"

echo "===== [2/3] full bf16 load + generate =====" | tee "$LOG/load.log"
python training/scripts/phase0_qwen_smoke.py --model "$MODEL" --mode load --dtype bfloat16 --device-map auto 2>&1 | tee -a "$LOG/load.log"
LOAD_RC=${PIPESTATUS[0]}
echo "" >> "$REPORT"; echo "## Full load + generate (rc=$LOAD_RC)" >> "$REPORT"
echo '```json' >> "$REPORT"; tail -n 60 "$LOG/load.log" >> "$REPORT"; echo '```' >> "$REPORT"

if [[ "$LOAD_RC" != "0" ]]; then
  echo "" >> "$REPORT"
  echo "**GO/NO-GO: NO-GO on load.** Consider the Qwen3-32B fallback or newer transformers/kernels." >> "$REPORT"
  echo "PHASE0 LOAD FAILED (rc=$LOAD_RC). See $LOG/load.log" >&2
  exit "$LOAD_RC"
fi

echo "===== [3/3] tiny LoRA train step =====" | tee "$LOG/train.log"
# minimal smoke dataset (assistant turns carry <think>) + config
mkdir -p data/processed data/decontam
cat > data/processed/_smoke_train.jsonl <<'JSONL'
{"id":"s1","messages":[{"role":"system","content":"Authorized security research and education context."},{"role":"user","content":"Is strcpy into a fixed buffer safe?"},{"role":"assistant","content":"<think>strcpy does no bounds checking; a long source overflows a fixed buffer.</think>\n\nNo. It can overflow the destination (CWE-787). Use strncpy/snprintf with the buffer size."}]}
{"id":"s2","messages":[{"role":"system","content":"Authorized security research and education context."},{"role":"user","content":"What does use-after-free mean?"},{"role":"assistant","content":"<think>Accessing memory after free leaves a dangling pointer; reuse leads to corruption.</think>\n\nUsing a pointer after its memory was freed (CWE-416); it can crash or be exploited."}]}
JSONL
cp data/processed/_smoke_train.jsonl data/processed/_smoke_val.jsonl

cat > training/configs/_phase0_smoke.yaml <<'YAML'
run: {name: phase0_smoke, stage: phase0_smoke, output_dir: /workspace/checkpoints/_phase0_smoke, seed: 1337}
model:
  name_or_path: Qwen/Qwen3.6-27B
  trust_remote_code: true
  dtype: bfloat16
  attn_implementation: sdpa
  freeze_vision_tower: true
  freeze_mtp_head: true
data:
  train_jsonl: data/processed/_smoke_train.jsonl
  validation_jsonl: data/processed/_smoke_val.jsonl
  max_seq_length: 1024
  packing: false
  require_think_blocks: true
training:
  method: lora
  lora: {r: 8, alpha: 16, dropout: 0.0, use_rslora: true, target_modules: all-linear, exclude_modules: [visual, vision_tower, multi_modal_projector, mtp, lm_head]}
  optimizer: paged_adamw_8bit
  learning_rate: 0.0002
  lr_scheduler_type: constant
  warmup_ratio: 0.0
  num_train_epochs: 1
  per_device_train_batch_size: 1
  gradient_accumulation_steps: 1
  gradient_checkpointing: true
  logging_steps: 1
  bf16: true
  tf32: true
YAML

python training/scripts/run_sft.py --config training/configs/_phase0_smoke.yaml --allow-missing-cybergym-baseline 2>&1 | tee -a "$LOG/train.log"
TRAIN_RC=${PIPESTATUS[0]}
echo "" >> "$REPORT"; echo "## Tiny LoRA train step (rc=$TRAIN_RC)" >> "$REPORT"
echo '```' >> "$REPORT"; tail -n 40 "$LOG/train.log" >> "$REPORT"; echo '```' >> "$REPORT"

echo "" >> "$REPORT"
if [[ "$TRAIN_RC" == "0" ]]; then
  echo "**GO/NO-GO: GO.** Loads, generates with thinking, and LoRA-trains on this stack." >> "$REPORT"
else
  echo "**GO/NO-GO: PARTIAL.** Loads + generates, but the LoRA train smoke failed (rc=$TRAIN_RC) — inspect $LOG/train.log." >> "$REPORT"
fi
echo "Phase-0 report written to $REPORT"
exit 0