from __future__ import annotations

import argparse
import json
import subprocess
import time
import urllib.parse
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import sys

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from core.h5_sig3 import parse_h5_sig3


APP_ID = "com.kuaishou.nebula"
H5_PRODUCT = "h5_jsbridge"
H5_SDK = "5bbcf3cd-727b-48ab-b4b4-5f01e61ee9a5"
DEFAULT_HAR = Path("nebula.kuaishou.com_2026_07_10_17_15_37.har")
DEFAULT_ADB = Path(r"D:\Tools\platform-tools\adb.exe")
DEFAULT_FRIDA = Path(r"C:\Users\youfak\.vfox\sdks\python\Scripts\frida.exe")


@dataclass(frozen=True)
class H5Sample:
    index: int
    method: str
    path: str
    target_crc32: int
    query_pairs: tuple[tuple[str, str], ...]
    body: str

    @property
    def key(self) -> str:
        return f"{self.index}:{self.method}:{self.path}"


def _load_h5_samples(har_path: Path) -> list[H5Sample]:
    with har_path.open("r", encoding="utf-8", errors="ignore") as f:
        har = json.load(f)

    samples: list[H5Sample] = []
    for i, entry in enumerate(har.get("log", {}).get("entries", [])):
        req = entry.get("request") or {}
        url = req.get("url") or ""
        split = urllib.parse.urlsplit(url)
        pairs = tuple(urllib.parse.parse_qsl(split.query, keep_blank_values=True))
        qs = dict(pairs)
        sig3 = qs.get("__NS_sig3", "")
        if len(sig3) != 68:
            continue
        if "unionTask" not in split.path and "overview" not in split.path:
            continue
        post = req.get("postData") or {}
        samples.append(
            H5Sample(
                index=i,
                method=req.get("method", "GET"),
                path=split.path,
                target_crc32=parse_h5_sig3(sig3).crc32,
                query_pairs=pairs,
                body=post.get("text") or "",
            )
        )
    return samples


def _without_sig3(pairs: tuple[tuple[str, str], ...], keep_sigcat: bool) -> list[tuple[str, str]]:
    out: list[tuple[str, str]] = []
    for key, value in pairs:
        if key == "__NS_sig3":
            continue
        if key == "sigCatVer" and not keep_sigcat:
            continue
        out.append((key, value))
    return out


def _urlenc(pairs: list[tuple[str, str]]) -> str:
    return urllib.parse.urlencode(pairs)


def _sorted_urlenc(pairs: list[tuple[str, str]]) -> str:
    return urllib.parse.urlencode(sorted(pairs))


def _sorted_no_sep(pairs: list[tuple[str, str]]) -> str:
    return "".join(f"{k}={v}" for k, v in sorted(pairs))


def _compact_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def _json_candidates_from_pairs(pairs: list[tuple[str, str]]) -> list[str]:
    if not pairs:
        return []
    obj = {k: v for k, v in pairs}
    out = [_compact_json(obj)]
    for key, value in pairs:
        if not value:
            continue
        if value[:1] not in "{[":
            continue
        try:
            out.append(value)
            out.append(_compact_json(json.loads(value)))
        except Exception:
            pass
    return out


def _body_json_candidates(body: str) -> list[str]:
    if not body:
        return []
    out = [body]
    try:
        parsed = json.loads(body)
        compact = _compact_json(parsed)
        out.append(compact)
        if isinstance(parsed, dict):
            for value in parsed.values():
                if isinstance(value, str) and value[:1] in "{[":
                    try:
                        out.append(value)
                        out.append(_compact_json(json.loads(value)))
                    except Exception:
                        pass
    except Exception:
        pass
    return out


def _candidate_map(samples: list[H5Sample]) -> dict[str, str]:
    candidates: dict[str, str] = {}

    def add(sample: H5Sample, name: str, plain: str) -> None:
        key = f"{sample.index}:{name}"
        candidates[key] = plain

    for sample in samples:
        add(sample, "empty", "")
        add(sample, "path", sample.path)
        add(sample, "method_path", f"{sample.method} {sample.path}")

        for keep_sigcat in (True, False):
            tag = "keep_sigcat" if keep_sigcat else "no_sigcat"
            pairs = _without_sig3(sample.query_pairs, keep_sigcat)
            raw_q = _urlenc(pairs)
            sorted_q = _sorted_urlenc(pairs)
            no_sep = _sorted_no_sep(pairs)
            for name, plain in (
                (f"{tag}_raw_q", raw_q),
                (f"{tag}_sorted_q", sorted_q),
                (f"{tag}_sorted_no_sep", no_sep),
                (f"{tag}_path_raw_q", sample.path + ("?" + raw_q if raw_q else "")),
                (f"{tag}_path_sorted_q", sample.path + ("?" + sorted_q if sorted_q else "")),
                (f"{tag}_path_no_sep", sample.path + no_sep),
            ):
                add(sample, name, plain)
            for idx, plain in enumerate(_json_candidates_from_pairs(pairs)):
                add(sample, f"{tag}_json_{idx}", plain)
            if sample.body:
                add(sample, f"{tag}_raw_q_body", (raw_q + "&" if raw_q else "") + sample.body)
                add(sample, f"{tag}_path_raw_q_body", sample.path + ("?" + raw_q if raw_q else "") + sample.body)

        for idx, plain in enumerate(_body_json_candidates(sample.body)):
            add(sample, f"body_{idx}", plain)
            add(sample, f"path_body_{idx}", sample.path + plain)

    return candidates


def _json_script(candidates: dict[str, str]) -> str:
    payload = [{"id": key, "plain": value} for key, value in candidates.items()]
    return f"""
"use strict";

const PRODUCT = {json.dumps(H5_PRODUCT)};
const SDK = {json.dumps(H5_SDK)};
const CANDIDATES = {json.dumps(payload, ensure_ascii=False)};

function log(tag, obj) {{
  try {{
    console.log(tag + " " + JSON.stringify(obj));
  }} catch (e) {{
    console.log(tag + " {{\\"error\\":\\"" + e + "\\"}}");
  }}
}}

Java.perform(function () {{
  try {{
    const MXSec = Java.use("com.middleware.security.MXSec");
    const wrapper = MXSec.get().getWrapper();
    log("[H5_ORACLE_START]", {{ count: CANDIDATES.length }});
    for (let i = 0; i < CANDIDATES.length; i++) {{
      const item = CANDIDATES[i];
      try {{
        const ret = String(wrapper.atlasSign(PRODUCT, SDK, 0, item.plain));
        log("[H5_ORACLE]", {{ id: item.id, ret: ret, plain: item.plain }});
      }} catch (e) {{
        log("[H5_ORACLE_ERR]", {{ id: item.id, error: String(e), plain: item.plain }});
      }}
    }}
    log("[H5_ORACLE_DONE]", {{ count: CANDIDATES.length }});
  }} catch (e) {{
    log("[H5_ORACLE_FATAL]", {{ error: String(e) }});
  }}
}});
"""


def _run(cmd: list[str], timeout: int = 15) -> str:
    return subprocess.check_output(cmd, text=True, stderr=subprocess.STDOUT, timeout=timeout).strip()


def _ensure_app(adb: Path) -> int:
    try:
        out = _run([str(adb), "shell", "pidof", APP_ID], timeout=5)
        if out:
            return int(out.split()[0])
    except Exception:
        pass
    subprocess.run(
        [str(adb), "shell", "monkey", "-p", APP_ID, "-c", "android.intent.category.LAUNCHER", "1"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        check=False,
    )
    time.sleep(8)
    out = _run([str(adb), "shell", "pidof", APP_ID], timeout=5)
    return int(out.split()[0])


def _run_frida(frida: Path, pid: int, script_path: Path, log_path: Path, timeout: int) -> None:
    proc = subprocess.Popen(
        [str(frida), "-U", "-p", str(pid), "-l", str(script_path), "-o", str(log_path)],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    try:
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            if proc.poll() is not None:
                return
            if log_path.exists():
                text = log_path.read_text("utf-8", errors="ignore")
                if "[H5_ORACLE_DONE]" in text or "[H5_ORACLE_FATAL]" in text:
                    break
            time.sleep(0.25)
    except subprocess.TimeoutExpired:
        pass
    finally:
        if proc.poll() is None:
            proc.terminate()
            try:
                proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                proc.kill()


def _parse_oracle_log(log_path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    if not log_path.exists():
        return rows
    for line in log_path.read_text("utf-8", errors="ignore").splitlines():
        if not line.startswith("[H5_ORACLE] "):
            continue
        try:
            item = json.loads(line.split(" ", 1)[1])
            ret = item.get("ret") or ""
            if len(ret) == 68:
                item["kind"] = "h5_68"
                item["crc32"] = parse_h5_sig3(ret).crc32
            elif len(ret) == 64 and ret.startswith("5a54"):
                item["kind"] = "atlas64"
                item["head8"] = ret[:16]
                item["body_hex"] = ret[16:]
            else:
                item["kind"] = f"len_{len(ret)}"
            rows.append(item)
        except Exception:
            pass
    return rows


def main() -> int:
    parser = argparse.ArgumentParser(description="Call APP H5 atlasSign as a secPlain oracle")
    parser.add_argument("--har", default=str(DEFAULT_HAR))
    parser.add_argument("--adb", default=str(DEFAULT_ADB))
    parser.add_argument("--frida", default=str(DEFAULT_FRIDA))
    parser.add_argument("--out-dir", default="out")
    parser.add_argument("--timeout", type=int, default=25)
    args = parser.parse_args()

    samples = _load_h5_samples(Path(args.har))
    candidates = _candidate_map(samples)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    ts = time.strftime("%Y%m%d_%H%M%S")
    script_path = out_dir / f"h5_atlas_oracle_{ts}.js"
    log_path = out_dir / f"h5_atlas_oracle_{ts}.log"
    script_path.write_text(_json_script(candidates), encoding="utf-8")

    pid = _ensure_app(Path(args.adb))
    _run_frida(Path(args.frida), pid, script_path, log_path, args.timeout)

    rows = _parse_oracle_log(log_path)
    target_by_sample = {str(s.index): s for s in samples}
    hits: list[tuple[H5Sample, str, str]] = []
    for row in rows:
        row_id = str(row.get("id", ""))
        sample_id, _, candidate_name = row_id.partition(":")
        sample = target_by_sample.get(sample_id)
        if sample is None:
            continue
        if row.get("crc32") == sample.target_crc32:
            hits.append((sample, candidate_name, row.get("plain", "")))

    h5_rows = [row for row in rows if row.get("kind") == "h5_68"]
    atlas_rows = [row for row in rows if row.get("kind") == "atlas64"]
    atlas_heads = sorted({str(row.get("head8")) for row in atlas_rows})
    print(
        f"samples={len(samples)} candidates={len(candidates)} "
        f"oracle_rows={len(rows)} h5_68_rows={len(h5_rows)} atlas64_rows={len(atlas_rows)}"
    )
    if atlas_heads:
        print("atlas64_head8=" + ",".join(atlas_heads))
    print(f"script={script_path}")
    print(f"log={log_path}")
    if hits:
        print("hits:")
        for sample, name, plain in hits:
            preview = plain if len(plain) <= 160 else plain[:160] + "..."
            print(f"- #{sample.index} {sample.path} crc={sample.target_crc32:08x} {name}: {preview}")
    else:
        print("hits=0")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())