"""Short-text simulation: UDP between UE tunnel IPs, with recipient receipts."""
import argparse
import json
import math
import re
import socket
import time
from collections import OrderedDict, deque
from pathlib import Path
from apps.phones import SMS_PHONES

PHONES = SMS_PHONES
PORT = 9002

def validate_envelope(message, receiver, peer_ip):
    if not isinstance(message, dict):
        raise ValueError("Invalid SMS envelope")
    sender = message.get("sender")
    if sender not in PHONES or PHONES[sender] != peer_ip:
        raise ValueError("Sender does not match the source tunnel IP")
    if message.get("recipient") != receiver or sender == receiver:
        raise ValueError("Incorrect recipient")
    text = message.get("text")
    if not isinstance(text, str) or not text.strip() or len(text) > 160:
        raise ValueError("SMS must contain 1 to 160 characters")
    if not isinstance(message.get("id"), str) or not re.fullmatch(r"[0-9a-f]{32}", message["id"]):
        raise ValueError("Invalid message ID")
    sent_at = message.get("sent_at")
    if not isinstance(sent_at, (int, float)) or not math.isfinite(sent_at) or sent_at <= 0:
        raise ValueError("Invalid timestamp")
    if message.get("kind", "sms") not in ("sms", "probe"):
        raise ValueError("Invalid message kind")
    envelope = {key: message[key] for key in ("id", "sender", "recipient", "text", "sent_at")}
    if message.get("kind") == "probe":
        envelope["kind"] = "probe"
    return envelope

def load_receipts(inbox):
    receipts = OrderedDict()
    if inbox.exists():
        with inbox.open(encoding="utf-8") as file:
            for line in deque(file, maxlen=200):
                try:
                    receipt = json.loads(line)
                    if receipt.get("status") == "delivered":
                        receipts[receipt["id"]] = receipt
                except (ValueError, KeyError, TypeError):
                    continue
    return receipts

def receipt_for(message, inbox, receipts):
    # Readiness probes exercise the same bidirectional UE path without inbox entries.
    if message.get("kind") == "probe":
        return {**message, "status": "ready"}
    receipt = receipts.get(message["id"])
    if receipt is None:
        receipt = {**message, "status": "delivered", "delivered_at": time.time()}
        # Acknowledge only after this UE has saved the message.
        with inbox.open("a", encoding="utf-8") as file:
            file.write(json.dumps(receipt, ensure_ascii=False) + "\n")
        receipts[message["id"]] = receipt
        if len(receipts) > 200:
            receipts.popitem(last=False)
        print(json.dumps(receipt, ensure_ascii=False), flush=True)
    elif any(receipt.get(key) != message[key] for key in ("sender", "recipient", "text", "sent_at")):
        raise ValueError("Message ID already used by a different envelope")
    return receipt

def receive(receiver, inbox=Path("/tmp/sms-inbox.jsonl")):
    from apps.network import bound_socket, tunnel_ip
    receipts = load_receipts(inbox)
    while True:
        try:
            with bound_socket(socket.SOCK_DGRAM, timeout=1, port=PORT) as sock:
                address = sock.getsockname()[0]
                interface_index = socket.if_nametoindex("uesimtun0")
                if address != PHONES[receiver]:
                    raise RuntimeError("Receiver tunnel does not match its subscriber address")
                while tunnel_ip(timeout=0) == address and socket.if_nametoindex("uesimtun0") == interface_index:
                    try:
                        data, peer = sock.recvfrom(8192)
                    except socket.timeout:
                        continue
                    try:
                        message = validate_envelope(json.loads(data), receiver, peer[0])
                        receipt = receipt_for(message, inbox, receipts)
                    except (ValueError, TypeError, UnicodeDecodeError):
                        continue
                    sock.sendto(json.dumps(receipt, ensure_ascii=False).encode(), peer)
        except (OSError, RuntimeError) as error:
            print(f"SMS receiver waiting for tunnel: {error}", flush=True)
            time.sleep(1)

def send(message):
    from apps.network import bound_socket
    validate_envelope(message, message["recipient"], PHONES[message["sender"]])
    with bound_socket(socket.SOCK_DGRAM, timeout=6) as sock:
        if sock.getsockname()[0] != PHONES[message["sender"]]:
            raise RuntimeError("Sender tunnel does not match its subscriber address")
        destination = (PHONES[message["recipient"]], PORT)
        # UDP connect also restricts received receipts to this destination.
        sock.connect(destination)
        started = time.perf_counter()
        payload = json.dumps(message, ensure_ascii=False).encode()
        timeout = 2 if message.get("kind") == "probe" else 6
        deadline = time.monotonic() + timeout
        receipt = None
        while time.monotonic() < deadline:
            sock.settimeout(min(1, max(.01, deadline - time.monotonic())))
            try:
                sock.send(payload)
                receipt = json.loads(sock.recv(8192))
                break
            except socket.timeout:
                continue
            except ConnectionRefusedError:
                time.sleep(.25)
        if receipt is None:
            raise RuntimeError(f"No delivery receipt within {timeout} seconds")
        expected = "ready" if message.get("kind") == "probe" else "delivered"
        if any(receipt.get(key) != message[key] for key in ("id", "sender", "recipient", "text")) or receipt.get("status") != expected:
            raise RuntimeError("Invalid delivery receipt")
        receipt["rtt_ms"] = round((time.perf_counter() - started) * 1000, 2)
        receipt["sender_ip"] = sock.getsockname()[0]
        receipt["recipient_ip"] = destination[0]
        print(json.dumps(receipt, ensure_ascii=False), flush=True)
        return receipt

def main():
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--listen", choices=PHONES)
    group.add_argument("--send", help="JSON message envelope")
    args = parser.parse_args()
    if args.listen:
        receive(args.listen)
    else:
        try:
            send(json.loads(args.send))
        except (OSError, ValueError, RuntimeError) as error:
            raise SystemExit(f"SMS delivery failed: {error}. Check phone sessions and UPF connectivity.")

if __name__ == "__main__":
    main()