ws-api

#!/usr/bin/env python3
"""
ws-api - Remote workspace feels like local. Zero shell escaping issues.
All data: json.dumps() -> HTTPS/WSS -> json.loads(), never through shell parsing.

Transport: WSS (default, persistent connection) or HTTPS (fallback).
  Set WS_TRANSPORT=http to force HTTPS mode.

== Transparent exec (just prefix ws-api, like running locally) ==
  ws-api ls -la /root/workspace
  ws-api git status
  ws-api go build ./...
  ws-api cat /etc/hostname
  ws-api docker ps

== Background exec ==
  ws-api go build ./... --bg                    # returns job ID immediately
  ws-api exec "make build" --bg                 # explicit exec, background
  ws-api jobs                                   # list all jobs
  ws-api job <id>                               # get job output (real-time if running)
  ws-api job <id> --clear                       # get output & remove job
  ws-api kill <id>                              # kill a running bg job

== File operations ==
  ws-api read <path> [-n] [--offset N] [--limit N]
  ws-api write <path> < file                         # stdin
  ws-api write <path> --content "text"               # inline
  ws-api write <path> --file local.txt               # from local file

== File transfer (binary-safe, base64) ==
  ws-api upload <local-path> <remote-path> [--mode 0755]
  ws-api download <remote-path> <local-path>

== Text editing (find & replace) ==
  ws-api edit <path> "old text" "new text" [--all]    # simple text (no $`! chars)
  ws-api edit <path> --old "old" --new "new" [--all]  # named args
  ws-api edit <path> --old-file o.txt --new-file n.txt # shell-safe for regex/$`!
  ws-api edit <path> --dry-run ...                    # preview without writing
  echo '{"old":"...","new":"..."}' | ws-api edit <path> # pipe JSON (also shell-safe)

== Line-based editing (by line numbers) ==
  ws-api edit-lines <path> --delete 10-20
  ws-api edit-lines <path> --insert 15 < content
  ws-api edit-lines <path> --insert 15 --content "new line"
  ws-api edit-lines <path> --replace 10-20 < content

== Batch editing (multiple replacements in one call) ==
  echo '[{"old":"x","new":"y"},{"old":"a","new":"b"}]' | ws-api patch <path>

== Search ==
  ws-api glob "*.go" [--path /root/workspace]
  ws-api grep "pattern" [--path /root/workspace] [--glob "*.go"] [--context 2]

== Explicit exec (with options) ==
  ws-api exec "command" [--dir /path] [--timeout 60] [--bg]

== Other ==
  ws-api ping
  ws-api raw <endpoint> < json_body
  ws-api help
"""
import base64
import json
import os
import shlex
import sys
import urllib.request
import urllib.error
import uuid

API_BASE = "https://hcwkapi.bygsoga.cc"
WS_URL = "wss://hcwkapi.bygsoga.cc/ws"
AUTH_TOKEN = "capy-workspace-7f3a9b2e"

# Transport: "wss" (default) or "http"
TRANSPORT = os.environ.get("WS_TRANSPORT", "wss").lower()

# Singleton WebSocket connection for WSS mode
_ws_conn = None

def _get_ws():
    """Get or create a persistent WebSocket connection."""
    global _ws_conn
    if _ws_conn is not None:
        return _ws_conn
    try:
        import websocket
        _ws_conn = websocket.create_connection(
            f"{WS_URL}?token={AUTH_TOKEN}",
            timeout=300,
        )
        return _ws_conn
    except Exception:
        return None

def _ws_call(action, data):
    """Send a request over WebSocket and return the response."""
    ws = _get_ws()
    if ws is None:
        return None
    msg_id = uuid.uuid4().hex[:12]
    msg = {"id": msg_id, "action": action, "data": data}
    try:
        ws.send(json.dumps(msg))
        resp = ws.recv()
        return json.loads(resp)
    except Exception:
        global _ws_conn
        _ws_conn = None
        return None

def _http_call(endpoint, data):
    """Send a request over HTTPS and return the response."""
    url = f"{API_BASE}/api/{endpoint}"
    payload = json.dumps(data).encode("utf-8")
    req = urllib.request.Request(url, data=payload, method="POST")
    req.add_header("Content-Type", "application/json")
    req.add_header("Authorization", f"Bearer {AUTH_TOKEN}")
    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8")
        try:
            return json.loads(body)
        except:
            return {"ok": False, "error": f"HTTP {e.code}: {body}"}
    except Exception as e:
        return {"ok": False, "error": str(e)}

def api_call(endpoint, data):
    """Unified API call: try WSS first (if enabled), fall back to HTTPS."""
    if TRANSPORT == "wss":
        result = _ws_call(endpoint, data)
        if result is not None:
            return result
        # WSS failed, fall back to HTTPS silently
    return _http_call(endpoint, data)

def cmd_ping():
    r = api_call("ping", {})
    r["transport"] = "wss" if (TRANSPORT == "wss" and _ws_conn is not None) else "https"
    print(json.dumps(r, indent=2))

def cmd_read(args):
    if not args:
        print("ERROR: path is required", file=sys.stderr); sys.exit(1)
    data = {"path": args[0]}
    numbered = False
    i = 1
    while i < len(args):
        if args[i] in ("-n", "--numbered"):
            numbered = True; data["numbered"] = True; i += 1
        elif args[i] == "--offset" and i+1 < len(args):
            data["offset"] = int(args[i+1]); i += 2
        elif args[i] == "--limit" and i+1 < len(args):
            data["limit"] = int(args[i+1]); i += 2
        else:
            i += 1
    r = api_call("read", data)
    if r.get("ok"):
        print(r["content"])
        total = r.get("total_lines", "?")
        size = r.get("size", "?")
        if numbered or data.get("offset") or data.get("limit"):
            start = r.get("start_line", 1)
            shown = r["content"].count("\n") + 1
            print(f"\n--- {total} lines total, {size} bytes, showing from line {start} ({shown} lines) ---", file=sys.stderr)
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_write(args):
    if not args:
        print("ERROR: path is required", file=sys.stderr); sys.exit(1)
    path = args[0]
    content = None
    i = 1
    while i < len(args):
        if args[i] == "--content" and i+1 < len(args):
            content = args[i+1]; i += 2
        elif args[i] == "--file" and i+1 < len(args):
            with open(args[i+1], "r") as f: content = f.read()
            i += 2
        else:
            i += 1
    if content is None:
        if not sys.stdin.isatty():
            content = sys.stdin.read()
        else:
            print("ERROR: provide content via stdin, --content, or --file", file=sys.stderr)
            sys.exit(1)
    r = api_call("write", {"path": path, "content": content})
    if r.get("ok"):
        print(f"Written {r['written']} bytes to {path}")
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_edit(args):
    if not args:
        print("ERROR: path is required", file=sys.stderr); sys.exit(1)
    path = args[0]
    old = new = None
    replace_all = False
    dry_run = False
    use_stdin = False
    positional = []
    i = 1
    while i < len(args):
        if args[i] == "--old" and i+1 < len(args):
            old = args[i+1]; i += 2
        elif args[i] == "--new" and i+1 < len(args):
            new = args[i+1]; i += 2
        elif args[i] == "--old-file" and i+1 < len(args):
            with open(args[i+1], "r") as f: old = f.read()
            i += 2
        elif args[i] == "--new-file" and i+1 < len(args):
            with open(args[i+1], "r") as f: new = f.read()
            i += 2
        elif args[i] in ("--all", "--replace-all"):
            replace_all = True; i += 1
        elif args[i] == "--dry-run":
            dry_run = True; i += 1
        elif args[i] == "--stdin":
            use_stdin = True; i += 1
        else:
            positional.append(args[i]); i += 1
    # Support positional: ws-api edit <path> "old" "new"
    if old is None and len(positional) >= 1:
        old = positional[0]
    if new is None and len(positional) >= 2:
        new = positional[1]
    # Try stdin only if explicitly requested or if we still need old/new
    if old is None and new is None:
        if use_stdin or not sys.stdin.isatty():
            stdin_text = sys.stdin.read()
            if stdin_text.strip():
                try:
                    stdin_data = json.loads(stdin_text)
                    old = stdin_data.get("old", old)
                    new = stdin_data.get("new", new)
                    if "replace_all" in stdin_data:
                        replace_all = stdin_data["replace_all"]
                except json.JSONDecodeError as e:
                    print(f"ERROR: Invalid JSON on stdin: {e}", file=sys.stderr)
                    sys.exit(1)
    if old is None or new is None:
        print("ERROR: provide old/new via positional args, --old/--new flags, or pipe JSON via stdin", file=sys.stderr)
        sys.exit(1)
    # Legacy compat: unescape \! from older shell workarounds
    old = old.replace("\\!", "!")
    new = new.replace("\\!", "!")
    data = {"path": path, "old_string": old, "new_string": new, "replace_all": replace_all}
    if dry_run:
        data["dry_run"] = True
    r = api_call("edit", data)
    if r.get("ok"):
        if r.get("dry_run"):
            count = r.get("replacements", 0)
            locs = r.get("locations", [])
            delta = r.get("size_delta", 0)
            print(f"[dry-run] Would replace {count} occurrence(s)")
            if locs:
                print(f"  Lines: {', '.join(str(l) for l in locs)}")
            print(f"  Old length: {r.get('old_length', '?')} chars")
            print(f"  New length: {r.get('new_length', '?')} chars")
            print(f"  Size delta: {'+' if delta >= 0 else ''}{delta} bytes")
        else:
            print(f"Replaced {r['replacements']} occurrence(s)")
            locs = r.get("locations")
            if locs:
                print(f"  Lines: {', '.join(str(l) for l in locs)}", file=sys.stderr)
    else:
        err = r.get('error', 'unknown')
        print(f"ERROR: {err}", file=sys.stderr)
        if "not found" in err.lower() or "not unique" in err.lower():
            has_shell_chars = any(c in (old + new) for c in '$`!')
            from_positional = len(positional) >= 1
            if from_positional and not has_shell_chars:
                print("", file=sys.stderr)
                print("HINT: If your text contained $, `, or ! characters, they may have been", file=sys.stderr)
                print("      eaten by the shell. Use one of these shell-safe methods instead:", file=sys.stderr)
                print("", file=sys.stderr)
                print("  # Method 1: --old-file / --new-file (most reliable)", file=sys.stderr)
                print("  ws-api edit <path> --old-file /tmp/old.txt --new-file /tmp/new.txt", file=sys.stderr)
                print("", file=sys.stderr)
                print("  # Method 2: pipe JSON via stdin", file=sys.stderr)
                print('  echo \'{"old":"...","new":"..."}\' | ws-api edit <path>', file=sys.stderr)
        sys.exit(1)

def cmd_edit_lines(args):
    """Line-based editing: delete, insert, replace by line numbers."""
    if not args:
        print("ERROR: path is required", file=sys.stderr); sys.exit(1)
    path = args[0]
    action = None
    start = end = 0
    content = None
    i = 1
    while i < len(args):
        if args[i] == "--delete" and i+1 < len(args):
            action = "delete"
            start, end = _parse_range(args[i+1]); i += 2
        elif args[i] == "--insert" and i+1 < len(args):
            action = "insert"
            start = int(args[i+1]); end = start; i += 2
        elif args[i] == "--replace" and i+1 < len(args):
            action = "replace"
            start, end = _parse_range(args[i+1]); i += 2
        elif args[i] == "--content" and i+1 < len(args):
            content = args[i+1]; i += 2
        else:
            i += 1
    if action is None:
        print("ERROR: --delete, --insert, or --replace is required", file=sys.stderr)
        sys.exit(1)
    if action in ("insert", "replace") and content is None:
        if not sys.stdin.isatty():
            content = sys.stdin.read()
        else:
            print("ERROR: content required via stdin or --content", file=sys.stderr)
            sys.exit(1)
    data = {"path": path, "action": action, "start": start, "end": end}
    if content is not None:
        data["content"] = content
    r = api_call("edit-lines", data)
    if r.get("ok"):
        print(r.get("info", "done"))
        print(f"  {r.get('old_lines','?')} -> {r.get('new_lines','?')} lines", file=sys.stderr)
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def _parse_range(s):
    """Parse '10-20' or '10' into (start, end)."""
    if "-" in s:
        parts = s.split("-", 1)
        return int(parts[0]), int(parts[1])
    n = int(s)
    return n, n

def cmd_patch(args):
    """Apply multiple edits atomically: ws-api patch <path> < edits.json"""
    if not args:
        print("ERROR: path is required", file=sys.stderr); sys.exit(1)
    path = args[0]
    if not sys.stdin.isatty():
        try:
            edits = json.loads(sys.stdin.read())
        except json.JSONDecodeError as e:
            print(f"ERROR: Invalid JSON: {e}", file=sys.stderr)
            sys.exit(1)
    else:
        print("ERROR: pipe edits JSON via stdin", file=sys.stderr)
        print('  Example: echo \'[{"old":"x","new":"y"},{"old":"a","new":"b"}]\' | ws-api patch file', file=sys.stderr)
        sys.exit(1)
    if isinstance(edits, list):
        edits_list = edits
    elif isinstance(edits, dict) and "edits" in edits:
        edits_list = edits["edits"]
    else:
        print("ERROR: expected JSON array or {\"edits\":[...]}", file=sys.stderr)
        sys.exit(1)
    normalized = []
    for e in edits_list:
        normalized.append({
            "old_string": e.get("old_string", e.get("old", "")),
            "new_string": e.get("new_string", e.get("new", "")),
        })
    r = api_call("patch", {"path": path, "edits": normalized})
    if r.get("ok"):
        print(f"Applied {r.get('edits_applied', '?')} edit(s)")
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def _handle_exec_result(r):
    """Process exec result: print stdout to stdout, stderr to stderr, exit with code."""
    if r.get("ok"):
        if r.get("stdout") is not None:
            if r["stdout"]:
                print(r["stdout"], end="", file=sys.stdout)
            if r.get("stderr"):
                print(r["stderr"], end="", file=sys.stderr)
        elif r.get("output"):
            print(r["output"], end="")
        if r.get("timed_out"):
            print("\n[ws-api] Command timed out (partial output above)", file=sys.stderr)
        sys.exit(r.get("exit_code", 0))
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_exec(args):
    if not args:
        print("ERROR: command is required", file=sys.stderr); sys.exit(1)
    command = args[0]
    data = {"command": command}
    bg = False
    i = 1
    while i < len(args):
        if args[i] in ("--dir", "-C") and i+1 < len(args):
            data["dir"] = args[i+1]; i += 2
        elif args[i] == "--timeout" and i+1 < len(args):
            data["timeout"] = int(args[i+1]); i += 2
        elif args[i] == "--bg":
            bg = True; i += 1
        else:
            i += 1
    if bg:
        _handle_exec_bg(data)
    else:
        _handle_exec_result(api_call("exec", data))

def _handle_exec_bg(data):
    """Submit background job and print job ID."""
    r = api_call("exec-bg", data)
    if r.get("ok"):
        job_id = r.get("job_id", "?")
        print(f"Background job started: {job_id}")
        print(f"  Command: {r.get('command', '?')}")
        print(f"  Check status: ws-api job {job_id}")
        print(f"  List all:     ws-api jobs")
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_exec_transparent(all_args):
    """Transparent exec: join all args into a command, preserving quoting."""
    bg = False
    filtered = []
    timeout = None
    i = 0
    while i < len(all_args):
        if all_args[i] == "--bg":
            bg = True; i += 1
        elif all_args[i] == "--timeout" and i+1 < len(all_args):
            timeout = int(all_args[i+1]); i += 2
        else:
            filtered.append(all_args[i]); i += 1
    command = " ".join(shlex.quote(a) for a in filtered)
    data = {"command": command}
    if timeout is not None:
        data["timeout"] = timeout
    if bg:
        _handle_exec_bg(data)
    else:
        _handle_exec_result(api_call("exec", data))

def cmd_jobs(args):
    """List background jobs."""
    r = api_call("jobs", {})
    if r.get("ok"):
        jobs = r.get("jobs", [])
        if not jobs:
            print("No background jobs")
            return
        print(f"{'ID':<16} {'STATUS':<10} {'EXIT':<6} {'OUTPUT':<10} {'COMMAND'}")
        print("-" * 80)
        for j in jobs:
            status = j.get("status", "?")
            if j.get("timed_out"):
                status = "timeout"
            eid = str(j.get("exit_code", "")) if status != "running" else ""
            if status == "running":
                out_bytes = j.get("stdout_bytes", 0) + j.get("stderr_bytes", 0)
                out_info = f"{out_bytes}B" if out_bytes > 0 else "-"
            else:
                out_info = ""
            print(f"{j['id']:<16} {status:<10} {eid:<6} {out_info:<10} {j.get('command', '?')}")
        print(f"\n{len(jobs)} job(s)", file=sys.stderr)
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_job(args):
    """Get job detail/output."""
    if not args:
        print("ERROR: job id is required", file=sys.stderr); sys.exit(1)
    job_id = args[0]
    clear = "--clear" in args
    r = api_call("job", {"id": job_id, "clear": clear})
    if r.get("ok"):
        status = r.get("status", "?")
        status_extra = ""
        if r.get("timed_out"):
            status_extra = " (timed out)"
        elif status == "killed":
            status_extra = " (manually killed)"
        print(f"Job:     {r.get('id')}")
        print(f"Status:  {status}{status_extra}")
        print(f"Command: {r.get('command')}")
        if r.get("dir"):
            print(f"Dir:     {r['dir']}")
        if status != "running":
            print(f"Exit:    {r.get('exit_code', '?')}")
        stdout = r.get("stdout", "")
        stderr = r.get("stderr", "")
        if stdout:
            label = "stdout (live)" if status == "running" else "stdout"
            print(f"\n--- {label} ---")
            print(stdout, end="")
        if stderr:
            label = "stderr (live)" if status == "running" else "stderr"
            print(f"\n--- {label} ---", file=sys.stderr)
            print(stderr, end="", file=sys.stderr)
        if not stdout and not stderr:
            if status == "running":
                print("\n(running, no output yet)")
            else:
                print("\n(no output)")
        if clear:
            print(f"\n[job cleared]", file=sys.stderr)
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_kill(args):
    """Kill a running background job."""
    if not args:
        print("ERROR: job id is required", file=sys.stderr); sys.exit(1)
    job_id = args[0]
    r = api_call("job-kill", {"id": job_id})
    if r.get("ok"):
        print(f"Job {job_id} killed")
        print(f"  Use 'ws-api job {job_id}' to see final output")
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_upload(args):
    """Upload local file to remote: ws-api upload <local> <remote> [--mode 0755]"""
    if len(args) < 2:
        print("ERROR: usage: ws-api upload <local-path> <remote-path> [--mode 0755]", file=sys.stderr)
        sys.exit(1)
    local_path = args[0]
    remote_path = args[1]
    mode = 0o644
    i = 2
    while i < len(args):
        if args[i] == "--mode" and i+1 < len(args):
            mode = int(args[i+1], 8); i += 2
        else:
            i += 1
    if not os.path.isfile(local_path):
        print(f"ERROR: local file not found: {local_path}", file=sys.stderr)
        sys.exit(1)
    with open(local_path, "rb") as f:
        data = f.read()
    encoded = base64.b64encode(data).decode("ascii")
    local_size = len(data)
    print(f"Uploading {local_path} ({local_size} bytes) -> {remote_path}")
    r = api_call("upload", {"path": remote_path, "content": encoded, "mode": mode})
    if r.get("ok"):
        print(f"Uploaded {r.get('written', '?')} bytes to {remote_path}")
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_download(args):
    """Download remote file to local: ws-api download <remote> <local>"""
    if len(args) < 2:
        print("ERROR: usage: ws-api download <remote-path> <local-path>", file=sys.stderr)
        sys.exit(1)
    remote_path = args[0]
    local_path = args[1]
    print(f"Downloading {remote_path} -> {local_path}")
    r = api_call("download", {"path": remote_path})
    if r.get("ok"):
        data = base64.b64decode(r["content"])
        local_dir = os.path.dirname(local_path)
        if local_dir:
            os.makedirs(local_dir, exist_ok=True)
        with open(local_path, "wb") as f:
            f.write(data)
        print(f"Downloaded {len(data)} bytes to {local_path}")
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_glob(args):
    if not args:
        print("ERROR: pattern is required", file=sys.stderr); sys.exit(1)
    data = {"pattern": args[0]}
    i = 1
    while i < len(args):
        if args[i] == "--path" and i+1 < len(args):
            data["path"] = args[i+1]; i += 2
        else:
            i += 1
    r = api_call("glob", data)
    if r.get("ok"):
        for f in r.get("files", []):
            print(f)
        print(f"\n{r['count']} file(s) found", file=sys.stderr)
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_grep(args):
    if not args:
        print("ERROR: pattern is required", file=sys.stderr); sys.exit(1)
    data = {"pattern": args[0]}
    i = 1
    while i < len(args):
        if args[i] == "--path" and i+1 < len(args):
            data["path"] = args[i+1]; i += 2
        elif args[i] == "--glob" and i+1 < len(args):
            data["glob"] = args[i+1]; i += 2
        elif args[i] == "--context" and i+1 < len(args):
            data["context"] = int(args[i+1]); i += 2
        else:
            i += 1
    r = api_call("grep", data)
    if r.get("ok"):
        if r.get("output"):
            print(r["output"], end="")
    else:
        print(f"ERROR: {r.get('error', 'unknown')}", file=sys.stderr)
        sys.exit(1)

def cmd_raw(args):
    """Raw JSON mode: ws-api raw <endpoint> < json_body"""
    if not args:
        print("ERROR: endpoint is required", file=sys.stderr); sys.exit(1)
    endpoint = args[0]
    data = json.loads(sys.stdin.read())
    r = api_call(endpoint, data)
    print(json.dumps(r, indent=2))

COMMANDS = {
    "ping":       lambda a: cmd_ping(),
    "read":       lambda a: cmd_read(a),
    "write":      lambda a: cmd_write(a),
    "edit":       lambda a: cmd_edit(a),
    "edit-lines": lambda a: cmd_edit_lines(a),
    "patch":      lambda a: cmd_patch(a),
    "exec":       lambda a: cmd_exec(a),
    "glob":       lambda a: cmd_glob(a),
    "grep":       lambda a: cmd_grep(a),
    "raw":        lambda a: cmd_raw(a),
    "jobs":       lambda a: cmd_jobs(a),
    "job":        lambda a: cmd_job(a),
    "kill":       lambda a: cmd_kill(a),
    "upload":     lambda a: cmd_upload(a),
    "download":   lambda a: cmd_download(a),
    "help":       lambda a: print(__doc__),
}

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    cmd = sys.argv[1]
    args = sys.argv[2:]

    if cmd in COMMANDS:
        COMMANDS[cmd](args)
    elif cmd.startswith("-"):
        print(__doc__)
        sys.exit(1)
    else:
        # Transparent exec: ws-api ls -la /root -> exec "ls -la /root"
        cmd_exec_transparent(sys.argv[1:])