"""
Frida Manager — 管理 Frida 与微信进程的连接生命周期

职责：
- 连接设备（USB / remote / gadget 三种模式）
- attach 到微信进程
- 加载 JS 脚本
- 维护 rpc.exports 代理
- 自动重连

连接模式：
- usb:    需要 Root，通过 frida-server 连接（传统方式）
- gadget: 无需 Root，通过注入到 APK 中的 frida-gadget 连接
- remote: 通过 TCP 连接远程 frida-server

技术栈：Frida 16.5.6（17.x 在 Android 11 上 Java 桥失效）
"""

import os
import re
import time
import logging
import threading
import subprocess
from typing import Optional, Callable, Dict, Any, Literal

logger = logging.getLogger(__name__)

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_SCRIPT = os.path.join(SCRIPT_DIR, "wechat_hook_v2.js")
WECHAT_PACKAGE = "com.tencent.mm"
import random as _random
_ANTI_DETECT_PORT_RANGE = (10000, 60000)
GADGET_PORT = _random.randint(*_ANTI_DETECT_PORT_RANGE)

ConnectionMode = Literal["usb", "gadget", "remote"]

_SNAKE_RE = re.compile(r'(?<=[a-z0-9])([A-Z])')


def _random_port() -> int:
    return _random.randint(*_ANTI_DETECT_PORT_RANGE)


class FridaManager:
    """Frida 会话管理器，支持 USB / Gadget / Remote 三种连接模式"""

    def __init__(
        self,
        device_serial: Optional[str] = None,
        script_path: str = DEFAULT_SCRIPT,
        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
        on_detach: Optional[Callable[[str], None]] = None,
        auto_reconnect: bool = True,
        reconnect_interval: float = 5.0,
        mode: ConnectionMode = "gadget",
        gadget_host: str = "127.0.0.1",
        gadget_port: int = 0,
    ):
        if gadget_port == 0:
            gadget_port = _random_port()
        self.device_serial = device_serial
        self.script_path = script_path
        self.on_event = on_event
        self.on_detach = on_detach
        self.auto_reconnect = auto_reconnect
        self.reconnect_interval = reconnect_interval
        self.mode = mode
        self.gadget_host = gadget_host
        self.gadget_port = gadget_port

        self._device = None
        self._session = None
        self._script = None
        self._rpc = None
        self._running = False
        self._detaching = False
        self._lock = threading.Lock()

    @property
    def connected(self) -> bool:
        return self._session is not None and self._script is not None

    @property
    def rpc(self):
        return self._rpc

    def start(self) -> bool:
        try:
            import frida
        except ImportError:
            logger.error("frida 未安装，请执行: pip install frida-tools")
            return False

        with self._lock:
            if self.connected:
                return True
            return self._do_attach()

    def _setup_gadget_forward(self):
        """为 Gadget 模式设置 ADB 端口转发"""
        try:
            cmd = ["adb"]
            if self.device_serial:
                cmd += ["-s", self.device_serial]
            cmd += ["forward", f"tcp:{self.gadget_port}", f"tcp:{self.gadget_port}"]
            subprocess.run(cmd, capture_output=True, timeout=5)
            logger.info(f"ADB 端口转发已设置: tcp:{self.gadget_port}")
        except Exception as e:
            logger.warning(f"端口转发设置失败: {e}")

    def _do_attach(self) -> bool:
        import frida

        try:
            if self.mode == "gadget":
                self._setup_gadget_forward()
                addr = f"{self.gadget_host}:{self.gadget_port}"
                mgr = frida.get_device_manager()
                self._device = mgr.add_remote_device(addr)
                logger.info(f"Frida Gadget 已连接: {addr}")

                self._session = self._device.attach("Gadget")
                self._session.on("detached", self._on_session_detached)
                logger.info("已 attach 到 Gadget 进程（无 Root 模式）")

            elif self.mode == "remote":
                self._setup_gadget_forward()
                addr = f"{self.gadget_host}:{self.gadget_port}"
                mgr = frida.get_device_manager()
                self._device = mgr.add_remote_device(addr)
                logger.info(f"Frida Remote 已连接: {addr}")

                pid = self._find_wechat_pid()
                if pid is None and self.device_serial:
                    logger.warning("微信未运行，尝试 adb 启动后 attach")
                    try:
                        subprocess.run(
                            ["adb", "-s", self.device_serial, "shell", "am", "start", "-n",
                             f"{WECHAT_PACKAGE}/.ui.LauncherUI"],
                            capture_output=True, timeout=15,
                        )
                        time.sleep(5)
                        pid = self._find_wechat_pid()
                    except Exception as launch_err:
                        logger.warning(f"adb 启动微信失败: {launch_err}")
                if pid is None:
                    logger.warning("微信未运行，Remote 模式尝试 spawn 启动")
                    try:
                        pid = self._device.spawn([WECHAT_PACKAGE])
                        self._device.resume(pid)
                        time.sleep(3)
                    except Exception as spawn_err:
                        logger.error(f"微信 spawn 失败: {spawn_err}")
                        return False
                self._session = self._device.attach(pid)
                self._session.on("detached", self._on_session_detached)
                logger.info(f"已 attach 到微信进程 pid={pid}")

            else:  # usb (需要 Root)
                if self.device_serial:
                    self._device = frida.get_device(self.device_serial)
                else:
                    self._device = frida.get_usb_device(timeout=10)
                logger.info(f"Frida USB 已连接设备: {self._device.name}")

                pid = self._find_wechat_pid()
                if pid is None:
                    logger.warning("微信未运行，尝试 spawn 启动")
                    pid = self._device.spawn([WECHAT_PACKAGE])
                    self._device.resume(pid)
                    time.sleep(3)
                self._session = self._device.attach(pid)
                self._session.on("detached", self._on_session_detached)
                logger.info(f"已 attach 到微信进程 pid={pid}")

            self._load_script()
            self._running = True
            return True

        except Exception as e:
            logger.error(f"Frida attach 失败 [{self.mode}]: {e}")
            self._do_cleanup()
            return False

    def _find_wechat_pid(self) -> Optional[int]:
        # 优先 ADB pidof（Remote/Gadget 经端口转发时 enumerate 常漏进程）
        if self.device_serial:
            for attempt in range(3):
                try:
                    result = subprocess.run(
                        ["adb", "-s", self.device_serial, "shell", "pidof", WECHAT_PACKAGE],
                        capture_output=True, text=True, timeout=8,
                    )
                    out = (result.stdout or "").strip()
                    if out:
                        return int(out.split()[0])
                except Exception:
                    pass
                if attempt < 2:
                    time.sleep(2)
        try:
            for proc in self._device.enumerate_processes():
                if proc.name in (WECHAT_PACKAGE, "WeChat") or getattr(proc, 'identifier', '') == WECHAT_PACKAGE:
                    return proc.pid
        except Exception:
            pass
        try:
            for app in self._device.enumerate_applications():
                if app.identifier == WECHAT_PACKAGE and app.pid > 0:
                    return app.pid
        except Exception:
            pass
        try:
            result = subprocess.run(
                ["adb", "-s", self.device_serial or "", "shell", "pidof", WECHAT_PACKAGE],
                capture_output=True, text=True, timeout=5,
            )
            if result.stdout.strip():
                return int(result.stdout.strip().split()[0])
        except Exception:
            pass
        return None

    def _load_script(self):
        path = self.script_path
        if not os.path.exists(path):
            fallback = os.path.join(SCRIPT_DIR, "wechat_hook_v1.js")
            if os.path.exists(fallback):
                path = fallback
            else:
                raise FileNotFoundError(f"Hook 脚本不存在: {path}")

        with open(path, "r", encoding="utf-8") as f:
            source = f.read()

        self._script = self._session.create_script(source)
        self._script.on("message", self._on_message)
        self._script.load()
        self._rpc = self._script.exports_sync
        logger.info(f"Hook 脚本已加载: {os.path.basename(path)}")

        try:
            pong = self._rpc.ping()
            logger.info(f"RPC ping: {pong}")
        except Exception as e:
            logger.warning(f"RPC ping 失败: {e}")

    def _on_message(self, message: dict, data):
        if message.get("type") == "send":
            payload = message.get("payload", {})
            msg_type = payload.get("type", "")

            if msg_type == "hook_event" and self.on_event:
                self.on_event(payload)
            elif msg_type == "log":
                level = payload.get("level", "info")
                tag = payload.get("tag", "hook")
                text = payload.get("message", "")
                getattr(logger, level, logger.info)(f"[{tag}] {text}")
        elif message.get("type") == "error":
            logger.error(f"Frida 脚本错误: {message.get('description', '')}")

    def _on_session_detached(self, reason: str):
        """Frida 内部线程回调 — 只清理引用，不重复 detach/unload，不获取 _lock（避免死锁）"""
        logger.warning(f"Frida 会话断开: {reason}")
        self._detaching = True
        self._nullify_refs()
        self._detaching = False

        if self.on_detach:
            try:
                self.on_detach(reason)
            except Exception as e:
                logger.warning(f"on_detach 回调异常: {e}")

        if self.auto_reconnect and self._running and reason != "application-requested":
            threading.Thread(target=self._reconnect_loop, daemon=True).start()

    def _reconnect_loop(self):
        while self._running and not self.connected:
            logger.info(f"尝试重连 Frida（{self.reconnect_interval}s 后）...")
            time.sleep(self.reconnect_interval)
            with self._lock:
                if self.connected:
                    break
                try:
                    if self._do_attach():
                        logger.info("Frida 重连成功")
                        break
                except Exception as e:
                    logger.warning(f"重连失败: {e}")

    @staticmethod
    def _to_snake(name: str) -> str:
        return _SNAKE_RE.sub(r'_\1', name).lower()

    def call_rpc(self, method: str, params: dict = None) -> Dict[str, Any]:
        with self._lock:
            if not self.connected or not self._rpc:
                return {"success": False, "error": "Frida 未连接"}
            rpc = self._rpc
        try:
            snake = self._to_snake(method)
            fn = getattr(rpc, snake, None) or getattr(rpc, method, None)
            if fn is None:
                return {"success": False, "error": f"RPC 方法不存在: {method} ({snake})"}
            result = fn(params) if params else fn({})
            if isinstance(result, dict):
                result.setdefault("success", True)
                return result
            return {"success": True, "data": result}
        except Exception as e:
            logger.error(f"RPC 调用失败 [{method}]: {e}")
            return {"success": False, "error": str(e)}

    def stop(self):
        self._running = False
        with self._lock:
            self._do_cleanup()
        logger.info("FridaManager 已停止")

    def _nullify_refs(self):
        """清空引用（不执行 unload/detach），用于 detach 回调内"""
        self._script = None
        self._session = None
        self._rpc = None

    def _do_cleanup(self):
        """主动断开时完整清理（unload 脚本 + detach 会话）"""
        if self._detaching:
            return
        try:
            if self._script:
                self._script.unload()
        except Exception:
            pass
        try:
            if self._session:
                self._session.detach()
        except Exception:
            pass
        self._nullify_refs()

    def get_status(self) -> Dict[str, Any]:
        return {
            "connected": self.connected,
            "mode": self.mode,
            "device": self._device.name if self._device else None,
            "device_serial": self.device_serial,
            "script": os.path.basename(self.script_path),
            "running": self._running,
            "gadget_endpoint": f"{self.gadget_host}:{self.gadget_port}" if self.mode == "gadget" else None,
        }