"""Client registration flow for first-time setup without a token."""

import asyncio
import json
import logging
import platform
from typing import Literal

from websockets import connect as ws_connect
from websockets.exceptions import ConnectionClosed

from bamboo_claw_client.config import ClientConfig

logger = logging.getLogger(__name__)

RegistrationOutcome = Literal["approved", "retry", "denied"]


async def run_registration(config: ClientConfig) -> RegistrationOutcome:
    """Connect to server without token and wait for admin approval.

    Returns:
      - "approved": token received and persisted
      - "retry": temporary failure (network/server), caller may retry
      - "denied": admin rejected/expired, requires manual setup
    """
    hostname = config.hostname or platform.node()
    client_id = config.ensure_client_id()
    ws_url = config.ws_url

    print(f"\nConnecting to {ws_url}...")
    try:
        ws = await ws_connect(ws_url, ping_interval=None, ping_timeout=None)
    except Exception as e:
        print(f"Failed to connect: {e}")
        return "retry"

    try:
        await ws.send(json.dumps({
            "type": "register_request",
            "hostname": hostname,
            "client_id": client_id,
            "ip": None,
            "gateway_info": {
                "url": config.gateway_url,
            },
        }))

        response = json.loads(await ws.recv())

        if response.get("type") == "error":
            print(f"Server error: {response.get('message')}")
            return "retry"

        if response.get("type") != "pending":
            print(f"Unexpected response: {response}")
            return "retry"

        approval_code = response["approval_code"]
        print(f"\nRegistration submitted!")
        print(f"  Hostname: {hostname}")
        print(f"  Approval Code: {approval_code}")
        print(f"\n  Ask your admin to run:")
        print(f"    POST /api/admin/registrations/{approval_code}/approve")
        print(f"\nWaiting for approval... (Ctrl+C to cancel)")

        while True:
            try:
                raw = await asyncio.wait_for(ws.recv(), timeout=30)
                data = json.loads(raw)
            except asyncio.TimeoutError:
                await ws.send(json.dumps({"type": "heartbeat"}))
                continue

            if data.get("type") == "heartbeat_ack":
                continue

            if data.get("type") == "approved":
                token = data["token"]
                config.token = token
                config.save()
                print(f"\nApproved! Token saved to config.")
                print(f"Run 'bambooclaw run' to start the agent.")
                return "approved"

            if data.get("type") == "rejected":
                print(f"\nRegistration rejected: {data.get('reason', 'No reason given')}")
                return "denied"

            if data.get("type") == "expired":
                print(f"\nRegistration expired. Please retry with 'bambooclaw setup'.")
                return "denied"

    except ConnectionClosed:
        print("\nConnection closed by server.")
        return "retry"
    except KeyboardInterrupt:
        print("\nRegistration cancelled.")
        return "denied"
    finally:
        try:
            await ws.close()
        except Exception:
            pass

    return "retry"
