gravity_control/bridge.py

"""Bridge protocol — communication between Discord bot and Antigravity.

Bridge directory: ~/.gemini/antigravity/bridge/
Structure:
  bridge/
    pending/     ← Bot writes approval requests for Discord
    response/    ← Bot writes user responses from Discord
    commands/    ← Bot writes user text input from Discord

Protocol:
1. VS Code Extension detects pending approval → writes JSON to pending/
2. Bot reads pending/ → sends Discord message with ✅/❌ buttons
3. User clicks button → Bot writes JSON to response/
4. VS Code Extension reads response/ → executes action

Transport layer:
  LocalTransport — file-based (default, single-PC)
  RemoteTransport — HTTP-based (future: multi-PC collector mode)
"""

import json
import time
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from dataclasses import dataclass, asdict
from enum import Enum

from config import Config

logger = logging.getLogger(__name__)


class ApprovalStatus(Enum):
    PENDING = "pending"
    APPROVED = "approved"
    REJECTED = "rejected"
    TIMEOUT = "timeout"


@dataclass
class ApprovalRequest:
    """An approval request from Antigravity."""
    request_id: str
    conversation_id: str
    command: str           # The command/action needing approval
    description: str       # Human-readable description
    timestamp: float
    status: str = "pending"
    discord_message_id: int = 0
    project_name: str = ""  # Project routing key
    step_type: str = ""  # e.g. 'diff_review', passed through to response


@dataclass
class UserResponse:
    """A user response from Discord."""
    request_id: str
    approved: bool
    user_input: str = ""
    timestamp: float = 0
    button_index: int = -1  # -1 = legacy (approve/reject), 0+ = specific button index
    step_type: str = ""  # pass through from pending for extension routing
    project_name: str = ""  # for multi-project: extension uses this when pending file is missing


# ─── Transport Abstraction ───

class BridgeTransport(ABC):
    """Abstract transport for bridge I/O.

    Implementations handle reading/writing JSON files for the bridge protocol,
    regardless of whether the storage is local filesystem or remote HTTP.
    """

    @abstractmethod
    def list_json_files(self, subdir: str) -> list[str]:
        """List JSON filenames in a subdirectory (e.g. 'pending', 'response')."""
        ...

    @abstractmethod
    def read_json(self, subdir: str, filename: str) -> dict | None:
        """Read and parse a JSON file. Returns None if not found or corrupt."""
        ...

    @abstractmethod
    def write_json(self, subdir: str, filename: str, data: dict) -> None:
        """Write data as JSON to a file in the given subdirectory."""
        ...

    @abstractmethod
    def delete_file(self, subdir: str, filename: str) -> bool:
        """Delete a file. Returns True if deleted, False if not found."""
        ...

    @abstractmethod
    def ensure_dirs(self) -> None:
        """Ensure all required subdirectories exist."""
        ...


class LocalTransport(BridgeTransport):
    """File-system based transport (default, single-PC mode).

    Reads/writes directly to the bridge directory on local disk.
    This is the existing behavior, extracted into a transport class.
    """

    def __init__(self, bridge_dir: Path):
        self.bridge_dir = bridge_dir

    def list_json_files(self, subdir: str) -> list[str]:
        d = self.bridge_dir / subdir
        if not d.exists():
            return []
        return [f.name for f in d.glob("*.json")]

    def read_json(self, subdir: str, filename: str) -> dict | None:
        fp = self.bridge_dir / subdir / filename
        if not fp.exists():
            return None
        try:
            return json.loads(fp.read_text(encoding="utf-8-sig"))
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"LocalTransport: bad file {subdir}/{filename}: {e}")
            return None

    def write_json(self, subdir: str, filename: str, data: dict) -> None:
        d = self.bridge_dir / subdir
        d.mkdir(parents=True, exist_ok=True)
        fp = d / filename
        fp.write_text(
            json.dumps(data, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    def delete_file(self, subdir: str, filename: str) -> bool:
        fp = self.bridge_dir / subdir / filename
        if fp.exists():
            try:
                fp.unlink()
                return True
            except OSError:
                return False
        return False

    def ensure_dirs(self) -> None:
        for sub in ("pending", "response", "commands"):
            (self.bridge_dir / sub).mkdir(parents=True, exist_ok=True)


class RemoteTransport(BridgeTransport):
    """HTTP-based transport for Collector → Gateway communication.

    Maps BridgeTransport methods to Gateway API endpoints:
      list_json_files("pending")  → GET /api/pending (returns list)
      write_json("pending", ...)  → POST /api/pending
      read_json("response", ...)  → GET /api/response/{rid}
      write_json("commands", ...) → (not used by Collector, Gateway pushes commands)
      etc.
    """

    def __init__(self, base_url: str, api_key: str = ""):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self._headers = {"Content-Type": "application/json"}
        if api_key:
            self._headers["Authorization"] = f"Bearer {api_key}"
        logger.info(f"RemoteTransport: {self.base_url} (auth={'yes' if api_key else 'no'})")

    def _request(self, method: str, path: str, data: dict | None = None) -> dict | None:
        """Make HTTP request to Gateway API."""
        import urllib.request
        import urllib.error

        url = f"{self.base_url}{path}"
        body = json.dumps(data, ensure_ascii=False).encode("utf-8") if data else None
        req = urllib.request.Request(url, data=body, headers=self._headers, method=method)

        try:
            with urllib.request.urlopen(req, timeout=10) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            logger.warning(f"RemoteTransport: {method} {path} → {e.code} {e.reason}")
            return None
        except (urllib.error.URLError, OSError, json.JSONDecodeError) as e:
            logger.warning(f"RemoteTransport: {method} {path} → {e}")
            return None

    def list_json_files(self, subdir: str) -> list[str]:
        """List pending requests from Gateway."""
        if subdir == "pending":
            result = self._request("GET", "/api/pending")
            if result and isinstance(result, list):
                return [f"{r['request_id']}.json" for r in result]
        elif subdir == "commands":
            # Commands are polled per-project (handled separately)
            return []
        return []

    def read_json(self, subdir: str, filename: str) -> dict | None:
        """Read a JSON file from Gateway."""
        rid = filename.replace(".json", "")
        if subdir == "response":
            return self._request("GET", f"/api/response/{rid}")
        elif subdir == "pending":
            # Pending data comes from list, not individual read
            result = self._request("GET", "/api/pending")
            if result and isinstance(result, list):
                for r in result:
                    if r.get("request_id") == rid:
                        return r
        return None

    def write_json(self, subdir: str, filename: str, data: dict) -> None:
        """Write data to Gateway via API."""
        if subdir == "pending":
            self._request("POST", "/api/pending", data)
        elif subdir == "response":
            rid = data.get("request_id", filename.replace(".json", ""))
            self._request("POST", f"/api/response/{rid}", data)
        elif subdir == "commands":
            # Commands go through write_command in BridgeProtocol
            self._request("POST", "/api/chat", data)

    def delete_file(self, subdir: str, filename: str) -> bool:
        """Delete not needed for remote — Gateway manages cleanup."""
        return True

    def ensure_dirs(self) -> None:
        """No local dirs needed for remote transport."""
        pass

    def poll_commands(self, project: str) -> list[dict]:
        """Poll Gateway for commands (Collector-specific, not in ABC)."""
        result = self._request("GET", f"/api/commands/{project}")
        if result and isinstance(result, dict):
            return result.get("commands", [])
        return []

    def register_session(self, conversation_id: str, project_name: str) -> None:
        """Register session → project mapping on Gateway."""
        self._request("POST", "/api/register", {
            "conversation_id": conversation_id,
            "project_name": project_name,
        })

    def send_chat(self, project_name: str, content: str) -> None:
        """Push chat snapshot to Gateway for relay to Discord."""
        self._request("POST", "/api/chat", {
            "project_name": project_name,
            "content": content,
        })


# ─── Bridge Protocol (uses Transport) ───

class BridgeProtocol:
    """Manages the bridge protocol via a pluggable transport."""

    def __init__(self, transport: BridgeTransport | None = None):
        if transport is None:
            bridge_dir = Config.BRAIN_PATH.parent / "bridge"
            transport = LocalTransport(bridge_dir)
        self.transport = transport

        # Legacy attributes for backward compatibility
        # (bot.py uses self.bridge.pending_dir etc. in some places)
        if isinstance(transport, LocalTransport):
            self.bridge_dir = transport.bridge_dir
            self.pending_dir = transport.bridge_dir / "pending"
            self.response_dir = transport.bridge_dir / "response"
            self.commands_dir = transport.bridge_dir / "commands"

        # Ensure directories exist
        self.transport.ensure_dirs()

        # Startup cleanup: purge stale pending files (> 5 min old)
        self._cleanup_stale_pending()

        logger.info(f"Bridge protocol initialized: transport={type(transport).__name__}")

    def _cleanup_stale_pending(self, max_age_seconds: int = 300):
        """Remove pending files older than max_age_seconds on startup."""
        now = time.time()
        cleaned = 0
        for fname in self.transport.list_json_files("pending"):
            data = self.transport.read_json("pending", fname)
            if data is None:
                self.transport.delete_file("pending", fname)
                cleaned += 1
                continue
            ts = data.get("timestamp", 0)
            if now - ts > max_age_seconds:
                self.transport.delete_file("pending", fname)
                cleaned += 1
        if cleaned:
            logger.info(f"Startup cleanup: removed {cleaned} stale pending files")

    def get_pending_requests(self) -> list[ApprovalRequest]:
        """Read all pending approval requests. Skips files older than 30 minutes."""
        requests = []
        fields = {f.name for f in ApprovalRequest.__dataclass_fields__.values()}
        now = time.time()
        MAX_AGE = 1800  # 30 minutes (matches Discord button timeout)
        for fname in self.transport.list_json_files("pending"):
            data = self.transport.read_json("pending", fname)
            if data is None:
                continue
            ts = data.get("timestamp", 0)
            if now - ts > MAX_AGE:
                # Too old — mark expired and skip
                data["status"] = "expired"
                self.transport.write_json("pending", fname, data)
                continue
            if data.get("status") == "pending":
                # Filter to known fields only
                filtered = {k: v for k, v in data.items() if k in fields}
                try:
                    requests.append(ApprovalRequest(**filtered))
                except TypeError as e:
                    logger.warning(f"Bad pending request {fname}: {e}")
        return requests

    def read_pending_request(self, request_id: str) -> ApprovalRequest | None:
        """Re-read a specific pending request (to get merged data)."""
        fname = f"{request_id}.json"
        data = self.transport.read_json("pending", fname)
        if data is None:
            return None
        fields = {fn.name for fn in ApprovalRequest.__dataclass_fields__.values()}
        filtered = {k: v for k, v in data.items() if k in fields}
        try:
            return ApprovalRequest(**filtered)
        except TypeError:
            return None

    def write_response(self, response: UserResponse):
        """Write a user response to the response directory."""
        response.timestamp = time.time()
        fname = f"{response.request_id}.json"

        self.transport.write_json("response", fname, asdict(response))
        logger.info(f"Response written: {fname} (approved={response.approved})")

        # Delete pending file after processing (prevents re-processing and accumulation)
        self.transport.delete_file("pending", fname)

    def write_command(self, conversation_id: str, text: str, *, project_name: str = ""):
        """Write a user text command for Antigravity to consume."""
        cmd_id = f"{int(time.time() * 1000)}"
        fname = f"{cmd_id}.json"

        data = {
            "id": cmd_id,
            "conversation_id": conversation_id,
            "project_name": project_name,
            "text": text,
            "timestamp": time.time(),
            "consumed": False,
        }

        self.transport.write_json("commands", fname, data)
        logger.info(f"Command written: {cmd_id} → project={project_name}")
        return cmd_id