"""NTCP2 connector script for podman integration tests.

Waits for the listener's key file, connects to it, performs the
Noise_XK handshake, exchanges frames, writes results to a JSON file.

Usage:
    python router_connector.py --host 127.0.0.1 --port 6000 \
        --key-file /shared/listener_key.bin \
        --result-file /shared/connector_result.json
"""

import argparse
import asyncio
import json
import os
import struct
import sys
import time
import traceback

sys.path.insert(0, "/app/src")

from i2p_crypto.x25519 import X25519DH
from i2p_transport.ntcp2 import NTCP2Frame, FrameType
from i2p_transport.ntcp2_connection import NTCP2Connection
from i2p_transport.ntcp2_handshake import NTCP2Handshake


async def _send_hs(writer, msg):
    writer.write(struct.pack("!H", len(msg)) + msg)
    await writer.drain()


async def _recv_hs(reader):
    lb = await reader.readexactly(2)
    length = struct.unpack("!H", lb)[0]
    return await reader.readexactly(length)


def wait_for_key_file(path, timeout=30):
    """Block until the listener's key file appears."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        if os.path.exists(path) and os.path.getsize(path) == 32:
            with open(path, "rb") as f:
                return f.read()
        time.sleep(0.2)
    raise TimeoutError(f"Key file {path} not found after {timeout}s")


async def run_connector(host, port, key_file, result_file):
    result = {"status": "error", "error": "unknown"}

    try:
        peer_static_pub = wait_for_key_file(key_file)
        print(f"Got listener key, connecting to {host}:{port}", flush=True)

        static = X25519DH.generate_keypair()
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(host, port), timeout=10.0
        )

        hs = NTCP2Handshake(
            our_static=static, peer_static_pub=peer_static_pub, initiator=True
        )
        msg1 = hs.create_message_1()
        await _send_hs(writer, msg1)
        msg2 = await asyncio.wait_for(_recv_hs(reader), timeout=10.0)
        msg3 = hs.process_message_2(msg2)
        await _send_hs(writer, msg3)

        send_cipher, recv_cipher = hs.split()
        conn = NTCP2Connection(
            reader=reader, writer=writer,
            cipher_send=send_cipher, cipher_recv=recv_cipher,
            remote_hash=peer_static_pub,
        )
        print("Handshake complete", flush=True)

        # Send a frame to the listener
        frame = NTCP2Frame(FrameType.I2NP, b"hello from connector")
        await conn.send_frame(frame)

        # Receive the listener's reply
        received = await asyncio.wait_for(conn.recv_frame(), timeout=10.0)
        print(f"Received: type={received.frame_type}, payload={received.payload!r}", flush=True)

        result = {
            "status": "ok",
            "handshake": "complete",
            "received_type": received.frame_type.value,
            "received_payload": received.payload.decode("utf-8", errors="replace"),
            "sent_frame": True,
        }

        conn._writer.close()
    except Exception as e:
        result = {"status": "error", "error": str(e)}
        traceback.print_exc()
    finally:
        with open(result_file, "w") as f:
            json.dump(result, f)
        print(f"Result: {result}", flush=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument("--key-file", required=True)
    parser.add_argument("--result-file", required=True)
    args = parser.parse_args()
    asyncio.run(run_connector(args.host, args.port, args.key_file, args.result_file))


if __name__ == "__main__":
    main()