"""Connect to a Java I2P router via NTCP2.

Reads the NTCP2 connection info (host, port, static key) from a JSON
file, attempts the Noise_XK handshake, and reports the result.

Usage:
    python connect_to_java.py --info-file /shared/java_ntcp2_info.json \
        --result-file /shared/interop_result.json
"""

import argparse
import asyncio
import json
import struct
import sys
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)


async def attempt_handshake(host, port, peer_static_pub):
    """Try to perform an NTCP2 handshake with the Java router."""
    our_static = X25519DH.generate_keypair()

    print(f"Connecting to {host}:{port}...", flush=True)
    reader, writer = await asyncio.wait_for(
        asyncio.open_connection(host, port), timeout=10.0
    )
    print("TCP connected", flush=True)

    hs = NTCP2Handshake(
        our_static=our_static,
        peer_static_pub=peer_static_pub,
        initiator=True,
    )

    # Send message 1
    msg1 = hs.create_message_1()
    print(f"Sending msg1 ({len(msg1)} bytes)...", flush=True)
    await _send_hs(writer, msg1)

    # Read message 2
    print("Waiting for msg2...", flush=True)
    msg2 = await asyncio.wait_for(_recv_hs(reader), timeout=15.0)
    print(f"Got msg2 ({len(msg2)} bytes)", flush=True)

    # Process msg2, send msg3
    msg3 = hs.process_message_2(msg2)
    print(f"Sending msg3 ({len(msg3)} bytes)...", flush=True)
    await _send_hs(writer, msg3)

    print("Handshake complete!", flush=True)

    # Try to receive a frame (Java router usually sends DateTime or RouterInfo)
    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,
    )

    received_frames = []
    try:
        for _ in range(3):  # Try to read up to 3 frames
            frame = await asyncio.wait_for(conn.recv_frame(), timeout=5.0)
            received_frames.append({
                "type": frame.frame_type.value,
                "type_name": frame.frame_type.name,
                "payload_len": len(frame.payload),
            })
            print(f"Received frame: type={frame.frame_type.name}, "
                  f"payload_len={len(frame.payload)}", flush=True)
    except asyncio.TimeoutError:
        print("No more frames (timeout)", flush=True)
    except Exception as e:
        print(f"Frame read error: {e}", flush=True)

    conn._writer.close()

    return {
        "handshake": "complete",
        "frames_received": received_frames,
    }


async def run(info_file, result_file):
    result = {"status": "error", "error": "unknown"}

    try:
        with open(info_file) as f:
            info = json.load(f)

        if info.get("status") != "ok":
            result = {"status": "error", "error": f"Bad info file: {info}"}
        else:
            host = info["host"]
            port = info["port"]
            static_key = bytes.fromhex(info["static_key_hex"])

            if len(static_key) != 32:
                result = {"status": "error",
                          "error": f"Invalid static key length: {len(static_key)}"}
            else:
                hs_result = await attempt_handshake(host, port, static_key)
                result = {"status": "ok", **hs_result}

    except asyncio.TimeoutError:
        result = {"status": "error", "error": "timeout during handshake"}
    except ConnectionRefusedError:
        result = {"status": "error", "error": "connection refused"}
    except Exception as e:
        result = {"status": "error", "error": str(e), "type": type(e).__name__}
        traceback.print_exc()

    with open(result_file, "w") as f:
        json.dump(result, f, indent=2)
    print(f"Result: {json.dumps(result, indent=2)}", flush=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--info-file", required=True)
    parser.add_argument("--result-file", required=True)
    args = parser.parse_args()
    asyncio.run(run(args.info_file, args.result_file))


if __name__ == "__main__":
    main()