"""ADB command executor — runs ADB commands locally and returns results.

Each method matches the protocol expected by RemoteAgentController on the server.
"""
from __future__ import annotations

import asyncio
import base64
import logging
import os
import re
import tempfile
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


class ADBExecutor:
    """Executes ADB commands on a local device."""

    def __init__(self, adb_path: str = "adb", device_serial: str | None = None):
        self._adb_path = adb_path
        self._device_serial = device_serial
        self._screen_size: tuple[int, int] | None = None
        # Temp dir for file transfers
        self._transfer_dir = Path(tempfile.mkdtemp(prefix="susatest_"))
        # Active file transfers: transfer_id -> {chunks: list, metadata: dict}
        self._transfers: dict[str, dict] = {}

    async def _run(self, *args: str, timeout: float = 30) -> tuple[str, str, int]:
        """Run an ADB command and return (stdout, stderr, returncode)."""
        cmd = [self._adb_path]
        if self._device_serial:
            cmd.extend(["-s", self._device_serial])
        cmd.extend(args)

        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
            return (
                stdout.decode("utf-8", errors="replace"),
                stderr.decode("utf-8", errors="replace"),
                proc.returncode or 0,
            )
        except asyncio.TimeoutError:
            return "", f"Command timed out after {timeout}s", 1
        except Exception as e:
            return "", str(e), 1

    async def get_devices(self) -> list[dict]:
        """List connected ADB devices with model and Android version."""
        stdout, _, _ = await self._run("devices", "-l")
        devices = []
        for line in stdout.strip().splitlines()[1:]:
            parts = line.split()
            if len(parts) >= 2 and parts[1] == "device":
                serial = parts[0]
                model = ""
                for p in parts[2:]:
                    if p.startswith("model:"):
                        model = p.split(":", 1)[1]
                # Get Android version
                android_ver = ""
                try:
                    out, _, _ = await self._run("-s", serial, "shell", "getprop", "ro.build.version.release")
                    android_ver = out.strip()
                except Exception:
                    pass
                devices.append({
                    "serial": serial,
                    "model": model,
                    "android_version": android_ver,
                    "status": "online",
                })
        return devices

    # ── File transfer handling ──

    def handle_file_transfer_start(self, msg: dict):
        """Start receiving a file transfer."""
        tid = msg["transfer_id"]
        self._transfers[tid] = {
            "filename": msg["filename"],
            "total_size": msg["total_size"],
            "total_chunks": msg["total_chunks"],
            "sha256": msg.get("sha256", ""),
            "chunks": [None] * msg["total_chunks"],
            "received": 0,
        }
        logger.info("[transfer] Started: %s (%.1f MB, %d chunks)",
                     msg["filename"], msg["total_size"] / 1e6, msg["total_chunks"])

    def handle_file_chunk(self, msg: dict):
        """Receive a chunk of a file transfer."""
        tid = msg["transfer_id"]
        transfer = self._transfers.get(tid)
        if not transfer:
            return
        idx = msg["chunk_index"]
        transfer["chunks"][idx] = base64.b64decode(msg["data_b64"])
        transfer["received"] += 1

    def handle_file_transfer_end(self, msg: dict) -> Path | None:
        """Finalize file transfer, write to disk, return path."""
        tid = msg["transfer_id"]
        transfer = self._transfers.pop(tid, None)
        if not transfer:
            return None

        # Reassemble
        out_path = self._transfer_dir / transfer["filename"]
        with open(out_path, "wb") as f:
            for chunk in transfer["chunks"]:
                if chunk:
                    f.write(chunk)

        logger.info("[transfer] Complete: %s -> %s", transfer["filename"], out_path)
        return out_path

    # ── Command execution (matches server protocol) ──

    async def execute(self, method: str, args: dict) -> dict:
        """Execute a command and return the result dict."""
        # Override device serial from args if provided
        device = args.pop("device_serial", None) or self._device_serial
        old_serial = self._device_serial
        self._device_serial = device

        try:
            handler = getattr(self, f"cmd_{method}", None)
            if not handler:
                return {"success": False, "error": f"Unknown command: {method}"}
            return await handler(args)
        except Exception as e:
            logger.exception("Command %s failed", method)
            return {"success": False, "error": str(e)}
        finally:
            self._device_serial = old_serial

    async def cmd_take_screenshot(self, args: dict) -> dict:
        pid = os.getpid()
        device_path = f"/sdcard/susa_screen_{pid}.png"
        local_path = self._transfer_dir / f"screen_{pid}.png"

        await self._run("shell", "screencap", "-p", device_path)
        await self._run("pull", device_path, str(local_path))
        await self._run("shell", "rm", device_path)

        if local_path.exists():
            data = local_path.read_bytes()
            local_path.unlink()
            return {"success": True, "data_b64": base64.b64encode(data).decode()}
        return {"success": False, "error": "Screenshot failed"}

    async def cmd_get_ui_hierarchy_xml(self, args: dict) -> dict:
        pid = os.getpid()
        device_path = f"/sdcard/susa_ui_{pid}.xml"
        await self._run("shell", "rm", "-f", device_path)
        await self._run("shell", "uiautomator", "dump", device_path, timeout=15)
        stdout, _, _ = await self._run("shell", "cat", device_path)
        await self._run("shell", "rm", "-f", device_path)
        if stdout.strip():
            return {"xml": stdout.strip()}
        return {"xml": None}

    async def cmd_get_current_activity(self, args: dict) -> dict:
        stdout, _, _ = await self._run("shell", "dumpsys", "activity", "activities")
        for line in stdout.splitlines():
            if "mResumedActivity" in line or "mFocusedActivity" in line:
                m = re.search(r"(\S+/\S+)", line)
                if m:
                    return {"activity": m.group(1)}
        return {"activity": None}

    async def cmd_tap(self, args: dict) -> dict:
        await self._run("shell", "input", "tap", str(args["x"]), str(args["y"]))
        return {"success": True}

    async def cmd_long_press(self, args: dict) -> dict:
        x, y = args["x"], args["y"]
        duration = args.get("duration_ms", 1000)
        await self._run("shell", "input", "swipe", str(x), str(y), str(x), str(y), str(duration))
        return {"success": True}

    async def cmd_swipe(self, args: dict) -> dict:
        await self._run("shell", "input", "swipe",
                         str(args["x1"]), str(args["y1"]),
                         str(args["x2"]), str(args["y2"]),
                         str(args.get("duration_ms", 300)))
        return {"success": True}

    async def cmd_scroll(self, args: dict) -> dict:
        direction = args.get("direction", "down")
        if not self._screen_size:
            stdout, _, _ = await self._run("shell", "wm", "size")
            m = re.search(r"(\d+)x(\d+)", stdout)
            if m:
                self._screen_size = (int(m.group(1)), int(m.group(2)))
            else:
                self._screen_size = (1080, 1920)

        w, h = self._screen_size
        cx, cy = w // 2, h // 2
        dist = h // 3
        coords = {
            "down": (cx, cy + dist // 2, cx, cy - dist // 2),
            "up": (cx, cy - dist // 2, cx, cy + dist // 2),
            "left": (cx + dist // 2, cy, cx - dist // 2, cy),
            "right": (cx - dist // 2, cy, cx + dist // 2, cy),
        }
        x1, y1, x2, y2 = coords.get(direction, coords["down"])
        await self._run("shell", "input", "swipe", str(x1), str(y1), str(x2), str(y2), "300")
        return {"success": True}

    async def cmd_type_text(self, args: dict) -> dict:
        text = args["text"]
        # Use ADB broadcast for reliable text input
        escaped = text.replace("'", "'\\''")
        await self._run("shell", "input", "text", escaped)
        return {"success": True}

    async def cmd_clear_focused_field(self, args: dict) -> dict:
        # Select all + delete
        await self._run("shell", "input", "keyevent", "KEYCODE_MOVE_HOME")
        await self._run("shell", "input", "keyevent", "--longpress", "KEYCODE_MOVE_END")
        await self._run("shell", "input", "keyevent", "KEYCODE_DEL")
        return {"success": True}

    async def cmd_press_back(self, args: dict) -> dict:
        await self._run("shell", "input", "keyevent", "4")
        return {"success": True}

    async def cmd_press_enter(self, args: dict) -> dict:
        await self._run("shell", "input", "keyevent", "66")
        return {"success": True}

    async def cmd_dismiss_keyboard(self, args: dict) -> dict:
        await self._run("shell", "input", "keyevent", "111")
        return {"success": True}

    async def cmd_check_crash(self, args: dict) -> dict:
        identifier = args["identifier"]
        stdout, _, _ = await self._run("shell", "logcat", "-d", "-t", "50", "-s", "AndroidRuntime:E")
        if identifier in stdout and ("FATAL" in stdout or "Exception" in stdout):
            return {"crash_log": stdout}
        # Also check for ANR dialog
        stdout2, _, _ = await self._run("shell", "dumpsys", "window", "windows")
        if "Application Not Responding" in stdout2 or "AppNotRespondingDialog" in stdout2:
            return {"crash_log": f"ANR detected for {identifier}"}
        return {"crash_log": None}

    async def cmd_check_soft_errors(self, args: dict) -> dict:
        package = args["package_name"]
        logcat = args.get("logcat_output") or ""
        if not logcat:
            logcat, _, _ = await self._run("shell", "logcat", "-d", "-t", "50")
        errors = []
        for line in logcat.splitlines():
            if package in line and ("Exception" in line or "Error" in line):
                if "FATAL" not in line:
                    errors.append(line.strip())
                    if len(errors) >= 3:
                        break
        return {"errors": errors}

    async def cmd_clear_logcat(self, args: dict) -> dict:
        await self._run("logcat", "-c")
        return {"success": True}

    async def cmd_get_logcat(self, args: dict) -> dict:
        lines = args.get("lines", 100)
        stdout, _, _ = await self._run("shell", "logcat", "-d", "-t", str(lines))
        return {"text": stdout}

    async def cmd_launch_app(self, args: dict) -> dict:
        identifier = args["identifier"]
        stdout, stderr, rc = await self._run(
            "shell", "monkey", "-p", identifier, "-c",
            "android.intent.category.LAUNCHER", "1"
        )
        success = rc == 0 and "error" not in stderr.lower()
        return {"success": success}

    async def cmd_force_stop(self, args: dict) -> dict:
        await self._run("shell", "am", "force-stop", args["identifier"])
        return {"success": True}

    async def cmd_install_apk(self, args: dict) -> dict:
        """Install APK that was transferred via file_transfer protocol."""
        tid = args.get("transfer_id", "")
        filename = args.get("filename", "")
        apk_path = self._transfer_dir / filename
        if not apk_path.exists():
            return {"success": False, "message": f"APK file not found: {filename}"}

        stdout, stderr, rc = await self._run("install", "-r", "-t", str(apk_path), timeout=120)
        success = rc == 0 and "Success" in stdout
        message = stdout.strip() if success else (stderr.strip() or stdout.strip())
        return {"success": success, "message": message}

    async def cmd_get_package_name(self, args: dict) -> dict:
        filename = args.get("filename", "")
        apk_path = self._transfer_dir / filename
        if not apk_path.exists():
            return {"package_name": None}
        # Try aapt
        for aapt in ["aapt", "aapt2"]:
            try:
                proc = await asyncio.create_subprocess_exec(
                    aapt, "dump", "badging", str(apk_path),
                    stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
                )
                stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=15)
                text = stdout.decode("utf-8", errors="replace")
                m = re.search(r"package:\s+name='([^']+)'", text)
                if m:
                    return {"package_name": m.group(1)}
            except Exception:
                continue
        return {"package_name": None}

    async def cmd_get_app_version(self, args: dict) -> dict:
        filename = args.get("filename", "")
        apk_path = self._transfer_dir / filename
        if not apk_path.exists():
            return {"version_name": "", "version_code": ""}
        for aapt in ["aapt", "aapt2"]:
            try:
                proc = await asyncio.create_subprocess_exec(
                    aapt, "dump", "badging", str(apk_path),
                    stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
                )
                stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=15)
                text = stdout.decode("utf-8", errors="replace")
                vn = re.search(r"versionName='([^']+)'", text)
                vc = re.search(r"versionCode='([^']+)'", text)
                return {
                    "version_name": vn.group(1) if vn else "",
                    "version_code": vc.group(1) if vc else "",
                }
            except Exception:
                continue
        return {"version_name": "", "version_code": ""}

    async def cmd_get_focused_field_info(self, args: dict) -> dict:
        stdout, _, _ = await self._run("shell", "dumpsys", "input_method")
        info = {}
        for line in stdout.splitlines():
            line = line.strip()
            if "inputType=" in line:
                m = re.search(r"inputType=0x([0-9a-fA-F]+)", line)
                if m:
                    info["inputType"] = m.group(1)
            if "hintText=" in line:
                m = re.search(r"hintText=(.+)", line)
                if m:
                    info["hintText"] = m.group(1).strip()
            if "fieldId=" in line:
                m = re.search(r"fieldId=(\d+)", line)
                if m:
                    info["fieldId"] = m.group(1)
            if "packageName=" in line:
                m = re.search(r"packageName=(\S+)", line)
                if m:
                    info["packageName"] = m.group(1)
        return {"info": info}

    async def cmd_get_recent_toasts(self, args: dict) -> dict:
        stdout, _, _ = await self._run("shell", "dumpsys", "notification")
        toasts = []
        for line in stdout.splitlines():
            if "toast" in line.lower() and "text=" in line.lower():
                toasts.append(line.strip())
        return {"toasts": toasts[-5:]}

    async def cmd_push_file(self, args: dict) -> dict:
        data = base64.b64decode(args["data_b64"])
        device_path = args["device_path"]
        local_tmp = self._transfer_dir / "push_tmp"
        local_tmp.write_bytes(data)
        _, stderr, rc = await self._run("push", str(local_tmp), device_path)
        local_tmp.unlink(missing_ok=True)
        return {"success": rc == 0}

    async def cmd_rotate_screen(self, args: dict) -> dict:
        landscape = args.get("landscape", True)
        await self._run("shell", "settings", "put", "system", "accelerometer_rotation", "0")
        val = "1" if landscape else "0"
        await self._run("shell", "settings", "put", "system", "user_rotation", val)
        return {"success": True}

    async def cmd_enable_auto_rotate(self, args: dict) -> dict:
        await self._run("shell", "settings", "put", "system", "accelerometer_rotation", "1")
        return {"success": True}
