"""
Frida 无线管理路由 — 服务端设备管理与指令分发
=============================================

功能：
1. 设备注册/注销（WiFi 直连）
2. Frida 连接管理（连接/断开/重连）
3. 指令分发（通过 WebSocket 发送到 Agent 的 WirelessBridge）
4. 部署脚本生成（Root/免Root/自动）
5. 局域网设备发现
6. 健康状态监控

所有操作通过 WiFi 完成，无需 USB 连接。
"""
from __future__ import annotations
import asyncio
import logging
import os
import sys
import time
from typing import Optional, List, Dict, Any
from datetime import datetime

from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel, Field

router = APIRouter(prefix="/frida", tags=["Frida无线管理"])
logger = logging.getLogger(__name__)


def _ensure_sdk_agent_path() -> None:
    """device 端 Hook 模块在 sdk/agent，非 sdk/app/agent。"""
    agent_dir = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "..", "agent"))
    if os.path.isdir(agent_dir) and agent_dir not in sys.path:
        sys.path.insert(0, agent_dir)


# ============================================================
# § 1  数据模型
# ============================================================

class DeviceRegisterRequest(BaseModel):
    """设备注册请求"""
    device_id: str
    ip: str
    frida_port: int = 27042
    mode: str = "auto"  # root_server / gadget / auto
    android_version: Optional[str] = None
    device_model: Optional[str] = None
    frida_version: Optional[str] = None
    wechat_version: Optional[str] = None


class DeviceConnectRequest(BaseModel):
    """设备连接请求"""
    device_id: str
    force: bool = False  # 强制重连


class CommandRequest(BaseModel):
    """指令执行请求"""
    device_id: str
    action: str
    params: dict = Field(default_factory=dict)
    timeout: int = 30


class BatchCommandRequest(BaseModel):
    """批量指令请求"""
    device_id: str
    commands: List[dict]
    timeout: int = 60


class DeployScriptRequest(BaseModel):
    """部署脚本生成请求"""
    mode: str = "auto"  # root / gadget / auto
    server_url: str = ""
    port: int = 0
    arch: str = "arm64"


class DiscoverRequest(BaseModel):
    """设备发现请求"""
    subnet: str = ""
    port: int = 27042
    timeout: float = 2.0


# ============================================================
# § 2  内存设备池（生产环境应持久化到 MongoDB）
# ============================================================

class FridaDevicePool:
    """Frida 设备连接池"""

    def __init__(self):
        self.devices: Dict[str, Dict[str, Any]] = {}

    def register(self, data: dict) -> dict:
        device_id = data["device_id"]
        now = datetime.now().isoformat()
        self.devices[device_id] = {
            **data,
            "status": "registered",
            "registered_at": now,
            "last_seen": now,
            "wechat_attached": False,
            "hook_ready": False,
            "supported_actions": 0,
        }
        logger.info(f"设备注册: {device_id} ({data.get('ip')}:{data.get('frida_port')})")
        return self.devices[device_id]

    def unregister(self, device_id: str):
        if device_id in self.devices:
            del self.devices[device_id]

    def get(self, device_id: str) -> Optional[dict]:
        return self.devices.get(device_id)

    def update(self, device_id: str, **kwargs):
        if device_id in self.devices:
            self.devices[device_id].update(kwargs)
            self.devices[device_id]["last_seen"] = datetime.now().isoformat()

    def list_all(self) -> List[dict]:
        return list(self.devices.values())

    def list_online(self) -> List[dict]:
        return [d for d in self.devices.values() if d.get("status") in ("connected", "registered")]


# 全局设备池
device_pool = FridaDevicePool()


# ============================================================
# § 3  辅助函数 — 通过 WebSocket Hub 发送指令到 Agent
# ============================================================

async def _send_command_to_agent(device_id: str, command: dict, timeout: int = 30) -> dict:
    """
    通过 WebSocket Hub 将指令发送到 Agent 端的 WirelessBridge
    
    流程：
    1. 查找设备的 WebSocket 连接
    2. 发送 command 消息
    3. 等待 Agent 返回结果
    """
    try:
        from services.ws_hub import ws_hub

        if not ws_hub.is_online(device_id):
            return {"success": False, "error": f"设备离线: {device_id}"}

        command["type"] = "command"

        # 统一走 ws_hub.send_command 的 pending/Future 机制。
        result = await ws_hub.send_command(device_id, command, timeout=timeout)
        if result.get("success") is not None:
            return result

        data = result.get("data") or {}
        if isinstance(data, dict) and data.get("success") is not None:
            return data

        code = int(result.get("code") or 500)
        return {
            "success": code == 200,
            "code": code,
            "message": result.get("message", ""),
            "data": data,
        }

    except asyncio.TimeoutError:
        return {"success": False, "error": f"指令超时 ({timeout}s)", "command_id": command.get("command_id")}
    except Exception as e:
        return {"success": False, "error": str(e)}


async def _wait_for_result(device_id: str, command_id: str) -> dict:
    """等待指令结果"""
    from services.ws_hub import ws_hub

    # 轮询等待结果
    for _ in range(600):  # 最多等 60 秒
        result = ws_hub.command_results.get(command_id)
        if result is not None:
            del ws_hub.command_results[command_id]
            return result
        await asyncio.sleep(0.1)

    return {"success": False, "error": "等待结果超时"}


# ============================================================
# § 4  API 端点
# ============================================================

# ---- 设备管理 ----

@router.post("/register", response_model=dict)
async def register_device(req: DeviceRegisterRequest):
    """
    注册 Frida 设备（手机端部署脚本自动调用）
    
    设备通过 WiFi 连接后，自动向服务器注册。
    """
    device = device_pool.register(req.dict())
    return {
        "success": True,
        "device_id": req.device_id,
        "message": f"设备已注册: {req.ip}:{req.frida_port} ({req.mode})",
        "device": device,
    }


@router.delete("/unregister/{device_id}", response_model=dict)
async def unregister_device(device_id: str):
    """注销设备"""
    device_pool.unregister(device_id)
    return {"success": True, "device_id": device_id, "message": "设备已注销"}


@router.get("/devices", response_model=dict)
async def list_devices():
    """列出所有已注册的 Frida 设备"""
    devices = device_pool.list_all()
    return {
        "success": True,
        "total": len(devices),
        "online": len(device_pool.list_online()),
        "devices": devices,
    }


@router.get("/device/{device_id}", response_model=dict)
async def get_device(device_id: str):
    """获取设备详情"""
    device = device_pool.get(device_id)
    if not device:
        raise HTTPException(status_code=404, detail=f"设备不存在: {device_id}")
    return {"success": True, "device": device}


# ---- 连接管理 ----

@router.post("/connect", response_model=dict)
async def connect_device(req: DeviceConnectRequest):
    """
    连接设备的 Frida（通过 WiFi TCP）
    
    向 Agent 发送连接指令，Agent 通过 WirelessBridge 连接本地 Frida。
    """
    device = device_pool.get(req.device_id)
    if not device:
        raise HTTPException(status_code=404, detail=f"设备未注册: {req.device_id}")

    result = await _send_command_to_agent(req.device_id, {
        "action": "frida_connect",
        "params": {
            "ip": device.get("ip", "127.0.0.1"),
            "port": device.get("frida_port", 27042),
            "mode": device.get("mode", "remote"),
            "force": req.force,
        },
    })

    if result.get("success"):
        device_pool.update(req.device_id,
                           status="connected",
                           wechat_attached=result.get("wechat_attached", False),
                           hook_ready=True,
                           supported_actions=result.get("supported_actions", 0))

    return result


@router.post("/disconnect/{device_id}", response_model=dict)
async def disconnect_device(device_id: str):
    """断开设备 Frida 连接"""
    result = await _send_command_to_agent(device_id, {
        "action": "frida_disconnect",
    })
    device_pool.update(device_id, status="disconnected", wechat_attached=False, hook_ready=False)
    return result


# ---- 指令执行 ----

@router.post("/execute", response_model=dict)
async def execute_command(req: CommandRequest):
    """
    执行 Frida RPC 指令
    
    通过 WebSocket 将指令发送到 Agent，Agent 通过 WirelessBridge 执行 Frida RPC。
    支持 112 个微信操作（24 个模块）。
    """
    device = device_pool.get(req.device_id)
    if not device:
        raise HTTPException(status_code=404, detail=f"设备未注册: {req.device_id}")

    result = await _send_command_to_agent(req.device_id, {
        "action": req.action,
        "params": req.params,
    }, timeout=req.timeout)

    return result


@router.post("/execute/batch", response_model=dict)
async def execute_batch(req: BatchCommandRequest):
    """批量执行指令"""
    device = device_pool.get(req.device_id)
    if not device:
        raise HTTPException(status_code=404, detail=f"设备未注册: {req.device_id}")

    result = await _send_command_to_agent(req.device_id, {
        "type": "batch",
        "commands": req.commands,
    }, timeout=req.timeout)

    return result


# ---- 部署脚本 ----

@router.post("/deploy/script", response_model=dict)
async def generate_deploy_script(req: DeployScriptRequest):
    """
    生成 Frida 部署脚本
    
    返回可在 Termux 中直接执行的 shell 脚本。
    支持 Root（frida-server）和免Root（frida-gadget）两种模式。
    """
    try:
        # 动态导入避免循环依赖
        import sys
        sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'agent', 'hook'))
        _ensure_sdk_agent_path()
        from hook.wireless_deployer import WirelessDeployer, DeployConfig

        config = DeployConfig(frida_arch=req.arch)
        deployer = WirelessDeployer(config)

        if req.mode == "root":
            script = deployer.generate_root_deploy_script(port=req.port)
        elif req.mode == "gadget":
            script = deployer.generate_gadget_deploy_script(port=req.port)
        else:
            script = deployer.generate_auto_deploy_script(
                server_url=req.server_url,
                port=req.port,
            )

        return {
            "success": True,
            "mode": req.mode,
            "script": script,
            "instructions": [
                "1. 在手机上打开 Termux",
                "2. 复制并粘贴脚本内容",
                "3. 执行脚本",
                f"4. 脚本会自动{'启动 frida-server' if req.mode == 'root' else '注入 frida-gadget' if req.mode == 'gadget' else '检测 Root 并选择最佳方案'}",
                "5. 完成后设备会自动注册到服务器",
            ],
        }
    except Exception as e:
        return {"success": False, "error": str(e)}


@router.get("/deploy/install-script", response_model=dict)
async def get_install_command(server_url: str = ""):
    """
    获取一键安装命令
    
    返回可在 Termux 中直接执行的 curl 命令。
    """
    if not server_url:
        server_url = "ws://YOUR_SERVER_IP:8899/ws/device"

    api_base = server_url.replace("ws://", "http://").replace("wss://", "https://").split("/ws/")[0]

    return {
        "success": True,
        "command": f'curl -sL {api_base}/install.sh | bash -s -- --server {server_url} --frida',
        "manual_steps": [
            "1. 在手机上安装 Termux（从 F-Droid 下载）",
            "2. 打开 Termux",
            f"3. 执行: curl -sL {api_base}/install.sh | bash -s -- --server {server_url} --frida",
            "4. 等待安装完成",
            "5. Agent 会自动连接服务器并启动 Frida",
        ],
    }


# ---- 设备发现 ----

@router.post("/discover", response_model=dict)
async def discover_devices(req: DiscoverRequest):
    """
    扫描局域网内的 Frida 设备
    
    扫描指定子网内运行 frida-server 的设备。
    """
    try:
        _ensure_sdk_agent_path()
        from hook.wireless_deployer import wireless_deployer

        loop = asyncio.get_event_loop()
        devices = await loop.run_in_executor(
            None,
            wireless_deployer.discover_devices,
            req.subnet,
            req.port,
            req.timeout,
        )

        return {
            "success": True,
            "found": len(devices),
            "devices": devices,
        }
    except Exception as e:
        return {"success": False, "error": str(e), "found": 0, "devices": []}


# ---- 状态监控 ----

@router.get("/status", response_model=dict)
async def frida_status():
    """获取 Frida 无线管理总状态"""
    devices = device_pool.list_all()
    connected = [d for d in devices if d.get("status") == "connected"]
    hook_ready = [d for d in connected if d.get("hook_ready")]

    return {
        "success": True,
        "summary": {
            "total_devices": len(devices),
            "connected": len(connected),
            "hook_ready": len(hook_ready),
            "disconnected": len(devices) - len(connected),
        },
        "supported_actions": 112,
        "supported_modules": 24,
        "connection_mode": "WiFi TCP (无USB)",
        "frida_version": "16.5.6",
    }


@router.get("/actions", response_model=dict)
async def list_supported_actions():
    """列出所有支持的 Frida RPC 操作"""
    try:
        _ensure_sdk_agent_path()
        from hook.hook_executor import ACTION_TO_RPC, MODULE_NAMES

        # 按模块分组
        modules = {}
        for module_id, module_name in MODULE_NAMES.items():
            modules[module_id] = {
                "name": module_name,
                "actions": [],
            }

        # 简单分组（基于 ACTION_TO_RPC 的注释分组）
        module_action_map = {
            "H15": ["get_messages", "get_recent_messages", "search_messages"],
            "H16": ["get_contacts", "get_contact_info", "search_contacts"],
            "H17": ["send_message", "send_group_message"],
            "H18": ["get_friend_requests"],
            "H19": ["add_friend", "accept_friend", "delete_friend", "set_friend_remark", "add_friend_by_qr"],
            "H20": ["post_moments", "delete_moments"],
            "H21": ["get_moments", "like_moments", "comment_moments"],
            "H22": ["get_groups", "get_group_info", "get_group_members", "create_group",
                     "invite_to_group", "remove_from_group", "set_group_announcement",
                     "set_group_name", "quit_group"],
            "H23": ["get_profile", "check_account_status", "set_nickname", "set_signature",
                     "set_avatar", "set_sex", "set_region", "set_what_up"],
            "H24": ["unblock_self", "change_password", "bind_phone", "unbind_phone",
                     "get_login_devices", "remove_login_device", "enable_fingerprint",
                     "set_account_protection"],
            "H25": ["send_red_packet", "receive_red_packet", "send_transfer",
                     "receive_transfer", "get_wallet_balance", "get_transaction_history"],
            "H26": ["scan_qr_code", "generate_my_qr_code", "generate_group_qr_code", "add_friend_by_qr"],
            "H27": ["browse_channels", "like_channel_video", "comment_channel_video",
                     "follow_channel", "unfollow_channel", "share_channel_video"],
            "H28": ["get_labels", "create_label", "delete_label", "set_contact_label",
                     "get_contacts_by_label"],
            "H29": ["get_favorites", "add_favorite", "delete_favorite"],
            "H30": ["set_privacy", "set_notification", "clear_chat_history",
                     "set_chat_background", "set_do_not_disturb", "pin_chat"],
            "H31": ["global_search"],
            "H32": ["open_mini_program", "get_recent_mini_programs", "share_mini_program"],
            "H33": ["send_image", "send_video", "send_file", "send_voice",
                     "send_location", "send_card", "send_link"],
            "H34": ["forward_message", "forward_multiple", "revoke_message"],
            "H35": ["register_account", "login_by_password", "login_by_sms",
                     "logout", "switch_account", "auto_register", "check_login_state", "get_sim_phone"],
            "H36": ["get_official_accounts", "follow_official_account",
                     "unfollow_official_account", "get_official_account_articles"],
            "H37": ["send_emoji", "add_custom_emoji"],
            "H38": ["add_to_float", "remove_from_float"],
            "H39": ["get_device_info", "get_storage_info", "get_network_info"],
        }

        for module_id, actions in module_action_map.items():
            if module_id in modules:
                modules[module_id]["actions"] = actions

        return {
            "success": True,
            "total_actions": len(ACTION_TO_RPC),
            "total_modules": len(MODULE_NAMES),
            "modules": modules,
            "all_actions": sorted(ACTION_TO_RPC.keys()),
        }
    except Exception as e:
        return {"success": False, "error": str(e)}


# 需要 os 模块
import os
