"""Core agent logic — connects to SUSATest server via HTTP long polling.

Flow:
1. Discover local ADB devices
2. Register with server via POST /api/agent/register
3. Poll for commands via GET /api/agent/poll (blocks up to 30s)
4. Execute commands via ADB, post results via POST /api/agent/result
5. Send periodic heartbeats via POST /api/agent/heartbeat
6. Re-register on connection loss with exponential backoff
"""
from __future__ import annotations

import asyncio
import json
import logging
import sys
import time

import httpx

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 = "https://www.susatest.com",
    device_serial: str | None = None,
    adb_path: str = "adb",
):
    """Main agent entry point — registers and polls for 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 not device_serial and len(devices) == 1:
        executor._device_serial = devices[0]["serial"]

    # Normalize server URL
    base_url = server_url.rstrip("/")
    if base_url.startswith("wss://"):
        base_url = base_url.replace("wss://", "https://")
    elif base_url.startswith("ws://"):
        base_url = base_url.replace("ws://", "http://")

    backoff = 1
    max_backoff = 30

    while True:
        try:
            async with httpx.AsyncClient(
                base_url=base_url,
                timeout=httpx.Timeout(60, connect=10),
            ) as client:
                # Register
                logger.info("Connecting to %s ...", base_url)
                resp = await client.post("/api/agent/register", json={
                    "token": api_key,
                    "agent_version": "0.2.0",
                    "devices": devices,
                })

                if resp.status_code == 401:
                    logger.error("Authentication failed. Check your API key.")
                    sys.exit(1)
                resp.raise_for_status()

                data = resp.json()
                agent_id = data["agent_id"]
                logger.info("Connected! %s", data["message"])
                logger.info("Waiting for commands... (Ctrl+C to disconnect)")
                backoff = 1

                # Start heartbeat in background
                heartbeat_task = asyncio.create_task(
                    _heartbeat_loop(client, agent_id, executor, devices)
                )

                try:
                    await _poll_loop(client, agent_id, executor)
                finally:
                    heartbeat_task.cancel()
                    try:
                        await heartbeat_task
                    except asyncio.CancelledError:
                        pass
                    # Graceful disconnect
                    try:
                        await client.post("/api/agent/disconnect", json={"agent_id": agent_id})
                    except Exception:
                        pass

        except httpx.HTTPStatusError as e:
            if e.response.status_code == 401:
                logger.error("Authentication failed. Check your API key.")
                sys.exit(1)
            logger.warning("Server error: %s. Retrying in %ds...", e, backoff)
        except (httpx.ConnectError, httpx.ConnectTimeout, 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(
    client: httpx.AsyncClient,
    agent_id: str,
    executor: ADBExecutor,
    devices: list[dict],
):
    """Send periodic heartbeats to keep registration alive."""
    while True:
        await asyncio.sleep(30)
        try:
            fresh_devices = await executor.get_devices()
            if fresh_devices:
                devices.clear()
                devices.extend(fresh_devices)
            resp = await client.post("/api/agent/heartbeat", json={
                "agent_id": agent_id,
                "devices": devices,
            })
            if resp.status_code == 404:
                logger.warning("Server lost our registration. Reconnecting...")
                return  # Exit to trigger re-register
        except Exception:
            break


async def _poll_loop(client: httpx.AsyncClient, agent_id: str, executor: ADBExecutor):
    """Poll for commands and execute them."""
    while True:
        try:
            resp = await client.get(
                "/api/agent/poll",
                params={"agent_id": agent_id},
                timeout=httpx.Timeout(45, connect=10),  # > server's 30s poll timeout
            )

            if resp.status_code == 404:
                logger.warning("Server lost our registration. Reconnecting...")
                return

            resp.raise_for_status()
            msg = resp.json()

            if msg.get("type") == "no_command":
                continue  # Poll again

            if msg.get("type") == "command":
                # Execute in background so we can keep polling
                asyncio.create_task(
                    _handle_command(client, agent_id, executor, msg)
                )

        except httpx.ReadTimeout:
            # Normal — server poll timed out, just retry
            continue
        except (httpx.ConnectError, httpx.ConnectTimeout):
            raise  # Let outer loop handle reconnection


async def _handle_command(
    client: httpx.AsyncClient,
    agent_id: str,
    executor: ADBExecutor,
    msg: dict,
):
    """Execute a single command and post the result back."""
    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

        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 client.post("/api/agent/result", json={
            "agent_id": agent_id,
            "id": req_id,
            "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 client.post("/api/agent/result", json={
            "agent_id": agent_id,
            "id": req_id,
            "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 ""
