#!/usr/bin/env python3
"""Fix Colab install command."""
import subprocess, os, json

TOKEN = "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"
subprocess.run(["git", "clone", f"https://{TOKEN}@github.com/ticketguy/littlefig.git", "/app/littlefig"], check=True)
os.chdir("/app/littlefig")
subprocess.run(["git", "config", "user.name", "0xticketguy"], check=True)
subprocess.run(["git", "config", "user.email", "0xticketguy@harboria.dev"], check=True)

with open("Little_Fig_Colab.ipynb", "r") as f:
    colab = json.load(f)

# Fix the install cell
for cell in colab["cells"]:
    if cell["cell_type"] == "code" and "pip install" in "".join(cell["source"]):
        if "little-fig" in "".join(cell["source"]):
            cell["source"] = [
                "# Install Little Fig\n",
                "!pip install -q torch\n",
                "!git clone https://github.com/ticketguy/littlefig.git /content/littlefig --quiet\n",
                "!cd /content/littlefig && pip install -q -e \".[train]\"\n",
                "!pip install -q pyngrok uvicorn fastapi\n",
                "\n",
                "import torch\n",
                "print(f'\\n✅ Ready | PyTorch {torch.__version__} | GPU: {torch.cuda.get_device_name() if torch.cuda.is_available() else \"CPU only\"}')\n",
            ]
            break

# Also fix the server launch to use correct path
for cell in colab["cells"]:
    if cell["cell_type"] == "code" and "uvicorn" in "".join(cell["source"]):
        cell["source"] = [
            "# Launch Little Fig Studio\n",
            "import subprocess, time, sys\n",
            "sys.path.insert(0, '/content/littlefig/src')\n",
            "\n",
            "from pyngrok import ngrok\n",
            "\n",
            "# Start server\n",
            "server = subprocess.Popen(\n",
            "    [sys.executable, '-m', 'uvicorn', 'little_fig.web.server:app',\n",
            "     '--host', '0.0.0.0', '--port', '8888'],\n",
            "    cwd='/content/littlefig/src',\n",
            "    stdout=subprocess.PIPE, stderr=subprocess.PIPE\n",
            ")\n",
            "time.sleep(4)\n",
            "\n",
            "# Tunnel\n",
            "public_url = ngrok.connect(8888)\n",
            "print(f'\\n🍐 Little Fig Studio is LIVE!')\n",
            "print(f'\\n   👉 {public_url}')\n",
            "print(f'\\n   Pick any model. Configure training. Launch.')\n",
            "print(f'   Keep this cell running.')\n",
        ]
        break

with open("Little_Fig_Colab.ipynb", "w") as f:
    json.dump(colab, f, indent=2)

subprocess.run(["git", "add", "-A"], check=True)
subprocess.run(["git", "commit", "-m", "Fix Colab: clone repo + install [train] (embers-diaries not on PyPI)"], check=True)
subprocess.run(["git", "push", "origin", "main"], check=True)
print("✅ Colab install fixed")