import os
import sys
import json
import time
import uuid
import shlex
import signal
import shutil
import socket
import secrets
import pathlib
import subprocess
from datetime import datetime
RESET_QWENPAW = False
PORT = int(os.environ.get("QWENPAW_COLAB_PORT", "8088"))
ROOT = pathlib.Path("/content/qwenpaw_colab")
WORKING_DIR = ROOT / "working"
SECRET_DIR = ROOT / "secrets"
LOG_DIR = ROOT / "logs"
WORKSPACE_DIR = WORKING_DIR / "workspaces" / "default"
PID_FILE = ROOT / "qwenpaw_app.pid"
APP_LOG = LOG_DIR / "qwenpaw_app.log"
if RESET_QWENPAW and ROOT.exists():
   shutil.rmtree(ROOT)
for p in [ROOT, WORKING_DIR, SECRET_DIR, LOG_DIR, WORKSPACE_DIR]:
   p.mkdir(parents=True, exist_ok=True)
os.environ["QWENPAW_WORKING_DIR"] = str(WORKING_DIR)
os.environ["QWENPAW_SECRET_DIR"] = str(SECRET_DIR)
os.environ["QWENPAW_AUTH_ENABLED"] = "true"
os.environ["QWENPAW_AUTH_USERNAME"] = os.environ.get("QWENPAW_AUTH_USERNAME", "admin")
os.environ["QWENPAW_LOG_LEVEL"] = os.environ.get("QWENPAW_LOG_LEVEL", "info")
os.environ["QWENPAW_SKILL_SCAN_MODE"] = os.environ.get("QWENPAW_SKILL_SCAN_MODE", "warn")
os.environ["QWENPAW_TOOL_GUARD_ENABLED"] = os.environ.get("QWENPAW_TOOL_GUARD_ENABLED", "true")
password_file = SECRET_DIR / ".colab_ui_password"
if not password_file.exists():
   password_file.write_text("qpw-" + secrets.token_urlsafe(18), encoding="utf-8")
os.environ["QWENPAW_AUTH_PASSWORD"] = password_file.read_text(encoding="utf-8").strip()
def run(cmd, check=False, env=None, cwd=None, stream=False):
   if isinstance(cmd, str):
       display_cmd = cmd
       shell = True
   else:
       display_cmd = " ".join(shlex.quote(str(x)) for x in cmd)
       shell = False
   print(f"\n$ {display_cmd}")
   if stream:
       proc = subprocess.Popen(cmd, shell=shell, env=env, cwd=cwd, text=True)
       rc = proc.wait()
       if check and rc != 0:
           raise RuntimeError(f"Command failed with exit code {rc}: {display_cmd}")
       return rc, ""
   out = subprocess.run(
       cmd,
       shell=shell,
       env=env,
       cwd=cwd,
       text=True,
       stdout=subprocess.PIPE,
       stderr=subprocess.STDOUT,
   )
   print(out.stdout[-4000:])
   if check and out.returncode != 0:
       raise RuntimeError(f"Command failed with exit code {out.returncode}: {display_cmd}")
   return out.returncode, out.stdout
def port_open(host="127.0.0.1", port=8088, timeout=0.5):
   try:
       with socket.create_connection((host, port), timeout=timeout):
           return True
   except OSError:
       return False
def wait_for_port(port, seconds=90):
   start = time.time()
   while time.time() - start < seconds:
       if port_open("127.0.0.1", port):
           return True
       time.sleep(1)
   return False
def stop_previous_app():
   if PID_FILE.exists():
       try:
           pid = int(PID_FILE.read_text().strip())
           os.kill(pid, signal.SIGTERM)
           time.sleep(2)
           try:
               os.kill(pid, 0)
               os.kill(pid, signal.SIGKILL)
           except OSError:
               pass
       except Exception:
           pass
       PID_FILE.unlink(missing_ok=True)
def qwenpaw_cmd(*args):
   exe = shutil.which("qwenpaw")
   if exe:
       return [exe, *args]
   return [sys.executable, "-m", "qwenpaw", *args]
def colab_secret_or_env(name):
   value = os.environ.get(name, "")
   try:
       from google.colab import userdata
       secret_value = userdata.get(name)
       if secret_value:
           value = secret_value
   except Exception:
       pass
   return value or ""
print("Python:", sys.version)
assert sys.version_info >= (3, 10), "QwenPaw needs Python 3.10+."
pip_spec = os.environ.get("QWENPAW_PIP_SPEC", "qwenpaw")
run([sys.executable, "-m", "pip", "install", "-q", "-U", "pip", "setuptools", "wheel"], check=False)
run([sys.executable, "-m", "pip", "install", "-q", "-U", pip_spec, "requests"], check=True)
try:
   import requests
except Exception:
   run([sys.executable, "-m", "pip", "install", "-q", "-U", "requests"], check=True)
   import requests

if not (WORKING_DIR / "config.json").exists():
   run(qwenpaw_cmd("init", "--defaults"), check=False)
else:
   print("QwenPaw working directory already initialized:", WORKING_DIR)
provider_candidates = [
   {
       "env": "OPENAI_API_KEY",
       "provider_id": "openai",
       "name": "OpenAI",
       "base_url": "https://api.openai.com/v1",
       "model": os.environ.get("QWENPAW_MODEL", "gpt-4o-mini"),
       "chat_model": "OpenAIChatModel",
       "prefix": "sk-",
   },
   {
       "env": "OPENROUTER_API_KEY",
       "provider_id": "openrouter",
       "name": "OpenRouter",
       "base_url": "https://openrouter.ai/api/v1",
       "model": os.environ.get("QWENPAW_MODEL", "openai/gpt-4o-mini"),
       "chat_model": "OpenAIChatModel",
       "prefix": "sk-or-",
   },
   {
       "env": "DASHSCOPE_API_KEY",
       "provider_id": "dashscope",
       "name": "DashScope",
       "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
       "model": os.environ.get("QWENPAW_MODEL", "qwen-plus"),
       "chat_model": "OpenAIChatModel",
       "prefix": "sk-",
   },
   {
       "env": "DEEPSEEK_API_KEY",
       "provider_id": "deepseek",
       "name": "DeepSeek",
       "base_url": "https://api.deepseek.com",
       "model": os.environ.get("QWENPAW_MODEL", "deepseek-chat"),
       "chat_model": "OpenAIChatModel",
       "prefix": "sk-",
   },
   {
       "env": "GEMINI_API_KEY",
       "provider_id": "gemini",
       "name": "Google Gemini",
       "base_url": "https://generativelanguage.googleapis.com",
       "model": os.environ.get("QWENPAW_MODEL", "gemini-2.5-flash"),
       "chat_model": "GeminiChatModel",
       "prefix": "",
   },
   {
       "env": "GOOGLE_API_KEY",
       "provider_id": "gemini",
       "name": "Google Gemini",
       "base_url": "https://generativelanguage.googleapis.com",
       "model": os.environ.get("QWENPAW_MODEL", "gemini-2.5-flash"),
       "chat_model": "GeminiChatModel",
       "prefix": "",
   },
]
selected = None
for candidate in provider_candidates:
   api_key = colab_secret_or_env(candidate["env"])
   if api_key:
       selected = {**candidate, "api_key": api_key}
       break
def read_json(path, default):
   try:
       if path.exists():
           return json.loads(path.read_text(encoding="utf-8"))
   except Exception:
       pass
   return default
def write_json(path, data):
   path.parent.mkdir(parents=True, exist_ok=True)
   path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
config_path = WORKING_DIR / "config.json"
config = read_json(config_path, {})
config.setdefault("agents", {})
config["agents"].setdefault("active_agent", "default")
config["agents"].setdefault("agent_order", ["default"])
config["agents"].setdefault("profiles", {})
config["agents"]["profiles"].setdefault("default", {})
config["agents"]["profiles"]["default"].update(
   {
       "id": "default",
       "name": "Colab Research Assistant",
       "description": "A QwenPaw agent configured for Google Colab tutorials, local files, custom skills, and API testing.",
       "workspace_dir": str(WORKSPACE_DIR),
       "enabled": True,
   }
)
config["last_api"] = {"host": "127.0.0.1", "port": PORT}
config["show_tool_details"] = True
config["user_timezone"] = "Asia/Kolkata"
write_json(config_path, config)

agent_dir = WORKING_DIR / "agents" / "default"
agent_dir.mkdir(parents=True, exist_ok=True)
agent_path = agent_dir / "agent.json"
agent = read_json(agent_path, {})
agent.update(
   {
       "id": "default",
       "name": "Colab Research Assistant",
       "description": "Advanced QwenPaw tutorial agent for Colab: file-aware, skill-aware, API-testable, and guarded.",
       "language": "en",
       "workspace_dir": str(WORKSPACE_DIR),
       "enabled": True,
       "channels": {
           "console": {
               "enabled": True
           }
       },
       "running": {
           "max_iters": 30,
           "llm_retry_enabled": True,
           "stream_output": True
       },
       "security": {
           "tool_guard": True,
           "file_guard": True,
           "skill_scanner": True,
           "skill_scan_mode": "warn"
       },
       "tool_filter": {
           "enabled": False,
           "allow": [],
           "deny": []
       },
       "memory": {
           "enabled": True
       }
   }
)
if selected:
   provider_dir = SECRET_DIR / "providers" / "builtin"
   provider_dir.mkdir(parents=True, exist_ok=True)
   provider_payload = {
       "id": selected["provider_id"],
       "name": selected["name"],
       "base_url": selected["base_url"],
       "api_key": selected["api_key"],
       "chat_model": selected["chat_model"],
       "models": [],
       "extra_models": [
           {
               "id": selected["model"],
               "name": selected["model"],
               "supports_image": None,
               "supports_video": None,
               "supports_multimodal": None,
               "is_free": False,
               "max_tokens": int(os.environ.get("QWENPAW_MAX_TOKENS", "2048")),
               "max_input_length": int(os.environ.get("QWENPAW_MAX_INPUT_LENGTH", "131072")),
               "generate_kwargs": {
                   "temperature": float(os.environ.get("QWENPAW_TEMPERATURE", "0.2")),
                   "max_tokens": int(os.environ.get("QWENPAW_MAX_TOKENS", "2048")),
               },
           }
       ],
       "api_key_prefix": selected["prefix"],
       "is_local": False,
       "freeze_url": True,
       "require_api_key": True,
       "is_custom": False,
       "support_model_discovery": False,
       "support_connection_check": False,
       "generate_kwargs": {
           "temperature": float(os.environ.get("QWENPAW_TEMPERATURE", "0.2")),
           "max_tokens": int(os.environ.get("QWENPAW_MAX_TOKENS", "2048")),
       },
       "custom_headers": {},
       "auth_mode": "api_key",
       "meta": {},
   }
   write_json(provider_dir / f"{selected['provider_id']}.json", provider_payload)
   write_json(
       SECRET_DIR / "providers" / "active_model.json",
       {"provider_id": selected["provider_id"], "model": selected["model"]},
   )
   agent["active_model"] = {"provider_id": selected["provider_id"], "model": selected["model"]}
   print(f"Configured model provider: {selected['name']} / {selected['model']}")
else:
   print(
       "No model key found. The web app will still launch, but chat requires a configured model.\n"
       "Add one Colab secret or environment variable such as OPENAI_API_KEY, OPENROUTER_API_KEY, "
       "DASHSCOPE_API_KEY, DEEPSEEK_API_KEY, GEMINI_API_KEY, or GOOGLE_API_KEY, then rerun."
   )
write_json(agent_path, agent)

skill_dir = WORKSPACE_DIR / "skills" / "research_brief"
skill_dir.mkdir(parents=True, exist_ok=True)
(skill_dir / "SKILL.md").write_text(
   """---
name: research_brief
description: Create rigorous research briefs from user questions, local notes, uploaded files, and available tools.
---
# Research Brief Skill
Use this skill when the user asks for research, product analysis, market mapping, technical due diligence, paper analysis, repo analysis, or a decision memo.
## Procedure
1. Restate the user's objective in one sentence.
2. Identify the most important entities, assumptions, and constraints.
3. Search available local workspace files first.
4. Use tools only when they are relevant and allowed.
5. Separate verified facts from inference.
6. Produce a compact brief with:
  - answer
  - evidence
  - risks or caveats
  - recommended next step
## Output Style
Prefer clear sections, short paragraphs, and explicit uncertainty.
Do not invent citations, file contents, commands, or results.
""",
   encoding="utf-8",
)
demo_dir = WORKSPACE_DIR / "demo_knowledge"
demo_dir.mkdir(parents=True, exist_ok=True)
(demo_dir / "qwenpaw_colab_notes.md").write_text(
   f"""# QwenPaw Colab Demo Notes
Created: {datetime.now().isoformat(timespec="seconds")}
This workspace is prepared by a Google Colab tutorial.
The tutorial demonstrates:
- QwenPaw installation and initialization
- provider auto-configuration from Colab secrets or environment variables
- authenticated Console launch
- custom workspace skill creation
- local workspace knowledge files
- streaming REST API calls
- optional public tunnel exposure
Recommended first prompt in the Console:
"Read my workspace notes and explain what this QwenPaw Colab setup can do. Then use the research_brief skill style to propose three advanced experiments."
""",
   encoding="utf-8",
)
(WORKSPACE_DIR / "README_COLAB_TUTORIAL.md").write_text(
   """# QwenPaw Advanced Colab Workspace
This workspace is intentionally small but structured like a real assistant workspace.
Suggested experiments:
1. Ask QwenPaw to inspect the demo_knowledge folder.
2. Ask it to use the research_brief skill style.
3. Use the REST API client in this notebook for automated tests.
4. Add more SKILL.md folders under workspace/skills.
5. Add more notes, CSVs, markdown files, or task briefs under workspace folders.
""",
   encoding="utf-8",
)
print("\nWorkspace prepared:")
print("Working dir:", WORKING_DIR)
print("Secret dir :", SECRET_DIR)
print("Workspace  :", WORKSPACE_DIR)
print("Skill file :", skill_dir / "SKILL.md")
run(qwenpaw_cmd("daemon", "version"), check=False)
run(qwenpaw_cmd("models", "list"), check=False)
run(qwenpaw_cmd("skills", "list", "--agent-id", "default"), check=False)

stop_previous_app()
APP_LOG.parent.mkdir(parents=True, exist_ok=True)
log_fh = APP_LOG.open("w", encoding="utf-8")
app_proc = subprocess.Popen(
   qwenpaw_cmd("app", "--host", "0.0.0.0", "--port", str(PORT), "--log-level", os.environ["QWENPAW_LOG_LEVEL"]),
   stdout=log_fh,
   stderr=subprocess.STDOUT,
   env=os.environ.copy(),
)
PID_FILE.write_text(str(app_proc.pid), encoding="utf-8")
if not wait_for_port(PORT, seconds=120):
   print("\nQwenPaw did not open the port. Last log lines:")
   try:
       print(APP_LOG.read_text(encoding="utf-8")[-6000:])
   except Exception as e:
       print("Could not read log:", e)
   raise RuntimeError("QwenPaw app failed to start.")
print(f"\nQwenPaw app is running on http://127.0.0.1:{PORT}")
print("Username:", os.environ["QWENPAW_AUTH_USERNAME"])
print("Password:", os.environ["QWENPAW_AUTH_PASSWORD"])
print("App log:", APP_LOG)
try:
   from google.colab import output
   proxy_url = output.eval_js(f"google.colab.kernel.proxyPort({PORT})")
   print("\nColab proxied Console URL:")
   print(proxy_url)
   try:
       output.serve_kernel_port_as_window(PORT)
   except Exception:
       pass
except Exception as e:
   print("\nNot running inside Google Colab proxy environment:", e)
def start_cloudflared_tunnel(port):
   system_bin = pathlib.Path("/usr/local/bin/cloudflared")
   local_bin = ROOT / "cloudflared"
   cloudflared = system_bin if system_bin.exists() else local_bin
   if not cloudflared.exists():
       url = "https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64"
       target = str(system_bin)
       rc, _ = run(f"wget -q {shlex.quote(url)} -O {shlex.quote(target)} && chmod +x {shlex.quote(target)}", check=False)
       if rc != 0 or not system_bin.exists():
           target = str(local_bin)
           rc, _ = run(f"wget -q {shlex.quote(url)} -O {shlex.quote(target)} && chmod +x {shlex.quote(target)}", check=False)
       cloudflared = pathlib.Path(target)
   if not cloudflared.exists():
       print("cloudflared tunnel unavailable. Use the Colab proxy URL above.")
       return None, None
   tunnel_log = LOG_DIR / "cloudflared.log"
   fh = tunnel_log.open("w", encoding="utf-8")
   proc = subprocess.Popen(
       [str(cloudflared), "tunnel", "--url", f"http://127.0.0.1:{port}", "--no-autoupdate"],
       stdout=fh,
       stderr=subprocess.STDOUT,
       text=True,
   )
   public_url = None
   start = time.time()
   while time.time() - start < 45:
       time.sleep(1)
       try:
           text = tunnel_log.read_text(encoding="utf-8", errors="ignore")
       except Exception:
           text = ""
       for token in text.replace("|", " ").split():
           if token.startswith("https://") and "trycloudflare.com" in token:
               public_url = token.strip()
               break
       if public_url:
           break
   if public_url:
       print("\nTemporary public tunnel URL:")
       print(public_url)
       print("Use the same username/password printed above.")
   else:
       print("\nCloudflare tunnel started but no URL was detected yet.")
       print("Tunnel log:", tunnel_log)
   return proc, public_url
ENABLE_CLOUDFLARE_TUNNEL = os.environ.get("ENABLE_QWENPAW_TUNNEL", "1") == "1"
cloudflared_proc, public_url = (None, None)
if ENABLE_CLOUDFLARE_TUNNEL:
   cloudflared_proc, public_url = start_cloudflared_tunnel(PORT)

def qwenpaw_chat(message, session_id=None, user_id="colab-user", agent_id="default", timeout=180):
   session_id = session_id or f"colab-{uuid.uuid4().hex[:10]}"
   url = f"http://127.0.0.1:{PORT}/api/console/chat"
   headers = {
       "Content-Type": "application/json",
       "X-Agent-Id": agent_id,
   }
   payload = {
       "message": message,
       "session_id": session_id,
       "user_id": user_id,
   }
   print("\nAPI request:")
   print(json.dumps({**payload, "message": message[:300]}, indent=2))
   with requests.post(url, headers=headers, json=payload, stream=True, timeout=timeout) as response:
       print("HTTP status:", response.status_code)
       if response.status_code >= 400:
           print(response.text[:4000])
           response.raise_for_status()
       last_text = ""
       final_text = ""
       raw_events_seen = 0
       for raw in response.iter_lines(decode_unicode=True):
           if not raw:
               continue
           if raw.startswith("data:"):
               raw_events_seen += 1
               data = raw[len("data:"):].strip()
               if data == "[DONE]":
                   break
               try:
                   event = json.loads(data)
               except Exception:
                   continue
               candidate_texts = []
               def walk(x):
                   if isinstance(x, dict):
                       for key, value in x.items():
                           if key in {"text", "content", "message", "delta"} and isinstance(value, str):
                               candidate_texts.append(value)
                           else:
                               walk(value)
                   elif isinstance(x, list):
                       for item in x:
                           walk(item)
               walk(event)
               if candidate_texts:
                   text = candidate_texts[-1]
                   if text and len(text) >= len(final_text):
                       final_text = text
                       if text.startswith(last_text):
                           print(text[len(last_text):], end="", flush=True)
                       else:
                           print("\n" + text, end="", flush=True)
                       last_text = text
       print("\n\nStreaming events seen:", raw_events_seen)
       return {"session_id": session_id, "text": final_text}
if selected:
   try:
       result = qwenpaw_chat(
           "Read the local workspace notes if available. Then explain this Colab QwenPaw setup in five bullets and suggest two advanced experiments.",
           session_id="qwenpaw-colab-demo",
       )
       print("\nFinal session_id:", result["session_id"])
   except Exception as e:
       print("\nAPI demo failed. The Console may still work; inspect the app log below.")
       print("Error:", repr(e))
       try:
           print(APP_LOG.read_text(encoding="utf-8")[-8000:])
       except Exception:
           pass
else:
   print(
       "\nSkipping API chat demo because no model provider key was found.\n"
       "Open the Console URL above, or add a Colab secret such as OPENAI_API_KEY / OPENROUTER_API_KEY / DASHSCOPE_API_KEY / DEEPSEEK_API_KEY / GEMINI_API_KEY and rerun."
   )
print("\nSummary")
print("Console username:", os.environ["QWENPAW_AUTH_USERNAME"])
print("Console password:", os.environ["QWENPAW_AUTH_PASSWORD"])
print("Local URL:", f"http://127.0.0.1:{PORT}")
if public_url:
   print("Public tunnel:", public_url)
print("Workspace:", WORKSPACE_DIR)
print("Logs:", APP_LOG)
print("\nTo stop the server manually, run:")
print(f"import os, signal; os.kill({app_proc.pid}, signal.SIGTERM)")