#!/usr/bin/env python3
"""Xtensa-vs-other ISA determination for the DMCUB payload, with in-tree
positive and negative controls (project rule: no unvalidated oracle).

Evidence collected per image:
  * op0 histogram   -- Xtensa's opcode selector is the low nibble of byte 0,
                       and narrow (2-byte) instructions are op0 in {8..D}.
  * ENTRY / RETW.N / RET.N / CALLX8 / L32R counts -- windowed-ABI fingerprints.
  * cosine similarity of the op0 histogram against each control.

Controls:
  positive  psp_out/smu_fw/smu_13_0_0_runtime_body.bin  (established Xtensa)
            psp_out/vcn_ucode.bin                       (established Xtensa)
  negative  psp_out/mec_11_0_0.bin                      (established RV64)
            psp_out/sdma_6_0_0.bin                      (F32, unknown ISA)
            psp_out/asd_code.bin                         (ARM/Thumb)
"""
import math
import os
import sys
import collections

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

CONTROLS = [
    ('+ SMU (Xtensa)', 'psp_out/smu_fw/smu_13_0_0_runtime_body.bin', 0),
    ('+ VCN (Xtensa)', 'psp_out/vcn_ucode.bin', 0),
    ('- MEC (RV64)', 'psp_out/mec_11_0_0.bin', 0),
    ('- SDMA (F32)', 'psp_out/sdma_6_0_0.bin', 0),
    ('- ASD (Thumb)', 'psp_out/asd_code.bin', 0),
]


def op0_hist(data):
    h = [0] * 16
    for b in data[::2]:          # sample byte0 of 2-byte-aligned slots
        h[b & 0xF] += 1
    n = sum(h) or 1
    return [x / n for x in h]


def cosine(a, b):
    num = sum(x * y for x, y in zip(a, b))
    da = math.sqrt(sum(x * x for x in a))
    db = math.sqrt(sum(y * y for y in b))
    return num / (da * db) if da and db else 0.0


def xtensa_marks(data):
    """Count windowed-ABI fingerprints.

    ENTRY   : RRI12, op0=6 op1=3 -> byte0 = 0x36, byte1 low nibble = s (stack reg,
              always a1 in a real prologue) => '36 41'
    RETW.N  : 0xF01D              => bytes '1d f0'
    RET.N   : 0xF00D              => bytes '0d f0'
    CALLX8  : CALLX op, byte0=0xE0 with op1/op2 selecting callx8
    L32R    : op0 = 1 (3-byte, 16-bit literal-pool offset)
    """
    m = collections.Counter()
    n = len(data)
    for i in range(n - 2):
        b0, b1, b2 = data[i], data[i + 1], data[i + 2]
        if b0 == 0x36 and (b1 & 0x0F) == 0x01:
            m['entry_a1'] += 1
        if b0 == 0x1D and b1 == 0xF0:
            m['retw.n'] += 1
        if b0 == 0x0D and b1 == 0xF0:
            m['ret.n'] += 1
        if b0 == 0xE0 and (b1 & 0xF0) == 0x00 and b2 == 0x00:
            m['callx8?'] += 1
    for b in data[::2]:
        if (b & 0xF) == 0x1:
            m['l32r_slots'] += 1
    return m


def report(label, data):
    h = op0_hist(data)
    m = xtensa_marks(data)
    print('%-34s size=%-8d' % (label, len(data)))
    print('   op0 hist: ' + ' '.join('%X:%4.1f' % (i, v * 100) for i, v in enumerate(h)))
    print('   marks:    entry(a1)=%-5d retw.n=%-5d ret.n=%-5d l32r_slots=%d' % (
        m['entry_a1'], m['retw.n'], m['ret.n'], m['l32r_slots']))
    return h


def main():
    target = sys.argv[1]
    off = int(sys.argv[2], 0) if len(sys.argv) > 2 else 0x100
    end = int(sys.argv[3], 0) if len(sys.argv) > 3 else None
    data = open(target, 'rb').read()[off:end]
    print('=== target')
    th = report('%s @0x%x' % (os.path.basename(target), off), data)
    print('=== controls (in-tree, ISA already established)')
    for label, rel, coff in CONTROLS:
        p = os.path.join(REPO, rel)
        if not os.path.exists(p):
            print('%-34s MISSING' % label)
            continue
        ch = report(label, open(p, 'rb').read()[coff:])
        print('   cosine(op0 hist, target) = %.4f' % cosine(th, ch))


if __name__ == '__main__':
    main()