xml-pipeline/agentserver/console/client.py

"""
Console client that connects to the agent server.

Provides SSH-style login with username/password authentication.
"""

from __future__ import annotations

import asyncio
import getpass
import json
import sys
from typing import Optional

try:
    import aiohttp
    AIOHTTP_AVAILABLE = True
except ImportError:
    AIOHTTP_AVAILABLE = False

try:
    from prompt_toolkit import PromptSession
    from prompt_toolkit.history import InMemoryHistory
    from prompt_toolkit.styles import Style
    PROMPT_TOOLKIT_AVAILABLE = True
except ImportError:
    PROMPT_TOOLKIT_AVAILABLE = False


DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8765
MAX_LOGIN_ATTEMPTS = 3


class ConsoleClient:
    """
    Text-based console client for the agent server.

    Usage:
        client = ConsoleClient()
        asyncio.run(client.run())
    """

    def __init__(self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT):
        self.host = host
        self.port = port
        self.base_url = f"http://{host}:{port}"
        self.ws_url = f"ws://{host}:{port}/ws"
        self.token: Optional[str] = None
        self.username: Optional[str] = None
        self.session: Optional[aiohttp.ClientSession] = None
        self.ws: Optional[aiohttp.ClientWebSocketResponse] = None
        self.running = False

    async def login(self) -> bool:
        """
        Perform SSH-style login.

        Returns:
            True if login successful, False otherwise
        """
        print(f"Connecting to {self.host}:{self.port}...")

        for attempt in range(1, MAX_LOGIN_ATTEMPTS + 1):
            try:
                username = input("Username: ")
                password = getpass.getpass("Password: ")
            except (EOFError, KeyboardInterrupt):
                print("\nLogin cancelled.")
                return False

            if not username or not password:
                print("Username and password required.")
                continue

            try:
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        f"{self.base_url}/auth/login",
                        json={"username": username, "password": password},
                    ) as resp:
                        data = await resp.json()

                        if resp.status == 200:
                            self.token = data["token"]
                            self.username = username
                            print(f"Welcome, {username}!")
                            return True
                        else:
                            error = data.get("error", "Authentication failed")
                            remaining = MAX_LOGIN_ATTEMPTS - attempt
                            if remaining > 0:
                                print(f"{error}. {remaining} attempt(s) remaining.")
                            else:
                                print(f"{error}. No attempts remaining.")
            except aiohttp.ClientError as e:
                print(f"Connection error: {e}")
                return False

        return False

    async def connect_ws(self) -> bool:
        """Connect to WebSocket after authentication."""
        if not self.token:
            return False

        try:
            self.session = aiohttp.ClientSession(
                headers={"Authorization": f"Bearer {self.token}"}
            )
            self.ws = await self.session.ws_connect(self.ws_url)

            # Wait for connected message
            msg = await self.ws.receive_json()
            if msg.get("type") == "connected":
                return True
            return False
        except Exception as e:
            print(f"WebSocket connection failed: {e}")
            return False

    async def send_command(self, cmd: str) -> Optional[dict]:
        """Send a command via WebSocket and get response."""
        if not self.ws:
            return None

        await self.ws.send_json(cmd)
        return await self.ws.receive_json()

    def print_help(self):
        """Print available commands."""
        print("""
Available commands:
  /help       - Show this help
  /status     - Show server status
  /listeners  - List active listeners
  /quit       - Disconnect and exit

Send messages:
  @listener message   - Send message to a listener
  message            - Send to default listener
""")

    async def handle_command(self, line: str) -> bool:
        """
        Handle a command line.

        Returns:
            False if should quit, True otherwise
        """
        line = line.strip()
        if not line:
            return True

        if line == "/help":
            self.print_help()
        elif line == "/quit" or line == "/exit":
            return False
        elif line == "/status":
            resp = await self.send_command({"type": "status"})
            if resp:
                threads = resp.get("threads", 0)
                print(f"Active threads: {threads}")
        elif line == "/listeners":
            resp = await self.send_command({"type": "listeners"})
            if resp:
                listeners = resp.get("listeners", [])
                if listeners:
                    print("Active listeners:")
                    for name in listeners:
                        print(f"  - {name}")
                else:
                    print("No active listeners")
        elif line.startswith("/"):
            print(f"Unknown command: {line}")
        else:
            # Send as message
            # TODO: Implement message sending when pump is connected
            print(f"Message sending not yet implemented: {line}")

        return True

    async def run(self):
        """Main client loop."""
        if not AIOHTTP_AVAILABLE:
            print("Error: aiohttp not installed")
            sys.exit(1)

        # Login
        if not await self.login():
            print("Authentication failed.")
            sys.exit(1)

        # Connect WebSocket
        if not await self.connect_ws():
            print("Failed to connect to server.")
            sys.exit(1)

        print("Connected. Type /help for commands, /quit to exit.")

        self.running = True

        try:
            if PROMPT_TOOLKIT_AVAILABLE:
                await self._run_prompt_toolkit()
            else:
                await self._run_simple()
        finally:
            await self.cleanup()

    async def _run_prompt_toolkit(self):
        """Run with prompt_toolkit for better UX."""
        style = Style.from_dict({
            "prompt": "ansicyan bold",
        })

        session = PromptSession(
            history=InMemoryHistory(),
            style=style,
        )

        while self.running:
            try:
                line = await asyncio.get_event_loop().run_in_executor(
                    None,
                    lambda: session.prompt(f"{self.username}> ")
                )
                if not await self.handle_command(line):
                    break
            except (EOFError, KeyboardInterrupt):
                break

    async def _run_simple(self):
        """Run with simple input (fallback)."""
        while self.running:
            try:
                line = input(f"{self.username}> ")
                if not await self.handle_command(line):
                    break
            except (EOFError, KeyboardInterrupt):
                break

    async def cleanup(self):
        """Clean up connections."""
        if self.ws:
            await self.ws.close()
        if self.session:
            await self.session.close()
        print("Disconnected.")


def main():
    """Entry point."""
    import argparse

    parser = argparse.ArgumentParser(description="XML Pipeline Console")
    parser.add_argument("--host", default=DEFAULT_HOST, help="Server host")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="Server port")
    args = parser.parse_args()

    client = ConsoleClient(host=args.host, port=args.port)
    asyncio.run(client.run())


if __name__ == "__main__":
    main()