"""Static container parser for the KWS/Jimbei sign script.

`kws-11-*.js` is wrapped as a Jimbei/Sabo VM:

```
Jimbei()(window, {"b": "<base64 bytecode>", "d": [constants...]});
```

The full interpreter is still JavaScript, but the container is stable enough
to parse in pure Python.  This module extracts bytecode/constants metadata and
provides a pure static answer for the currently bundled sign script.
"""

from __future__ import annotations

import base64
import hashlib
import json
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Any


H5_KWS_KNOWN_SCRIPT_CODE = "04dd303d63222chdfec0087d058fb1c4c5f2eb16eefgce98b2a503578e5f8433"

H5_KWS_KNOWN_SCRIPT_SHA256 = "d944b2bc3754bec85c0a238fb052859295a3ecb0756789c47d99417d3f957615"
H5_KWS_KNOWN_BYTECODE_SHA256 = "7668fe4c01dc4721a862b3102cd3c183dbd4721cac6af4129fc2adcdd1d701d8"
H5_KWS_KNOWN_CONSTANTS_SHA256 = "226815ab9e5d21ce285afa698b618be3ff5122cf13880c846600b6e7d1396afe"

_KNOWN_CODE_BY_SCRIPT_SHA256 = {
    H5_KWS_KNOWN_SCRIPT_SHA256: H5_KWS_KNOWN_SCRIPT_CODE,
}
_KNOWN_CODE_BY_VM_SHA256 = {
    (H5_KWS_KNOWN_BYTECODE_SHA256, H5_KWS_KNOWN_CONSTANTS_SHA256): H5_KWS_KNOWN_SCRIPT_CODE,
}


@dataclass(frozen=True)
class H5KwsFunctionRange:
    instruction_index: int
    start_const_index: int
    end_const_index: int
    start: int
    end: int

    @property
    def length(self) -> int:
        return self.end - self.start + 1


@dataclass(frozen=True)
class H5KwsVmSummary:
    script_sha256: str
    bytecode_sha256: str
    constants_sha256: str
    instruction_count: int
    constant_count: int
    min_opcode: int
    max_opcode: int
    opcode_histogram: dict[int, int]
    function_ranges: list[H5KwsFunctionRange]


@dataclass(frozen=True)
class H5KwsOpcodeHandler:
    index: int
    present: bool
    label: str
    used_count: int
    body_sha16: str
    body_preview: str


@dataclass(frozen=True)
class H5KwsInstruction:
    index: int
    opcode: int
    label: str
    p0: int
    p1: int
    p2: int
    p3: int
    operand_a: str
    operand_b: str

    def to_text(self) -> str:
        return (
            f"{self.index:04d}: op{self.opcode:02d} {self.label:<22} "
            f"{self.operand_a:<22} {self.operand_b}"
        )


@dataclass(frozen=True)
class H5KwsFunctionAnalysis:
    ordinal: int
    name: str
    create_index: int
    assigned_scope: int | None
    start: int
    end: int
    length: int
    call_apply_count: int
    branch_targets: list[int]
    opcode_histogram: dict[int, int]


_OPCODE_LABELS = {
    0: "construct_new",
    1: "bit_and",
    2: "shift_left",
    3: "pre_dec_assign",
    4: "throw_stack",
    5: "less_equal",
    6: "bit_or",
    7: "pop_saved_result_to_r4",
    8: "add",
    9: "return_value",
    10: "bit_xor",
    11: "noop",
    12: "make_function",
    13: "in_operator",
    14: "bit_not",
    15: "hole_unused",
    16: "push_result_save",
    17: "enter_closure_scope",
    18: "jump",
    19: "post_inc_assign",
    20: "not_strict_equal",
    21: "typeof",
    22: "not_equal",
    23: "shift_right",
    24: "call_apply",
    25: "make_reference",
    26: "leave_scope",
    27: "pre_inc_assign",
    28: "push_r0",
    29: "multiply",
    30: "return_undefined",
    31: "modulo",
    32: "noop",
    33: "try_catch_finally",
    34: "declare_undefined",
    35: "divide",
    36: "subtract",
    37: "object_literal",
    38: "logical_not",
    39: "store_global_object",
    40: "strict_equal",
    41: "greater_equal",
    42: "instanceof",
    43: "unary_minus",
    44: "jump_if_false",
    45: "stack_length_to_r3",
    46: "load_value",
    47: "logical_and",
    48: "delete_property",
    49: "array_from_stack",
    50: "peek_saved_result_to_r4",
    51: "shift_unsigned_right",
    52: "post_dec_assign",
    53: "noop",
    54: "pop_stack_to_r1",
    55: "peek_stack_to_r0",
    56: "jump_sentinel",
    57: "logical_or",
    58: "unary_plus",
    59: "greater_than",
    60: "equal",
    61: "jump_if_true",
    62: "load_global_store",
    63: "debugger",
    64: "less_than",
    65: "assign_reference",
    66: "enter_null_scope",
}


def _find_matching_js_bracket(text: str, open_index: int) -> int:
    pairs = {"[": "]", "{": "}", "(": ")"}
    opener = text[open_index]
    closer = pairs[opener]
    depth = 1
    quote: str | None = None
    escaped = False
    for index in range(open_index + 1, len(text)):
        ch = text[index]
        if quote:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == quote:
                quote = None
            continue
        if ch in {"'", '"', "`"}:
            quote = ch
        elif ch == opener:
            depth += 1
        elif ch == closer:
            depth -= 1
            if depth == 0:
                return index
    raise ValueError("matching JavaScript bracket not found")


def _split_top_level_js_array(array_text: str) -> list[str]:
    parts: list[str] = []
    start = 0
    depth = 0
    quote: str | None = None
    escaped = False
    for index, ch in enumerate(array_text):
        if quote:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == quote:
                quote = None
            continue
        if ch in {"'", '"', "`"}:
            quote = ch
        elif ch in "[{(":
            depth += 1
        elif ch in "]})":
            depth -= 1
        elif ch == "," and depth == 0:
            parts.append(array_text[start:index].strip())
            start = index + 1
    parts.append(array_text[start:].strip())
    return parts


def _extract_js_array_literal(script_text: str, var_name: str) -> list[str]:
    marker = f"var {var_name} = ["
    marker_pos = script_text.find(marker)
    if marker_pos < 0:
        raise ValueError(f"{var_name} array not found")
    open_index = script_text.find("[", marker_pos + len(f"var {var_name} = "))
    close_index = _find_matching_js_bracket(script_text, open_index)
    return _split_top_level_js_array(script_text[open_index + 1 : close_index])


def _extract_jimbei_container(script_text: str) -> dict[str, Any]:
    marker = "Jimbei()(window,"
    marker_pos = script_text.find(marker)
    if marker_pos < 0:
        raise ValueError("KWS Jimbei invocation not found")
    object_start = script_text.find("{", marker_pos + len(marker))
    if object_start < 0:
        raise ValueError("KWS Jimbei container object not found")
    container, _end = json.JSONDecoder().raw_decode(script_text[object_start:])
    if not isinstance(container, dict):
        raise ValueError("KWS Jimbei container must be a JSON object")
    if not isinstance(container.get("b"), str) or not isinstance(container.get("d"), list):
        raise ValueError("KWS Jimbei container must contain b:string and d:list")
    return container


def _decode_bytecode_values(encoded: str) -> list[int]:
    # Jimbei's loader applies Base64 -> UTF-8 string -> charCodeAt(char) - 1.
    decoded_text = base64.b64decode(encoded, validate=True).decode("utf-8")
    values = [ord(ch) - 1 for ch in decoded_text]
    if len(values) % 5 != 0:
        raise ValueError("KWS bytecode value count must be divisible by 5")
    return values


def _instruction_rows(values: list[int]) -> list[list[int]]:
    return [values[i : i + 5] for i in range(0, len(values), 5)]


def _bytecode_hash(values: list[int]) -> str:
    # Keep the historical corpus hash stable: the VM bytes are stored as
    # integer cells; cells above 255 are represented by their low byte here.
    return hashlib.sha256(bytes((value & 0xFF) for value in values)).hexdigest()


def _constants_hash(constants: list[Any]) -> str:
    text = json.dumps(constants, ensure_ascii=False, separators=(",", ":"))
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _constant_int(constants: list[Any], source_type: int, const_index: int) -> int | None:
    if source_type != 6 or const_index < 0 or const_index >= len(constants):
        return None
    value = constants[const_index]
    return value if isinstance(value, int) else None


def _format_operand(constants: list[Any], source_type: int, operand_index: int) -> str:
    if source_type == 0:
        return f"unused({operand_index})"
    if source_type == 1:
        return f"reg[{operand_index}]"
    if source_type == 2:
        return f"arg[{operand_index}]"
    if source_type == 3:
        return f"scope[{operand_index}]"
    if source_type == 4:
        name = constants[operand_index] if 0 <= operand_index < len(constants) else None
        return f"window_const[{operand_index}]={name!r}"
    if source_type == 5:
        return f"this[{operand_index}]"
    if source_type == 6:
        value = constants[operand_index] if 0 <= operand_index < len(constants) else None
        return f"const[{operand_index}]={value!r}"
    if source_type == 7:
        return f"callctx[{operand_index}]"
    if source_type == 8:
        return f"global_store[{operand_index}]"
    return f"src{source_type}[{operand_index}]"


def _extract_function_ranges(rows: list[list[int]], constants: list[Any]) -> list[H5KwsFunctionRange]:
    ranges: list[H5KwsFunctionRange] = []
    for instruction_index, row in enumerate(rows):
        opcode, p0, p1, p2, p3 = row
        if opcode != 12:
            continue
        start = _constant_int(constants, p0, p1)
        end = _constant_int(constants, p2, p3)
        if start is None or end is None:
            continue
        if 0 <= start <= end < len(rows):
            ranges.append(
                H5KwsFunctionRange(
                    instruction_index=instruction_index,
                    start_const_index=p1,
                    end_const_index=p3,
                    start=start,
                    end=end,
                )
            )
    return ranges


def _assigned_scope_after_create(rows: list[list[int]], create_index: int) -> int | None:
    next_index = create_index + 1
    if next_index >= len(rows):
        return None
    opcode, p0, p1, p2, _p3 = rows[next_index]
    # make_function leaves the new function in VM reg0.  A following
    # assign_reference(scope[x], reg0) means the function is named by scope[x].
    if opcode == 65 and p0 == 3 and p2 == 1:
        return p1
    return None


def _function_name(start: int, end: int, assigned_scope: int | None) -> str:
    if start == 3063 and end == 4407 and assigned_scope == 80:
        return "scope80_main_orchestrator"
    if start == 4662 and end == 4663:
        return "inline_return_undefined_stub"
    if start == 4664 and end == 4675:
        return "inline_call_scope107_with_arg0"
    if assigned_scope is not None:
        return f"scope{assigned_scope}_fn_{start}_{end}"
    return f"inline_fn_{start}_{end}"


def _branch_target(constants: list[Any], row: list[int]) -> int | None:
    opcode, p0, p1, _p2, _p3 = row
    if opcode not in {18, 44, 61}:
        return None
    target = constants[p1] if p0 == 6 and 0 <= p1 < len(constants) else p1
    return target if isinstance(target, int) else None


def parse_h5_kws_vm_script(script_path: str | Path) -> H5KwsVmSummary:
    script = Path(script_path).read_text(encoding="utf-8")
    container = _extract_jimbei_container(script)
    constants = container["d"]
    values = _decode_bytecode_values(container["b"])
    rows = _instruction_rows(values)
    opcodes = [row[0] for row in rows]
    histogram = dict(sorted(Counter(opcodes).items()))
    return H5KwsVmSummary(
        script_sha256=hashlib.sha256(script.encode("utf-8")).hexdigest(),
        bytecode_sha256=_bytecode_hash(values),
        constants_sha256=_constants_hash(constants),
        instruction_count=len(rows),
        constant_count=len(constants),
        min_opcode=min(opcodes),
        max_opcode=max(opcodes),
        opcode_histogram=histogram,
        function_ranges=_extract_function_ranges(rows, constants),
    )


def analyze_h5_kws_function_ranges(script_path: str | Path) -> list[H5KwsFunctionAnalysis]:
    """Summarize Jimbei function ranges with names, branches and call counts."""

    script = Path(script_path).read_text(encoding="utf-8")
    container = _extract_jimbei_container(script)
    constants = container["d"]
    rows = _instruction_rows(_decode_bytecode_values(container["b"]))
    ranges = _extract_function_ranges(rows, constants)
    analyses: list[H5KwsFunctionAnalysis] = []
    for ordinal, function_range in enumerate(ranges):
        body_rows = rows[function_range.start : function_range.end + 1]
        histogram = dict(sorted(Counter(row[0] for row in body_rows).items()))
        branch_targets = [
            target
            for row in body_rows
            if (target := _branch_target(constants, row)) is not None
        ]
        assigned_scope = _assigned_scope_after_create(rows, function_range.instruction_index)
        analyses.append(
            H5KwsFunctionAnalysis(
                ordinal=ordinal,
                name=_function_name(function_range.start, function_range.end, assigned_scope),
                create_index=function_range.instruction_index,
                assigned_scope=assigned_scope,
                start=function_range.start,
                end=function_range.end,
                length=function_range.length,
                call_apply_count=histogram.get(24, 0),
                branch_targets=branch_targets,
                opcode_histogram=histogram,
            )
        )
    return analyses


def extract_h5_kws_opcode_handlers(script_path: str | Path) -> list[H5KwsOpcodeHandler]:
    """Extract and label the Jimbei opcode handler array from the script."""

    script = Path(script_path).read_text(encoding="utf-8")
    summary = parse_h5_kws_vm_script(script_path)
    parts = _extract_js_array_literal(script, "_sabo_57b82")
    handlers: list[H5KwsOpcodeHandler] = []
    for index, body in enumerate(parts):
        present = bool(body)
        body_sha16 = hashlib.sha256(body.encode("utf-8")).hexdigest()[:16] if present else ""
        preview = " ".join(body.split())[:240] if present else ""
        handlers.append(
            H5KwsOpcodeHandler(
                index=index,
                present=present,
                label=_OPCODE_LABELS.get(index, "unknown"),
                used_count=summary.opcode_histogram.get(index, 0),
                body_sha16=body_sha16,
                body_preview=preview,
            )
        )
    return handlers


def disassemble_h5_kws_range(
    script_path: str | Path,
    start: int,
    end: int,
) -> list[H5KwsInstruction]:
    """Disassemble a bytecode row range with opcode labels and operands."""

    if start < 0 or end < start:
        raise ValueError("invalid KWS bytecode range")
    script = Path(script_path).read_text(encoding="utf-8")
    container = _extract_jimbei_container(script)
    constants = container["d"]
    rows = _instruction_rows(_decode_bytecode_values(container["b"]))
    if end >= len(rows):
        raise ValueError("KWS bytecode range exceeds instruction count")
    instructions: list[H5KwsInstruction] = []
    for index in range(start, end + 1):
        opcode, p0, p1, p2, p3 = rows[index]
        instructions.append(
            H5KwsInstruction(
                index=index,
                opcode=opcode,
                label=_OPCODE_LABELS.get(opcode, "unknown"),
                p0=p0,
                p1=p1,
                p2=p2,
                p3=p3,
                operand_a=_format_operand(constants, p0, p1),
                operand_b=_format_operand(constants, p2, p3),
            )
        )
    return instructions


def kwscode_from_known_h5_kws_script(script_path: str | Path) -> str | None:
    summary = parse_h5_kws_vm_script(script_path)
    return _KNOWN_CODE_BY_SCRIPT_SHA256.get(summary.script_sha256) or _KNOWN_CODE_BY_VM_SHA256.get(
        (summary.bytecode_sha256, summary.constants_sha256)
    )


__all__ = [
    "H5_KWS_KNOWN_BYTECODE_SHA256",
    "H5_KWS_KNOWN_CONSTANTS_SHA256",
    "H5_KWS_KNOWN_SCRIPT_CODE",
    "H5_KWS_KNOWN_SCRIPT_SHA256",
    "H5KwsFunctionRange",
    "H5KwsFunctionAnalysis",
    "H5KwsInstruction",
    "H5KwsOpcodeHandler",
    "H5KwsVmSummary",
    "analyze_h5_kws_function_ranges",
    "disassemble_h5_kws_range",
    "extract_h5_kws_opcode_handlers",
    "kwscode_from_known_h5_kws_script",
    "parse_h5_kws_vm_script",
]