"""
Frida deep-runtime inspection MCP tools — generic, no per-app reverse engineering.

Independent module (E6 / AgentPhone parity). Registered into the existing
FastMCP instance by android_mcp._build_server() via:

    from frida_inspect_tools import register_frida_inspect_tools
    register_frida_inspect_tools(mcp)

This is phoneuse.md's Layer 5 "Deep Runtime" (Frida), scoped deliberately
narrow: only the two capabilities that work on ANY app without first
reverse-engineering it. String-decryption hooking (the third capability
discussed alongside these two) is NOT here — it requires locating a specific
method in a specific app version, which is per-app work (see TASKS.md S3-05/
S3-06: WeChat's obfuscation already defeated that approach once).

  inspect_foreground_intent  — dump the focused Activity's Intent (action/
                                data/component/extras) in a target process.
  dump_live_bitmaps          — enumerate every live Bitmap in a target
                                process's heap; export the largest N as PNG.

Both inject a fixed, audited JS payload (not caller-supplied script) into the
target process via the same Mac-side bridge the existing frida_call MCP tool
uses (frida_bridge.py, POST http://127.0.0.1:8767/frida_call). On-device this
reaches the Mac through `adb reverse tcp:8767 tcp:8767` (see
start_frida_bridge.sh); requires frida-server running as root on the device
and frida_bridge.py running on the Mac.

Known limitation (in frida_bridge.py, not fixed here — out of scope for this
module): its process-name matching does substring search against the
human-readable process label frida reports, not the package name, so
multi-process apps whose main process is labeled differently (e.g. WeChat's
main process shows as "WeChat", not "com.tencent.mm") can resolve to a
sub-process with no UI (e.g. "com.tencent.mm:push") instead of the real
foreground process. Both tools below detect the resulting empty result and
attach a `note` explaining this, instead of silently implying the target app
has no Activities/Bitmaps.
"""

from __future__ import annotations

import json
import urllib.request
from typing import Any

BRIDGE_URL = "http://127.0.0.1:8767/frida_call"

_INTENT_SCRIPT = """
Java.perform(function () {
    var results = [];
    Java.choose('android.app.Activity', {
        onMatch: function (instance) {
            try {
                var info = {};
                info.className = instance.getClass().getName();
                info.hasWindowFocus = instance.hasWindowFocus();
                info.isFinishing = instance.isFinishing();
                var intent = instance.getIntent();
                if (intent !== null) {
                    info.action = intent.getAction();
                    info.dataString = intent.getDataString();
                    var comp = intent.getComponent();
                    info.component = comp !== null ? comp.toString() : null;
                    var extrasObj = {};
                    var bundle = intent.getExtras();
                    if (bundle !== null) {
                        var iterator = bundle.keySet().iterator();
                        while (iterator.hasNext()) {
                            var key = iterator.next();
                            try {
                                var val = bundle.get(key);
                                extrasObj[key] = val !== null ? val.toString() : null;
                            } catch (innerErr) {
                                extrasObj[key] = "<error: " + innerErr + ">";
                            }
                        }
                    }
                    info.extras = extrasObj;
                } else {
                    info.action = null;
                }
                results.push(info);
            } catch (e) {
                results.push({ error: e.toString() });
            }
        },
        onComplete: function () {
            send({ activities: results, count: results.length });
        }
    });
});
"""

_BITMAP_SCRIPT = """
Java.perform(function () {
    var MAX_EXPORT = __MAX_EXPORT__;
    var CompressFormat = Java.use('android.graphics.Bitmap$CompressFormat');
    var ByteArrayOutputStream = Java.use('java.io.ByteArrayOutputStream');
    var Base64 = Java.use('android.util.Base64');
    var BASE64_NO_WRAP = 2;
    var candidates = [];

    Java.choose('android.graphics.Bitmap', {
        onMatch: function (bmp) {
            try {
                if (bmp.isRecycled()) return;
                candidates.push({
                    bmp: bmp,
                    width: bmp.getWidth(),
                    height: bmp.getHeight(),
                    config: bmp.getConfig() !== null ? bmp.getConfig().toString() : null,
                    allocationByteCount: bmp.getAllocationByteCount()
                });
            } catch (e) {}
        },
        onComplete: function () {
            candidates.sort(function (a, b) { return b.allocationByteCount - a.allocationByteCount; });
            var top = candidates.slice(0, MAX_EXPORT);
            var exported = [];
            for (var i = 0; i < top.length; i++) {
                var c = top[i];
                var item = { width: c.width, height: c.height, config: c.config,
                             allocationByteCount: c.allocationByteCount };
                try {
                    var baos = ByteArrayOutputStream.$new();
                    item.compressOk = c.bmp.compress(CompressFormat.PNG.value, 100, baos);
                    var bytes = baos.toByteArray();
                    item.pngByteLength = bytes.length;
                    item.png_base64 = Base64.encodeToString(bytes, BASE64_NO_WRAP);
                } catch (e) {
                    item.exportError = e.toString();
                }
                exported.push(item);
            }
            send({ totalLiveBitmaps: candidates.length, exportedCount: exported.length, bitmaps: exported });
        }
    });
});
"""

_EMPTY_RESULT_NOTE = (
    "Empty result. The target process may be a no-UI helper process rather "
    "than the app's main process — frida_bridge.py matches by substring "
    "against frida's process label (not the package name), which can pick "
    "the wrong process for multi-process apps (e.g. WeChat's main process is "
    "labeled 'WeChat', not 'com.tencent.mm', so the substring match falls "
    "back to a sub-process like 'com.tencent.mm:push'). This does not "
    "necessarily mean the app has no foreground Activity/Bitmaps."
)


def _frida_call(package: str, script: str, timeout: int) -> dict[str, Any]:
    """POST a JS payload to frida_bridge.py (Mac) via the adb-reverse tunnel."""
    payload = json.dumps({"package": package, "script": script, "timeout": timeout}).encode()
    try:
        req = urllib.request.Request(
            BRIDGE_URL, data=payload, headers={"Content-Type": "application/json"}
        )
        with urllib.request.urlopen(req, timeout=timeout + 10) as resp:
            return json.loads(resp.read().decode())
    except Exception as e:
        return {
            "error": f"frida_bridge unreachable: {e}. Ensure start_frida_bridge.sh is "
                     f"running on the Mac and adb reverse tcp:8767 tcp:8767 is set."
        }


def register_frida_inspect_tools(mcp) -> None:
    """Register the generic Frida deep-runtime inspection tools onto an existing FastMCP."""

    @mcp.tool()
    def inspect_foreground_intent(package: str, timeout: int = 10) -> str:
        """
        Inspect the focused Activity's Intent in a running app via Frida.

        Injects a fixed, generic script (no per-app reverse engineering) that
        enumerates every live android.app.Activity instance in the target
        process and returns each one's class name, hasWindowFocus/isFinishing
        state, and Intent (action, data, component, all Bundle extras
        stringified). The entry with hasWindowFocus=true is the true
        foreground Activity; others are backgrounded-but-alive.

        Requires root + frida-server on device, and frida_bridge.py running
        on the Mac (start_frida_bridge.sh) reachable via adb reverse
        tcp:8767 tcp:8767.

        Note: complex (Parcelable/Serializable) extras only show Java's
        default toString() (e.g. "SomeClass@a1b2c3"), not a deep dump.
        Multi-process apps may attach to the wrong sub-process — see this
        module's docstring.

        Args:
            package: Target app package name, e.g. 'com.android.settings'.
            timeout: Max seconds to wait for script output (default 10).
        """
        result = _frida_call(package, _INTENT_SCRIPT, timeout)
        if "error" in result:
            return json.dumps(result, ensure_ascii=False)
        output = result.get("output") or []
        activities = output[0].get("activities", []) if output else []
        if not activities:
            result["note"] = _EMPTY_RESULT_NOTE
        return json.dumps(result, ensure_ascii=False)

    @mcp.tool()
    def dump_live_bitmaps(package: str, max_export: int = 3, timeout: int = 15) -> str:
        """
        Enumerate live (non-recycled) Bitmap objects in a running app's heap
        via Frida, and export the largest few as base64-encoded PNGs.

        Injects a fixed, generic script (no per-app reverse engineering) that
        walks the target process's heap for every live android.graphics.Bitmap,
        collects width/height/config/allocationByteCount for all of them, then
        PNG-compresses and base64-encodes only the `max_export` largest ones
        (compression happens inside the target app's own process — keep this
        small; it is slow and memory-heavy for large bitmaps).

        Requires root + frida-server on device, and frida_bridge.py running
        on the Mac (start_frida_bridge.sh) reachable via adb reverse
        tcp:8767 tcp:8767.

        Note: this only captures whatever happens to be allocated in memory
        at the moment of the call — not a guarantee of "the photo you
        wanted," just a snapshot. Multi-process apps may attach to the wrong
        sub-process — see this module's docstring.

        Args:
            package: Target app package name, e.g. 'com.android.settings'.
            max_export: How many of the largest live bitmaps to export as PNG
                        (default 3, keep small — see note above).
            timeout: Max seconds to wait for script output (default 15).
        """
        script = _BITMAP_SCRIPT.replace("__MAX_EXPORT__", str(int(max_export)))
        result = _frida_call(package, script, timeout)
        if "error" in result:
            return json.dumps(result, ensure_ascii=False)
        output = result.get("output") or []
        total = output[0].get("totalLiveBitmaps", 0) if output else 0
        if not total:
            result["note"] = _EMPTY_RESULT_NOTE
        return json.dumps(result, ensure_ascii=False)