#!/usr/bin/env python3
"""
packets.py - Jazz Packet UDP Emission/Reception
================================================

Jazz Packets are UDP broadcast messages carrying symbolic payloads.
Port 57321 is the canonical Jazz port (non-root, high-port).

Consolidated from:
- receipts/.../tools/jazz_packet_emitter__v4.py
- jazz_packet_listener concepts

WO#NET_JAZZ_00
"""

from __future__ import annotations
import socket
import random
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Optional, Callable, Generator, Dict, Any

# Canonical Jazz UDP port (high port, no root needed)
JAZZ_PORT = 57321

# Jazz packet tag
JAZZ_TAG = "::JAZZ::"


@dataclass
class JazzPacket:
    """A Jazz network packet."""
    timestamp: str
    tag: str
    payload: str
    source_ip: Optional[str] = None
    source_port: Optional[int] = None

    def encode(self) -> bytes:
        """Encode packet for UDP transmission."""
        msg = f"[{self.timestamp}] {JAZZ_TAG}{self.tag}::{self.payload}"
        return msg.encode('utf-8')

    @classmethod
    def decode(cls, data: bytes, addr: tuple = None) -> "JazzPacket":
        """Decode packet from UDP data."""
        try:
            msg = data.decode('utf-8', errors='ignore')
            # Parse: [timestamp] ::JAZZ::tag::payload
            if JAZZ_TAG in msg:
                parts = msg.split(JAZZ_TAG, 1)
                timestamp_part = parts[0].strip()
                rest = parts[1] if len(parts) > 1 else ""

                # Extract timestamp from [timestamp]
                if timestamp_part.startswith("[") and "]" in timestamp_part:
                    timestamp = timestamp_part[1:timestamp_part.index("]")]
                else:
                    timestamp = datetime.now().isoformat()

                # Split tag::payload
                if "::" in rest:
                    tag, payload = rest.split("::", 1)
                else:
                    tag = "UNKNOWN"
                    payload = rest

                return cls(
                    timestamp=timestamp,
                    tag=tag,
                    payload=payload,
                    source_ip=addr[0] if addr else None,
                    source_port=addr[1] if addr else None,
                )
        except Exception:
            pass

        # Fallback
        return cls(
            timestamp=datetime.now().isoformat(),
            tag="RAW",
            payload=data.decode('utf-8', errors='ignore'),
            source_ip=addr[0] if addr else None,
            source_port=addr[1] if addr else None,
        )


def emit_jazz_packet(
    payload: str,
    tag: str = "HELIX",
    dest: str = "255.255.255.255",
    port: int = JAZZ_PORT,
) -> Dict[str, Any]:
    """
    Emit a Jazz packet via UDP broadcast.

    Args:
        payload: Message payload (max 512 chars recommended)
        tag: Source tag (e.g., "PETAL", "BLOOMCORE", "HELIX")
        dest: Destination IP (default: broadcast)
        port: Destination port (default: 57321)

    Returns:
        Dict with emission result
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)

    try:
        # Truncate payload if needed
        payload_trimmed = payload[-512:] if len(payload) > 512 else payload

        packet = JazzPacket(
            timestamp=datetime.now().isoformat(),
            tag=tag,
            payload=payload_trimmed,
        )

        sock.sendto(packet.encode(), (dest, port))

        return {
            "sent": True,
            "dest": dest,
            "port": port,
            "tag": tag,
            "payload_size": len(payload_trimmed),
            "timestamp": packet.timestamp,
        }

    except Exception as e:
        return {
            "sent": False,
            "error": str(e),
            "timestamp": datetime.now().isoformat(),
        }
    finally:
        sock.close()


def listen_jazz_packets(
    port: int = JAZZ_PORT,
    timeout: Optional[float] = None,
    max_packets: Optional[int] = None,
    on_packet: Optional[Callable[[JazzPacket], None]] = None,
) -> Generator[JazzPacket, None, None]:
    """
    Listen for Jazz packets on UDP.

    Args:
        port: Port to listen on (default: 57321)
        timeout: Socket timeout in seconds (None = blocking)
        max_packets: Max packets to receive (None = infinite)
        on_packet: Optional callback for each packet

    Yields:
        JazzPacket objects
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("", port))

    if timeout:
        sock.settimeout(timeout)

    count = 0
    try:
        while max_packets is None or count < max_packets:
            try:
                data, addr = sock.recvfrom(2048)
                packet = JazzPacket.decode(data, addr)

                if on_packet:
                    on_packet(packet)

                yield packet
                count += 1

            except socket.timeout:
                continue

    finally:
        sock.close()


def emit_from_file(
    source_path: Path,
    tag: str = "PETAL",
    dest: str = "255.255.255.255",
    port: int = JAZZ_PORT,
) -> Dict[str, Any]:
    """
    Emit a Jazz packet with payload from a file (random line).

    Args:
        source_path: Path to source file
        tag: Source tag
        dest: Destination IP
        port: Destination port

    Returns:
        Dict with emission result
    """
    if not source_path.exists():
        return {
            "sent": False,
            "error": f"Source file not found: {source_path}",
            "timestamp": datetime.now().isoformat(),
        }

    lines = source_path.read_text(errors='ignore').splitlines()
    if not lines:
        return {
            "sent": False,
            "error": "Source file is empty",
            "timestamp": datetime.now().isoformat(),
        }

    payload = random.choice(lines)
    return emit_jazz_packet(payload, tag=tag, dest=dest, port=port)


# ============================================================================
# CLI
# ============================================================================

def _cli():
    """Jazz Packets CLI."""
    import argparse

    parser = argparse.ArgumentParser(description="Jazz Packet Emitter/Listener")
    subparsers = parser.add_subparsers(dest="command")

    # Emit subcommand
    emit_parser = subparsers.add_parser("emit", help="Emit a Jazz packet")
    emit_parser.add_argument("payload", help="Payload to send")
    emit_parser.add_argument("--tag", default="HELIX", help="Source tag")
    emit_parser.add_argument("--port", type=int, default=JAZZ_PORT, help="Destination port")

    # Listen subcommand
    listen_parser = subparsers.add_parser("listen", help="Listen for Jazz packets")
    listen_parser.add_argument("--port", type=int, default=JAZZ_PORT, help="Port to listen on")
    listen_parser.add_argument("--timeout", type=float, help="Socket timeout")
    listen_parser.add_argument("--max", type=int, help="Max packets to receive")

    args = parser.parse_args()

    if args.command == "emit":
        result = emit_jazz_packet(args.payload, tag=args.tag, port=args.port)
        print(f"[JAZZ] Emit: {result}")

    elif args.command == "listen":
        print(f"[JAZZ] Listening on port {args.port}...")
        for packet in listen_jazz_packets(port=args.port, timeout=args.timeout, max_packets=args.max):
            print(f"[JAZZ] {packet.source_ip} :: {packet.tag} :: {packet.payload[:64]}...")

    else:
        parser.print_help()


if __name__ == "__main__":
    _cli()
