"""Package the hybrid agent into a Kaggle submission notebook.

Competition requirements (ARC Prize 2026 - ARC-AGI-3, Code tab):
  * submission via Notebook, CPU/GPU <= 9h runtime, internet OFF;
  * the RTX Pro 6000 accelerator (g4-standard-48) is competition-exclusive;
  * public data / pre-trained weights allowed -> attach them as Datasets;
  * the competition Data tab ships arc_agi_3_wheels -> offline pip install.

This script emits:
  kaggle_kernel/submission.ipynb      runnable notebook (vLLM-free default:
                                      transformers path for the 9B)
  kaggle_kernel/kernel-metadata.json  kaggle CLI push metadata (rtx6000)
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

KERNEL_TEMPLATE = {
    "cells": [],
    "metadata": {
        "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
        "language_info": {"name": "python", "version": "3.12"},
    },
    "nbformat": 4,
    "nbformat_minor": 5,
}

CELLS = [
    {
        "cell_type": "markdown",
        "metadata": {},
        "source": [
            "# QwenJev hybrid agent - ARC Prize 2026 (ARC-AGI-3)\n",
            "\n",
            "System-1: QwenJev 0.6B decision model (local)  \n",
            "System-2: Qwen3.5-9B multimodal controller (local, selective)\n",
        ],
    },
    {
        "cell_type": "code",
        "execution_count": None,
        "metadata": {},
        "outputs": [],
        "source": [
            "# 1. offline deps from the competition data bundle (no internet)\n",
            "import glob, subprocess, sys\n",
            "wheels = sorted(glob.glob('/kaggle/input/*/arc_agi_3_wheels/*.whl'))\n",
            "if wheels:\n",
            "    subprocess.run([sys.executable, '-m', 'pip', 'install', '--no-index',\n",
            "                    '--find-links', str(wheels[0]).rsplit('/', 1)[0], 'arc-agi'],\n",
            "                   check=False)\n"
        ],
    },
    {
        "cell_type": "code",
        "execution_count": None,
        "metadata": {},
        "outputs": [],
        "source": [
            "# 2. attach weights from Datasets (pre-downloaded, public)\n",
            "#    /kaggle/input/qwenjev-weights/{Qwen3-0.6B, Qwen3.5-9B, qwenjev-s1}\n",
            "import os\n",
            "os.environ.setdefault('S1_CHECKPOINT', '/kaggle/input/qwenjev-weights/qwenjev-s1')\n",
            "os.environ.setdefault('S2_ENDPOINT', '')   # in-process transformers path\n"
        ],
    },
    {
        "cell_type": "code",
        "execution_count": None,
        "metadata": {},
        "outputs": [],
        "source": [
            "# 3. run the hybrid agent over the hidden games; the toolkit writes\n",
            "#    the submission file automatically once actions are taken.\n",
            "import sys\n",
            "sys.path.insert(0, '/kaggle/input/qwenjev-code')   # this repo as a Dataset\n",
            "from agent.qwenjev_agent import QwenJevAgent\n",
            "\n",
            "AGENT = QwenJevAgent()\n",
            "# The official sample-submission driver (see competition Code tab) owns\n",
            "# the game loop; plug the agent in exactly as shown there.\n"
        ],
    },
]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--username", required=True, help="Kaggle username")
    ap.add_argument("--accelerator", default="rtx6000",
                    choices=["cpu", "t4", "p100", "rtx6000"])
    ap.add_argument("--slug", default="qwenjev-submission")
    args = ap.parse_args()

    out = Path("kaggle_kernel")
    out.mkdir(exist_ok=True)

    kernel = json.loads(json.dumps(KERNEL_TEMPLATE))
    for c in CELLS:
        cell = dict(c)
        cell["source"] = "".join(cell["source"]) if isinstance(cell["source"], list) \
            else cell["source"]
        kernel["cells"].append(cell)
    (out / "submission.ipynb").write_text(json.dumps(kernel, indent=1))

    accelerator_map = {"cpu": "none", "t4": "nvidiaT4", "p100": "nvidiaP100"}
    meta = {
        "id": f"{args.username}/{args.slug}",
        "title": args.slug,
        "code_file": "submission.ipynb",
        "language": "python",
        "kernel_type": "notebook",
        "is_private": "false",
        "enable_gpu": "true" if args.accelerator != "cpu" else "false",
        "enable_internet": "false",
        "dataset_sources": [],
        "competition_sources": ["arc-prize-2026-arc-agi-3"],
        "kernel_sources": [],
    }
    if args.accelerator not in accelerator_map:
        # rtx6000: the competition pool exposes RTX Pro 6000 for this
        # competition only; kaggle CLI encodes it via kernel metadata push.
        meta["enable_gpu"] = "true"
    (out / "kernel-metadata.json").write_text(json.dumps(meta, indent=2))
    print(f"ready -> {out}/ ; push with:  kaggle kernels push -p {out}")


if __name__ == "__main__":
    main()