"""Create a portable Colab walkthrough using the standard notebook JSON format."""
import ast
from contextlib import redirect_stdout
import io
import json
from pathlib import Path


def build():
    cells = []
    def markdown(text):
        cells.append({"cell_type": "markdown", "id": f"cell-{len(cells):02}", "metadata": {}, "source": text})
    def code(text):
        cells.append({"cell_type": "code", "id": f"cell-{len(cells):02}", "metadata": {},
                      "source": text, "execution_count": None, "outputs": []})
    markdown("""# PocketVision: cloud training for chat, code and writing

## Goal
Run a small QLoRA smoke experiment on a **Colab cloud GPU**, then inspect a candidate before scaling up.
The PC only displays this notebook. No trained 4B weights are included.

**Validation status:** the saved outputs show a local dry run only. Cloud installation, source authentication,
GPU memory, training, and model quality remain untested. Free Colab GPU availability is not guaranteed.
Choose a free runtime yourself; this notebook does not purchase compute or alter subscriptions.

English Wikipedia restricts AI-generated article prose. Use the reference-article task for other contexts,
and the outline/source-review task for Wikipedia preparation.
[Current guideline](https://en.wikipedia.org/wiki/Wikipedia:Writing_articles_with_large_language_models).
""")
    markdown("""## Setup
1. Upload this notebook to Colab and select a GPU runtime.
2. Set `DRY_RUN = False` below. Set `RUN_TRAINING = True` only when ready to use the allocated GPU.
3. Store an existing Hugging Face **read token** in Colab Secrets as `HF_TOKEN` and grant this notebook access.
   The source repository is private. Do not paste a token into a code cell or chat.
4. The Drive step asks you to mount Google Drive. This grants notebook code access to the mounted files;
   only the named PocketVision folder is used here for checkpoints.
""")
    code('''from pathlib import Path
import json
import os
import subprocess
import sys

DRY_RUN = True
RUN_TRAINING = False
SOURCE_REPO = "gamedevelop/pocketvision-chat-code"
SOURCE_REVISION = None  # None resolves and records one immutable commit before downloading.
PILOT_ROWS = 200        # Dependency/memory smoke test. Use 20000 only after reviewing the smoke run.
HOLDOUT_PER_SOURCE = 2
MAX_STEPS = 25
RUN_NAME = "writing-smoke-01"  # Change for a fresh run; existing output is never overwritten.
RESUME_CHECKPOINT = None       # Or an existing checkpoint-N directory from the same run.
cloud_python = Path("/content/pocketvision-venv/bin/python")
source_dir = None
run_dir = None
print("Dry-run validation only; no cloud operations." if DRY_RUN else "Cloud setup enabled; training still needs RUN_TRAINING=True.")
''')
    markdown("""### 1. Install an isolated cloud environment
Installation uses a separate environment, preserving the notebook kernel's packages.
The next check requires a CUDA GPU before any model weights are loaded.
""")
    code('''if not DRY_RUN:
    subprocess.run([sys.executable, "-m", "venv", str(cloud_python.parent.parent)], check=True)
    subprocess.run([str(cloud_python), "-m", "pip", "install", "torch==2.9.1",
                    "--index-url", "https://download.pytorch.org/whl/cu128"], check=True)
    subprocess.run([sys.executable, "-m", "pip", "install", "huggingface-hub==0.36.2"], check=True)
    subprocess.run([str(cloud_python), "-c",
                    "import torch; assert torch.cuda.is_available(), 'Select a Colab GPU runtime'; print(torch.cuda.get_device_name(0))"], check=True)
else:
    print("Skipped: dependency installation and GPU allocation checks.")
''')
    markdown("""### 2. Download only the versioned source
The commit is resolved once and logged; pretrained weights and source ZIPs are excluded.
The model and datasets have their own pinned revisions in the configuration.
""")
    code('''if not DRY_RUN:
    from google.colab import userdata
    from huggingface_hub import HfApi, snapshot_download
    os.environ["HF_TOKEN"] = userdata.get("HF_TOKEN")
    revision = HfApi().model_info(SOURCE_REPO, revision=SOURCE_REVISION or "main").sha
    source_dir = Path(snapshot_download(SOURCE_REPO, revision=revision, allow_patterns=[
        "app.py", "chat.py", "LICENSE*", "chatcode/*.py", "chatcode/*.json", "chatcode/*.txt", "chatcode/*.md"]))
    assert (source_dir / "chatcode/articles.py").is_file(), "This source commit predates article support."
    subprocess.run([str(cloud_python), "-m", "pip", "install", "-r",
                    str(source_dir / "chatcode/requirements-cloud.txt")], check=True)
    subprocess.run([str(cloud_python), "-m", "pip", "check"], check=True)
    print("Source commit:", revision)
else:
    print("Skipped: private Hub download and authentication. No token requested.")
''')
    markdown("""### 3. Use persistent storage
Colab's runtime disk is temporary. Drive contains the data manifest, optimizer checkpoints and candidate.
Only this new run folder is created; no existing files are deleted.
""")
    code('''if not DRY_RUN:
    from google.colab import drive
    drive.mount("/content/drive")
    checkpoint_root = Path("/content/drive/MyDrive/PocketVision-writing")
    run_dir = checkpoint_root / RUN_NAME
    if run_dir.exists() and any(run_dir.iterdir()) and not RESUME_CHECKPOINT:
        raise FileExistsError("Choose a fresh RUN_NAME or explicitly resume an existing checkpoint.")
    run_dir.mkdir(parents=True, exist_ok=True)
    (run_dir / "run-info.json").write_text(json.dumps({"source_repo": SOURCE_REPO,
        "source_revision": revision, "pilot_rows": PILOT_ROWS, "max_steps": MAX_STEPS}, indent=2))
    print("Persistent run directory:", run_dir)
else:
    print("Skipped: Drive authorization and persistent cloud writes.")
''')
    markdown("""## Steps
### 4. Prepare data and run the capped experiment
This cell streams selected chat, coding, rewriting and summarization data, adds the small authored behavior
packets, and runs 25 optimizer steps. It may consume significant GPU time even on a small sample.
The full 20,000-example experiment is a separate decision after measuring steps/second and memory.
""")
    code('''if not DRY_RUN and RUN_TRAINING:
    data_dir = run_dir / "data"
    if not (data_dir / "manifest.json").is_file():
        subprocess.run([str(cloud_python), "-m", "chatcode.data", "--output", str(data_dir),
                        "--train", str(PILOT_ROWS), "--holdout-per-source", str(HOLDOUT_PER_SOURCE)], cwd=source_dir, check=True)
    command = [str(cloud_python), "-m", "chatcode.train", "--data", str(data_dir),
               "--output", str(run_dir / "training"), "--max-steps", str(MAX_STEPS)]
    if RESUME_CHECKPOINT:
        command.extend(["--resume-from-checkpoint", str(RESUME_CHECKPOINT)])
    subprocess.run(command, cwd=source_dir, check=True)
else:
    print("Skipped: dataset preparation and weight training. No model weights downloaded.")
''')
    markdown("""## Checks
### 5. Compare base and candidate
The automated article check measures JSON/citation-ID validity on four held-out fictional cases.
It does **not** measure factual accuracy, source reliability, or Wikipedia suitability. Inspect all saved
outputs, and separately test chat and code regressions. Do not execute generated programs outside a sandbox.
""")
    code('''if not DRY_RUN and RUN_TRAINING:
    candidate_dir = run_dir / "training/adapter"
    subprocess.run([str(cloud_python), "-m", "chatcode.evaluate_articles", "--adapter", str(candidate_dir),
                    "--output", str(run_dir / "article-evaluation.json")], cwd=source_dir, check=True)
    print("Candidate and evaluation saved under:", run_dir)
else:
    print("Skipped: model evaluation. No accuracy or quality result is claimed.")
''')
    markdown("""## Next Steps
Review the cloud smoke run before raising `PILOT_ROWS`, holdout sizes or run duration. Training data changes
require a new run directory; do not resume a checkpoint against different data.
The source recipe alone does not improve weights. Compare factual support and neutrality with the base,
and check programming and conversation regressions before publishing a candidate.
No upload or Wikipedia publication is performed by this notebook. A later Hub upload needs a separate
write-scoped credential and a reviewed candidate. Do not share notebook outputs containing credentials.
""")
    notebook = {"nbformat": 4, "nbformat_minor": 5,
                "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
                             "language_info": {"name": "python", "version": "3.12"},
                             "pocketvision_validation": "local_dry_run_only_no_cloud_execution"}, "cells": cells}
    return notebook


def validate_and_dry_run(notebook):
    """Structural and Python syntax checks when nbformat is unavailable locally."""
    assert notebook["nbformat"] == 4 and notebook["nbformat_minor"] == 5
    namespace, ids, count = {}, set(), 0
    for cell in notebook["cells"]:
        assert cell["id"] not in ids and isinstance(cell["metadata"], dict) and isinstance(cell["source"], str)
        ids.add(cell["id"])
        assert cell["cell_type"] in {"markdown", "code"}
        if cell["cell_type"] == "code":
            compiled = compile(ast.parse(cell["source"]), "<notebook-cell>", "exec")
            output = io.StringIO()
            with redirect_stdout(output):
                exec(compiled, namespace)
            count += 1
            cell["execution_count"] = count
            cell["outputs"] = [{"output_type": "stream", "name": "stdout", "text": output.getvalue()}]
    assert namespace["DRY_RUN"] is True and namespace["RUN_TRAINING"] is False
    return notebook


if __name__ == "__main__":
    path = Path(__file__).with_name("cloud_training.ipynb")
    path.write_text(json.dumps(validate_and_dry_run(build()), ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Wrote {path}; all cells executed only in dry-run mode. GPU execution remains unverified.")