#!/usr/bin/env python3
"""
otp_roundtrip_test.py — end-to-end test of otp_encrypt / otp_decrypt verbs.

Sends framed JSON-RPC requests to nsigner --listen stdio and checks the
round-trip: plaintext -> otp_encrypt -> otp_decrypt -> recovered plaintext.

Framing: 4-byte big-endian length prefix + JSON payload.

Usage: python3 tools/otp_roundtrip_test.py
"""
import base64
import json
import os
import struct
import subprocess
import sys
import tempfile
import time

NSIGNER = "./build/nsigner"
PAD_DIR = "/media/user/Music/pads"
PAD_SPEC = "333e9902db839d9d"
MNEMONIC_FILE = ".test_mnemonic"


def send_framed(proc, obj):
    payload = json.dumps(obj).encode()
    proc.stdin.write(struct.pack(">I", len(payload)))
    proc.stdin.write(payload)
    proc.stdin.flush()


def recv_framed(proc):
    """Read a framed response, skipping any banner text the signer writes
    to stdout before the first frame. The banner is line-based ASCII; a
    valid frame starts with a 4-byte big-endian length followed by '{'."""
    # Read 4 bytes at a time, sliding window, until we find a frame header.
    buf = b""
    while True:
        b = proc.stdout.read(1)
        if not b:
            return None
        buf = (buf + b)[-4:]
        if len(buf) < 4:
            continue
        (length,) = struct.unpack(">I", buf)
        # Sanity: frame length should be reasonable (1..1MB) and the next
        # byte after the header should be '{' (start of JSON).
        if 1 <= length <= 1024 * 1024:
            # Peek: read one more byte to check for '{'.
            peek = proc.stdout.read(1)
            if peek == b'{':
                body = peek + proc.stdout.read(length - 1)
                return json.loads(body.decode())
            else:
                # Not a frame; prepend peek to the stream by including it
                # in the sliding window.
                buf = (buf + peek)[-4:]
        # Otherwise keep scanning.


def main():
    if not os.path.isfile(NSIGNER):
        print(f"ERROR: {NSIGNER} not found. Run 'make dev' first.")
        return 1

    with open(MNEMONIC_FILE) as f:
        mnemonic = f.read().strip()

    plaintext = b"Hello, OTP world!"
    pt_b64 = base64.b64encode(plaintext).decode()
    print(f"Plaintext: {plaintext.decode()}")
    print(f"Plaintext base64: {pt_b64}")

    # Write the mnemonic to a fixed temp file and pass it as fd 3 to nsigner
    # via a bash wrapper (so stdin stays free for framed requests).
    mnem_path = ".test_mnemonic_otp_roundtrip.tmp"
    with open(mnem_path, "w") as f:
        f.write(mnemonic + "\n")

    def run_one_request(req_obj):
        """Run nsigner in stdio mode for a single framed request/response.
        Returns the parsed JSON response or None."""
        shell_cmd = (
            f"exec 3<{mnem_path}; "
            f"exec {NSIGNER} --listen stdio --mnemonic-fd 3 "
            f"--otp-pad-dir {PAD_DIR} --otp-pad {PAD_SPEC} "
            f"--otp-allow-blkback --allow-all"
        )
        proc = subprocess.Popen(
            ["bash", "-c", shell_cmd],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            text=False,
        )
        time.sleep(1.5)  # let the signer bind the pad and be ready
        if proc.poll() is not None:
            err = proc.stderr.read().decode()
            print(f"ERROR: nsigner exited early (code {proc.returncode})")
            print(f"stderr: {err}")
            return None
        send_framed(proc, req_obj)
        resp = recv_framed(proc)
        proc.stdin.close()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
        return resp

    # --- otp_encrypt ---
    print("\n=== Sending otp_encrypt ===")
    resp = run_one_request({
        "id": "1",
        "method": "otp_encrypt",
        "params": [pt_b64, {"encoding": "ascii"}],
    })
    print(f"Encrypt response: {json.dumps(resp)}")

    if resp is None or "result" not in resp:
        print("ERROR: no result in encrypt response")
        try:
            os.unlink(mnem_path)
        except OSError:
            pass
        return 1

    result_obj = json.loads(resp["result"])
    ciphertext = result_obj["ciphertext"]
    off_before = result_obj["pad_offset_before"]
    off_after = result_obj["pad_offset_after"]
    print(f"Pad offset: {off_before} -> {off_after} "
          f"(consumed {off_after - off_before} bytes)")
    print(f"Ciphertext (first 80 chars): {ciphertext[:80]}...")

    # --- otp_decrypt (separate invocation; offset persists in .state) ---
    print("\n=== Sending otp_decrypt ===")
    resp2 = run_one_request({
        "id": "2",
        "method": "otp_decrypt",
        "params": [ciphertext, {"encoding": "ascii"}],
    })
    print(f"Decrypt response: {json.dumps(resp2)}")

    if resp2 is None or "result" not in resp2:
        print("ERROR: no result in decrypt response")
        return 1

    result2 = json.loads(resp2["result"])
    recovered_b64 = result2["plaintext"]
    recovered = base64.b64decode(recovered_b64)
    print(f"\nRecovered plaintext: {recovered.decode()}")

    if recovered == plaintext:
        print("\n=== ASCII ROUND-TRIP SUCCESS ===")
    else:
        print("\n=== ASCII ROUND-TRIP FAILED ===")
        print(f"Expected: {plaintext.decode()}")
        print(f"Got:      {recovered.decode()}")
        return 1

    # --- Binary encoding round-trip ---
    # Reset the pad offset for a clean binary test.
    state_path = (f"{PAD_DIR}/{result_obj['pad_chksum']}.state")
    with open(state_path, "w") as sf:
        sf.write("offset=32\n")

    print("\n=== Sending otp_encrypt (binary) ===")
    resp3 = run_one_request({
        "id": "3",
        "method": "otp_encrypt",
        "params": [pt_b64, {"encoding": "binary"}],
    })
    print(f"Binary encrypt response: {json.dumps(resp3)}")
    if resp3 is None or "result" not in resp3:
        print("ERROR: no result in binary encrypt response")
        return 1
    result3 = json.loads(resp3["result"])
    bin_b64 = result3["ciphertext"]
    # The binary ciphertext is base64-encoded in the JSON result.
    bin_blob = base64.b64decode(bin_b64)
    print(f"Binary blob size: {len(bin_blob)} bytes "
          f"(header 58 + padded data {len(bin_blob) - 58})")
    if not bin_blob[:4] == b"OTP\0":
        print("ERROR: binary blob missing OTP magic")
        return 1

    print("\n=== Sending otp_decrypt (binary) ===")
    resp4 = run_one_request({
        "id": "4",
        "method": "otp_decrypt",
        "params": [bin_b64, {"encoding": "binary"}],
    })
    print(f"Binary decrypt response: {json.dumps(resp4)}")
    if resp4 is None or "result" not in resp4:
        print("ERROR: no result in binary decrypt response")
        return 1
    result4 = json.loads(resp4["result"])
    recovered2 = base64.b64decode(result4["plaintext"])
    print(f"\nRecovered plaintext (binary path): {recovered2.decode()}")

    if recovered2 == plaintext:
        print("\n=== BINARY ROUND-TRIP SUCCESS ===")
        return 0
    else:
        print("\n=== BINARY ROUND-TRIP FAILED ===")
        print(f"Expected: {plaintext.decode()}")
        print(f"Got:      {recovered2.decode()}")
        return 1


if __name__ == "__main__":
    sys.exit(main())
