import argparse
import threading
import sys

from noise_xk.session import NoiseXKSession
from noise_xk.transport import P2PTransport
from dissononce.dh.x25519.public import PublicKey

def main():
    parser = argparse.ArgumentParser(prog="noise_xk")
    parser.add_argument(
        "--listen-port", type=int, required=True,
        help="Port on which to bind+listen"
    )
    parser.add_argument(
        "--peer-host", type=str, default="127.0.0.1",
        help="Peer host to dial (default: 127.0.0.1)"
    )
    parser.add_argument(
        "--peer-port", type=int, required=True,
        help="Peer port to dial"
    )
    args = parser.parse_args()

    # 1. Generate static keypair and print our static public key
    kp = NoiseXKSession.generate_keypair()
    # kp.public is a PublicKey; .data holds raw bytes
    local_priv = kp.private     # carried implicitly in NoiseXKSession
    local_pub = kp.public
    print(f"[My static pubkey:] {local_pub.data.hex()}")

    # 2. Read peer pubkey from user input
    peer_pubkey = None
    while True:
        line = input(">>> ").strip()
        if line.startswith("peer_pubkey "):
            hexstr = line.split(None, 1)[1]
            try:
                raw = bytes.fromhex(hexstr)
                peer_pubkey = PublicKey(raw)  # wrap raw bytes in PublicKey
                break
            except ValueError:
                print("Invalid hex; please retry.")
        else:
            print("Use: peer_pubkey <hex>")

    # 3. Establish P2P connection (race listen vs. dial)
    transport = P2PTransport(
        listen_port=args.listen_port,
        peer_host=args.peer_host,
        peer_port=args.peer_port
    )
    print(
        f"Racing connect/listen on ports "
        f"{args.listen_port} ⇆ {args.peer_host}:{args.peer_port}…"
    )
    sock, initiator = transport.connect()
    print(f"Connected (initiator={initiator}); performing handshake…")

    # 4. Perform Noise XK handshake
    session = NoiseXKSession(kp, peer_pubkey)
    session.handshake(sock, initiator)
    print("Handshake complete! You can now type messages.")

    # 5. Reader thread for incoming messages
    def reader():
        while True:
            try:
                pt = session.receive(sock)
                print(f"\n< {pt.decode()}")
            except Exception as e:
                print(f"\n[Receive error ({type(e).__name__}): {e!r}]")
                break

    thread = threading.Thread(target=reader, daemon=True)
    thread.start()

    # 6. Main loop: send user input
    try:
        for line in sys.stdin:
            text = line.rstrip("\n")
            if not text:
                continue
            session.send(sock, text.encode())
    except KeyboardInterrupt:
        pass
    finally:
        sock.close()

if __name__ == "__main__":
    main()
