# main.py
#!/usr/bin/env python3
import argparse
import socket
import numpy as np
import channel_helper as ch
from encoder import make_codebook, encode_message
from decoder import decode_blocks, count_errors
from channel import channel as external_channel


def send_and_recv(x: np.ndarray, host: str, port: int) -> np.ndarray:
    """Send samples x to server and receive output via TCP"""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.connect((host, port))
        header = b'0' + b'dUV'
        ch.send_msg(sock, header, x)
        _, Y = ch.recv_msg(sock)
    return Y


def main():
    p = argparse.ArgumentParser(description="PDC Tx/Rx local or server")
    p.add_argument("--message", required=True, help="40-character message to send")
    p.add_argument("--srv_hostname", help="Server hostname")
    p.add_argument("--srv_port", type=int, help="Server port")
    p.add_argument("--local", action='store_true', help="Use local channel simulation")
    args = p.parse_args()

    msg = args.message
    if len(msg) != 40:
        raise ValueError("Message must be exactly 40 characters.")
    num_blocks = len(msg)
    C = make_codebook(r=5, num_blocks=num_blocks, margin=0.95)
    x = encode_message(msg, C)
    print(f"→ Total samples = {x.size}, total energy = {np.sum(x*x):.1f}")

    if args.local:
        print("-- Local simulation mode --")
        Y = external_channel(x)
    else:
        if not args.srv_hostname or not args.srv_port:
            raise ValueError("Must specify --srv_hostname and --srv_port unless --local")
        Y = send_and_recv(x, args.srv_hostname, args.srv_port)

    msg_hat = decode_blocks(Y, C)
    print(f"↓ Decoded message: {msg_hat}")

    errors = count_errors(msg, msg_hat)
    print(f"Errors: {len(errors)} / {len(msg)} characters differ")
    if errors:
        for i, o, e in errors:
            print(f"  Pos {i}: sent '{o}' but got '{e}'")
    else:
        print("✔️  No decoding errors!")

if __name__ == "__main__":
    main()