#!/usr/bin/env python3
"""
Generate src/test_vectors.rs from exported Python Reticulum vectors JSON.

Usage:
  python3 tools/generate_test_vectors_rs.py
  python3 tools/generate_test_vectors_rs.py --input vectors/python_vectors.json
  python3 tools/generate_test_vectors_rs.py --input vectors/python_vectors.json --output src/test_vectors.rs
"""

from __future__ import annotations

import argparse
import json
import pathlib
import textwrap
from typing import Iterable, List


DEFAULT_INPUT = pathlib.Path("vectors/python_vectors.json")
DEFAULT_OUTPUT = pathlib.Path("src/test_vectors.rs")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate src/test_vectors.rs from exported Python Reticulum vectors."
    )
    parser.add_argument(
        "--input",
        type=pathlib.Path,
        default=DEFAULT_INPUT,
        help=f"Input JSON from export_python_vectors.py (default: {DEFAULT_INPUT})",
    )
    parser.add_argument(
        "--output",
        type=pathlib.Path,
        default=DEFAULT_OUTPUT,
        help=f"Output Rust file (default: {DEFAULT_OUTPUT})",
    )
    return parser.parse_args()


def require_hex(data: str, expected_len: int | None = None) -> bytes:
    raw = bytes.fromhex(data)
    if expected_len is not None and len(raw) != expected_len:
        raise ValueError(f"expected {expected_len} bytes, got {len(raw)} from {data!r}")
    return raw


def rust_byte_array(name: str, data: bytes) -> str:
    lines: List[str] = [f"pub const {name}: [u8; {len(data)}] = ["]
    width = 8
    for offset in range(0, len(data), width):
        chunk = data[offset:offset + width]
        items = ", ".join(f"0x{byte:02x}" for byte in chunk)
        lines.append(f"    {items},")
    lines.append("];")
    return "\n".join(lines)


def rust_str_const(name: str, value: str) -> str:
    return f'pub const {name}: &str =\n    "{value}";'


def build_module(payload: dict) -> str:
    fixed = payload["fixed_inputs"]
    vectors = payload["vectors"]

    dest_identity = require_hex(fixed["destination_identity_hex"], 64)
    link_owner_identity = require_hex(fixed["link_owner_identity_hex"], 64)
    announce_rng = require_hex(fixed["announce_rng_bytes_hex"], 32)
    announce_random_hash = require_hex(fixed["announce_random_hash_hex"], 16)
    forwarded_transport_id = require_hex(fixed["forwarded_transport_id_hex"], 16)
    lrproof_link_id = require_hex(fixed["lrproof_link_id_hex"], 16)
    lrproof_x25519_private = require_hex(fixed["lrproof_x25519_private_hex"], 32)
    lrproof_ed25519_private = require_hex(fixed["lrproof_ed25519_private_hex"], 32)

    destination_private = dest_identity[:32]
    destination_signing = dest_identity[32:]
    link_owner_private = link_owner_identity[:32]
    link_owner_signing = link_owner_identity[32:]

    announce_packet_hex = vectors["announce"]["raw_hex"]
    path_response_packet_hex = vectors["path_response"]["raw_hex"]
    forwarded_announce_packet_hex = vectors["forwarded_announce"]["raw_hex"]
    explicit_packet_proof_hex = vectors["explicit_packet_proof"]["raw_hex"]
    lrproof_packet_hex = vectors["lrproof"]["raw_hex"]
    path_request_no_transport_data_hex = vectors["path_request_without_transport"]["data_hex"]
    path_request_with_transport_data_hex = vectors["path_request_with_transport"]["data_hex"]

    fixed_announce_time = int(payload["meta"]["fixed_announce_time"])

    sections = [
        "// Generated by tools/generate_test_vectors_rs.py from Python Reticulum vectors.",
        "// Do not edit by hand; regenerate from exported JSON instead.",
        "",
        "use ed25519_dalek::SigningKey;",
        "use rand_core::{CryptoRng, Error, RngCore};",
        "use x25519_dalek::StaticSecret;",
        "",
        "use crate::identity::PrivateIdentity;",
        "",
        rust_byte_array("FIXED_IDENTITY_PRIVATE_KEY", destination_private),
        rust_byte_array("FIXED_IDENTITY_SIGNING_KEY", destination_signing),
        "",
        rust_byte_array("FIXED_LINK_OWNER_PRIVATE_KEY", link_owner_private),
        rust_byte_array("FIXED_LINK_OWNER_SIGNING_KEY", link_owner_signing),
        "",
        rust_byte_array("FIXED_ANNOUNCE_RNG_BYTES", announce_rng),
        rust_byte_array("FIXED_ANNOUNCE_RANDOM_HASH_BYTES", announce_random_hash),
        f"pub const FIXED_ANNOUNCE_TIMESTAMP: u64 = {fixed_announce_time:_};",
        rust_byte_array("FIXED_LRPROOF_LINK_ID", lrproof_link_id),
        rust_byte_array(
            "FIXED_FORWARDED_ANNOUNCE_TRANSPORT_ID", forwarded_transport_id
        ),
        rust_byte_array(
            "FIXED_LRPROOF_X25519_PRIVATE_KEY", lrproof_x25519_private
        ),
        rust_byte_array(
            "FIXED_LRPROOF_ED25519_PRIVATE_KEY", lrproof_ed25519_private
        ),
        "",
        rust_str_const("ANNOUNCE_PACKET_HEX", announce_packet_hex),
        rust_str_const("PATH_RESPONSE_PACKET_HEX", path_response_packet_hex),
        rust_str_const("FORWARDED_ANNOUNCE_PACKET_HEX", forwarded_announce_packet_hex),
        rust_str_const("EXPLICIT_PACKET_PROOF_HEX", explicit_packet_proof_hex),
        rust_str_const("LRPROOF_PACKET_HEX", lrproof_packet_hex),
        rust_str_const(
            "PATH_REQUEST_NO_TRANSPORT_DATA_HEX", path_request_no_transport_data_hex
        ),
        rust_str_const(
            "PATH_REQUEST_WITH_TRANSPORT_DATA_HEX", path_request_with_transport_data_hex
        ),
        "",
        textwrap.dedent(
            """
            #[derive(Copy, Clone)]
            pub struct FixedRng {
                bytes: [u8; 32],
                offset: usize,
            }

            impl FixedRng {
                pub fn new(bytes: [u8; 32]) -> Self {
                    Self { bytes, offset: 0 }
                }
            }

            impl RngCore for FixedRng {
                fn next_u32(&mut self) -> u32 {
                    let mut buf = [0u8; 4];
                    self.fill_bytes(&mut buf);
                    u32::from_le_bytes(buf)
                }

                fn next_u64(&mut self) -> u64 {
                    let mut buf = [0u8; 8];
                    self.fill_bytes(&mut buf);
                    u64::from_le_bytes(buf)
                }

                fn fill_bytes(&mut self, dest: &mut [u8]) {
                    for byte in dest {
                        *byte = self.bytes[self.offset % self.bytes.len()];
                        self.offset += 1;
                    }
                }

                fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
                    self.fill_bytes(dest);
                    Ok(())
                }
            }

            impl CryptoRng for FixedRng {}

            pub fn decode_hex(hex: &str) -> Vec<u8> {
                assert_eq!(hex.len() % 2, 0, "hex string must be even-length");
                (0..hex.len())
                    .step_by(2)
                    .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).expect("valid hex byte"))
                    .collect()
            }

            pub fn fixed_private_identity() -> PrivateIdentity {
                PrivateIdentity::new(
                    StaticSecret::from(FIXED_IDENTITY_PRIVATE_KEY),
                    SigningKey::from_bytes(&FIXED_IDENTITY_SIGNING_KEY),
                )
            }
            """
        ).strip(),
        "",
    ]

    return "\n".join(sections)


def main() -> int:
    args = parse_args()
    payload = json.loads(args.input.read_text(encoding="utf-8"))
    content = build_module(payload)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(content, encoding="utf-8")
    return 0


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