"""Core agent logic — connects to SUSATest server, executes commands on local device.

Flow:
1. Discover local ADB devices
2. Connect to server via WebSocket
3. Send heartbeat with device list
4. Receive commands, execute via ADB, send responses
5. Reconnect on disconnect with exponential backoff
"""
from __future__ import annotations

import asyncio
import json
import logging
import sys
import time

import websockets

from .adb_executor import ADBExecutor

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [susatest] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)


async def run_agent(
    api_key: str,
    server_url: str = "wss://www.susatest.com",
    device_serial: str | None = None,
    adb_path: str = "adb",
):
    """Main agent entry point — connects to server and processes commands."""
    executor = ADBExecutor(adb_path=adb_path, device_serial=device_serial)

    # Discover local devices
    devices = await executor.get_devices()
    if not devices:
        logger.error("No Android devices found. Connect a device via USB and enable USB debugging.")
        sys.exit(1)

    if device_serial:
        devices = [d for d in devices if d["serial"] == device_serial]
        if not devices:
            logger.error("Device %s not found. Available: %s",
                         device_serial, [d["serial"] for d in await executor.get_devices()])
            sys.exit(1)

    logger.info("Found %d device(s): %s", len(devices), [d["serial"] for d in devices])

    # If no specific device, use first one
    if not device_serial and len(devices) == 1:
        executor._device_serial = devices[0]["serial"]

    # Connect with reconnection
    backoff = 1
    max_backoff = 30
    ws_url = f"{server_url}/ws/agent?token={api_key}"

    while True:
        try:
            logger.info("Connecting to %s ...", server_url)
            async with websockets.connect(
                ws_url,
                ping_interval=20,
                ping_timeout=30,
                max_size=50 * 1024 * 1024,  # 50MB for APK transfer
            ) as ws:
                backoff = 1  # Reset on successful connect
                logger.info("Connected! Sending device list...")

                # Send initial heartbeat
                await ws.send(json.dumps({
                    "type": "heartbeat",
                    "agent_version": "0.1.0",
                    "devices": devices,
                }))

                # Wait for ACK
                ack = json.loads(await ws.recv())
                if ack.get("type") == "heartbeat_ack":
                    logger.info("Registered: %s", ack.get("message", "OK"))
                else:
                    logger.error("Unexpected response: %s", ack)
                    continue

                logger.info("Waiting for commands... (Ctrl+C to disconnect)")

                # Start heartbeat task
                heartbeat_task = asyncio.create_task(_heartbeat_loop(ws, executor, devices))

                try:
                    await _command_loop(ws, executor)
                finally:
                    heartbeat_task.cancel()
                    try:
                        await heartbeat_task
                    except asyncio.CancelledError:
                        pass

        except websockets.exceptions.ConnectionClosedError as e:
            if e.code == 4001:
                logger.error("Authentication failed. Check your API key.")
                sys.exit(1)
            logger.warning("Connection closed: %s. Reconnecting in %ds...", e, backoff)
        except (ConnectionRefusedError, OSError) as e:
            logger.warning("Cannot connect: %s. Retrying in %ds...", e, backoff)
        except Exception as e:
            logger.error("Unexpected error: %s. Retrying in %ds...", e, backoff)

        await asyncio.sleep(backoff)
        backoff = min(backoff * 2, max_backoff)


async def _heartbeat_loop(ws, executor: ADBExecutor, devices: list[dict]):
    """Send periodic heartbeats to keep connection alive."""
    while True:
        await asyncio.sleep(30)
        try:
            # Refresh device list
            fresh_devices = await executor.get_devices()
            if fresh_devices:
                devices.clear()
                devices.extend(fresh_devices)
            await ws.send(json.dumps({
                "type": "heartbeat",
                "agent_version": "0.1.0",
                "devices": devices,
            }))
        except Exception:
            break


async def _command_loop(ws, executor: ADBExecutor):
    """Receive and execute commands from server."""
    async for raw in ws:
        msg = json.loads(raw)
        msg_type = msg.get("type")

        if msg_type == "command":
            # Execute command asynchronously
            asyncio.create_task(_handle_command(ws, executor, msg))

        elif msg_type == "file_transfer_start":
            executor.handle_file_transfer_start(msg)

        elif msg_type == "file_chunk":
            executor.handle_file_chunk(msg)

        elif msg_type == "file_transfer_end":
            path = executor.handle_file_transfer_end(msg)
            if path:
                logger.info("File received: %s (%.1f MB)", path.name, path.stat().st_size / 1e6)

        elif msg_type == "heartbeat_ack":
            pass  # Server acknowledged our heartbeat

        else:
            logger.warning("Unknown message type: %s", msg_type)


async def _handle_command(ws, executor: ADBExecutor, msg: dict):
    """Execute a single command and send the response."""
    req_id = msg["id"]
    method = msg["method"]
    args = msg.get("args", {})

    t0 = time.monotonic()
    try:
        result = await executor.execute(method, args)
        elapsed = time.monotonic() - t0

        # Log with data size if applicable
        size_info = ""
        if "data_b64" in result:
            size_kb = len(result["data_b64"]) * 3 / 4 / 1024
            size_info = f" ({size_kb:.0f}KB)"

        success = result.get("success", True)
        status = "OK" if success else "FAIL"
        logger.info("[%s] %s(%s) -> %s%s (%.1fs)",
                     time.strftime("%H:%M:%S"), method,
                     _format_args(method, args), status, size_info, elapsed)

        await ws.send(json.dumps({
            "id": req_id,
            "type": "response",
            "success": result.get("success", True),
            "result": result,
        }))
    except Exception as e:
        elapsed = time.monotonic() - t0
        logger.error("[%s] %s -> ERROR: %s (%.1fs)",
                     time.strftime("%H:%M:%S"), method, e, elapsed)
        await ws.send(json.dumps({
            "id": req_id,
            "type": "response",
            "success": False,
            "error": str(e),
        }))


def _format_args(method: str, args: dict) -> str:
    """Format args for logging (exclude large data)."""
    if method == "tap":
        return f"{args.get('x')}, {args.get('y')}"
    if method == "swipe":
        return f"{args.get('x1')},{args.get('y1')} -> {args.get('x2')},{args.get('y2')}"
    if method == "type_text":
        text = args.get("text", "")
        return f"'{text[:20]}...'" if len(text) > 20 else f"'{text}'"
    if method in ("launch_app", "force_stop", "check_crash"):
        return args.get("identifier", "")
    return ""
