"""Main agent loop for the Bamboo Claw client.

Manages both gateway and server connections, processes commands,
and handles reconnection for both sides.
"""

import asyncio
import logging
import signal
from typing import Type

from bamboo_claw_client.config import ClientConfig, DEFAULT_CONFIG_DIR
from bamboo_claw_client.gateway import GatewayClient
from bamboo_claw_client.server_ws import ServerConnection

logger = logging.getLogger(__name__)
MAX_SERVER_BACKOFF = 60
AUTH_ERROR_MARKERS = (
    "token revoked",
    "token expired",
    "invalid token",
    "authentication failed",
    "forbidden",
)
REGISTRATION_OUTCOMES = {"approved", "retry", "denied"}


async def run_agent(
    config: ClientConfig,
    *,
    shutdown_event: asyncio.Event | None = None,
    enable_signal_handlers: bool = True,
    gateway_factory: Type[GatewayClient] = GatewayClient,
    server_factory: Type[ServerConnection] = ServerConnection,
):
    """Main entry point for the client agent."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    gateway = gateway_factory(
        url=config.gateway_url,
        token=config.gateway_token,
        data_dir=str(DEFAULT_CONFIG_DIR),
    )

    server = server_factory(
        server_url=config.server_url,
        token=config.token,
        client_id=config.client_id,
        hostname=config.hostname,
    )

    shutdown_event = shutdown_event or asyncio.Event()

    if enable_signal_handlers:
        loop = asyncio.get_event_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                loop.add_signal_handler(sig, shutdown_event.set)
            except (NotImplementedError, RuntimeError):
                pass

    gateway_connected = False
    gateway_reconnect_task: asyncio.Task | None = None
    try:
        await gateway.connect()
        gateway_connected = True
        logger.info("Connected to OpenClaw gateway at %s", config.gateway_url)
    except Exception as e:
        logger.warning("Could not connect to gateway: %s (will retry in background)", e)
        gateway_reconnect_task = asyncio.create_task(_gateway_reconnect_bg(gateway, server, config))

    backoff = 1

    try:
        while not shutdown_event.is_set():
            try:
                client_id = await server.connect()
                backoff = 1
                logger.info("Connected to server as %s", client_id)

                if config.client_id != client_id:
                    config.client_id = client_id
                    try:
                        config.save()
                    except Exception as exc:
                        logger.warning("Failed to persist updated client_id: %s", exc)
            except Exception as e:
                error_msg = str(e).lower()
                if _is_auth_error(error_msg):
                    logger.warning("Token appears invalid or expired. Attempting re-registration...")
                    config.token = ""
                    try:
                        config.save()
                    except Exception as exc:
                        logger.warning("Failed to persist cleared token state: %s", exc)
                    if not await _recover_auth(config=config, server=server, shutdown_event=shutdown_event):
                        if not shutdown_event.is_set():
                            logger.error(
                                "Re-registration denied or unavailable. "
                                "Run 'bambooclaw setup' and restart the client."
                            )
                        break
                    backoff = 1
                    continue
                logger.warning("Failed to connect to server: %s", e)
                if shutdown_event.is_set():
                    break
                await _wait_or_shutdown(shutdown_event, backoff)
                backoff = min(backoff * 2, MAX_SERVER_BACKOFF)
                continue

            try:
                gateway_info = {
                    "connected": gateway.connected or gateway_connected,
                    "info": {"url": config.gateway_url} if (gateway.connected or gateway_connected) else None,
                }
                await server.register(gateway_info=gateway_info)
                await server.start_heartbeat()

                async def handle_command(data: dict):
                    request_id = data.get("id", "")
                    command = data.get("command", "")
                    params = data.get("params", {})

                    if command == "chat":
                        await _handle_chat_command(gateway, server, request_id, params)
                    else:
                        await server.send_command_result(
                            request_id, ok=False, error=f"Unknown command: {command}"
                        )

                listen_task = asyncio.create_task(server.listen(handle_command))
                shutdown_task = asyncio.create_task(shutdown_event.wait())
                dropped_task = asyncio.create_task(_wait_for_server_disconnect(server))
                done, pending = await asyncio.wait(
                    [listen_task, shutdown_task, dropped_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )
                for task in pending:
                    task.cancel()
                await asyncio.gather(*pending, return_exceptions=True)

                await server.disconnect()

                if shutdown_task in done or shutdown_event.is_set():
                    break

                logger.warning("Server connection dropped. Reconnecting with backoff...")
            except asyncio.CancelledError:
                raise
            except Exception as e:
                logger.warning("Server loop error: %s", e)

            if shutdown_event.is_set():
                break
            await _wait_or_shutdown(shutdown_event, backoff)
            backoff = min(backoff * 2, MAX_SERVER_BACKOFF)
    finally:
        logger.info("Shutting down...")
        if gateway_reconnect_task and not gateway_reconnect_task.done():
            gateway_reconnect_task.cancel()
            await asyncio.gather(gateway_reconnect_task, return_exceptions=True)
        await server.disconnect()
        await gateway.disconnect()


async def _handle_chat_command(
    gateway: GatewayClient,
    server: ServerConnection,
    request_id: str,
    params: dict,
):
    """Handle a chat command by forwarding to the gateway and relaying events."""
    if not gateway.connected:
        await server.send_command_result(
            request_id, ok=False, error="Gateway not connected"
        )
        return

    message = params.get("message", "")
    session_key = params.get("session_key")
    attachments = params.get("attachments")

    async def on_event(event: dict):
        await server.send_chat_event(request_id, event)

    try:
        result = await gateway.send_chat_message(
            message=message,
            session_key=session_key,
            attachments=attachments,
            on_event=on_event,
        )
        await server.send_command_result(request_id, ok=True, result=result.get("content", ""))
    except Exception as e:
        logger.warning("Chat command error: %s", e)
        await server.send_chat_event(request_id, {"state": "error", "error": str(e)})


async def _gateway_reconnect_bg(
    gateway: GatewayClient,
    server: ServerConnection,
    config: ClientConfig,
):
    """Background task to reconnect to the gateway."""
    try:
        await gateway.reconnect_loop()
        await server.send_gateway_status(True)
        logger.info("Gateway reconnected in background")
    except asyncio.CancelledError:
        pass
    except Exception as e:
        logger.warning("Gateway reconnect background task failed: %s", e)


def _is_auth_error(error_msg: str) -> bool:
    return any(marker in error_msg for marker in AUTH_ERROR_MARKERS)


async def _run_registration(config: ClientConfig) -> object:
    from bamboo_claw_client.registration import run_registration

    return await run_registration(config)


def _normalize_registration_outcome(outcome: object, config: ClientConfig) -> str:
    if isinstance(outcome, str):
        normalized = outcome.strip().lower()
        if normalized in REGISTRATION_OUTCOMES:
            return normalized
    if outcome is True:
        return "approved"
    if outcome is False:
        return "retry"
    return "approved" if config.token else "retry"


async def _recover_auth(
    config: ClientConfig,
    server: ServerConnection,
    shutdown_event: asyncio.Event,
) -> bool:
    backoff = 1
    while not shutdown_event.is_set():
        outcome = _normalize_registration_outcome(await _run_registration(config), config)
        if outcome == "approved":
            if not config.token:
                logger.warning("Registration reported approved but no token persisted; retrying...")
            else:
                server.token = config.token
                return True
        elif outcome == "denied":
            return False

        logger.warning("Re-registration not approved yet; retrying in %ds", backoff)
        await _wait_or_shutdown(shutdown_event, backoff)
        backoff = min(backoff * 2, MAX_SERVER_BACKOFF)
    return False


async def _wait_or_shutdown(shutdown_event: asyncio.Event, seconds: int) -> None:
    try:
        await asyncio.wait_for(shutdown_event.wait(), timeout=seconds)
    except asyncio.TimeoutError:
        return


async def _wait_for_server_disconnect(server: ServerConnection) -> None:
    while server.connected:
        await asyncio.sleep(0.25)
