"""Client CLI: argument parsing, socket communication, output streaming."""

import argparse
import asyncio
import json
import os
import struct
import sys
from pathlib import Path

from rpipe.protocol import (
    read_frame,
    pack_frame,
    TYPE_RUN_REQ,
    TYPE_STDIN,
    TYPE_CANCEL,
    TYPE_META,
    TYPE_STDOUT,
    TYPE_STDERR,
    TYPE_EXIT,
    TYPE_ERROR,
)


async def main():
    """Main client entry point."""
    ap = argparse.ArgumentParser(
        description="rpipe client: execute commands on remote hosts via daemon",
    )
    ap.add_argument("host", help="Remote hostname")
    ap.add_argument("cmd", nargs=argparse.REMAINDER, help="Command to execute")
    ap.add_argument("--cwd", help="Working directory on remote host")
    ap.add_argument(
        "--env",
        action="append",
        default=[],
        help="Environment variable (KEY=VALUE, repeatable)",
    )
    ap.add_argument("--stdin", action="store_true", help="Enable stdin piping")
    ap.add_argument("--timeout", type=int, default=60, help="Timeout in seconds")
    ap.add_argument(
        "--read-only",
        action="store_true",
        default=True,
        help="Enable read-only mode (block writes)",
    )
    ap.add_argument(
        "--socket",
        default=os.path.join(os.path.expanduser("~"), ".rpipe", "daemon.sock"),
        help="Path to daemon socket",
    )

    args = ap.parse_args()

    if not args.cmd:
        ap.error("command is required")

    # Connect to daemon
    sock_path = args.socket
    if not Path(sock_path).exists():
        print(f"error: daemon socket not found: {sock_path}", file=sys.stderr)
        print(f"hint: start daemon with: python -m rpipe.daemon", file=sys.stderr)
        raise SystemExit(1)

    try:
        reader, writer = await asyncio.open_unix_connection(sock_path)
    except Exception as e:
        print(f"error: failed to connect to daemon: {e}", file=sys.stderr)
        raise SystemExit(1)

    # Parse environment variables
    env = {}
    for kv in args.env:
        if "=" not in kv:
            print(f"error: invalid env format: {kv} (expected KEY=VALUE)", file=sys.stderr)
            raise SystemExit(1)
        k, v = kv.split("=", 1)
        env[k] = v

    # Build RUN_REQ
    req = {
        "host": args.host,
        "cmd": " ".join(args.cmd),
        "cwd": args.cwd,
        "env": env if env else {},
        "stdin_mode": "pipe" if args.stdin else "none",
        "timeout_s": args.timeout,
        "read_only": args.read_only,
    }

    # Send RUN_REQ frame
    writer.write(pack_frame(TYPE_RUN_REQ, 0, json.dumps(req).encode("utf-8")))
    await writer.drain()

    run_id = None
    exit_code = 1

    # Task to pump stdin from client to daemon
    async def pump_stdin():
        nonlocal run_id
        # Wait for run_id to be allocated by daemon
        while run_id is None:
            await asyncio.sleep(0.01)

        loop = asyncio.get_running_loop()
        while True:
            try:
                data = await loop.run_in_executor(None, sys.stdin.buffer.read, 8192)
                if not data:
                    break
                writer.write(pack_frame(TYPE_STDIN, run_id, data))
                await writer.drain()
            except Exception:
                break

    # Start stdin task if enabled
    stdin_task = asyncio.create_task(pump_stdin()) if args.stdin else None

    # Read response frames
    try:
        while True:
            ftype, rid, payload = await read_frame(reader)

            # First frame should be META with allocated run_id
            if ftype == TYPE_META:
                meta = json.loads(payload.decode("utf-8"))
                run_id = meta["run_id"]
                continue

            # Ignore frames for other runs (shouldn't happen in current architecture)
            if run_id is not None and rid != run_id:
                continue

            # Output frames
            if ftype == TYPE_STDOUT:
                sys.stdout.buffer.write(payload)
                sys.stdout.buffer.flush()

            elif ftype == TYPE_STDERR:
                sys.stderr.buffer.write(payload)
                sys.stderr.buffer.flush()

            # Error frame
            elif ftype == TYPE_ERROR:
                error_msg = json.loads(payload.decode("utf-8"))
                print(f"error: {error_msg.get('error', 'unknown error')}", file=sys.stderr)
                exit_code = 1

            # Exit code
            elif ftype == TYPE_EXIT:
                (exit_code,) = struct.unpack("!i", payload)
                break

    except Exception as e:
        print(f"error: communication error: {e}", file=sys.stderr)
        exit_code = 1

    finally:
        # Cancel stdin task
        if stdin_task:
            stdin_task.cancel()
            try:
                await stdin_task
            except asyncio.CancelledError:
                pass

        # Close connection
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass

    raise SystemExit(exit_code)


if __name__ == "__main__":
    asyncio.run(main())
