#!/usr/bin/env python3
import os
import sys
import time
import socket
import struct
import argparse
import random
import tempfile
import subprocess
from tqdm import tqdm  # For progress bar

# Default configuration values
DEFAULT_MULTICAST_ADDR = "239.192.55.1"
DEFAULT_PORT = 1681
DEFAULT_ZONE_INFO = "100000000000000000000000"  # Zone 1
DEFAULT_SEND_DUPLICATES = True
DEFAULT_CHUNK_SIZE = 900
DEFAULT_TTL = 2
DEFAULT_HEADER = "4d454c"  # MEL header
DEFAULT_COMMAND = "070301"  # Stream command

# Audio quality presets
HIGH_QUALITY = {
    "bitrate": "64k",
    "sample_rate": 48000,
    "channels": 1,  # mono
    "sample_format": "s32",  # 16-bit samples (default for MP3)
    "description": "Ass but no Durchfall"
    # Note: Original software appears to use LAME 3.99.5
}

LOW_QUALITY = {
    "bitrate": "64k",
    "sample_rate": 32000,
    "channels": 1,  # mono
    "sample_format": "s16",  # 16-bit samples (default for MP3)
    "description": "Ass with durchfalls"
}

def compute_psa_checksum(data: bytes) -> bytes:
    """Compute PSA checksum for packet data."""
    var_e = 0x0000  # Starting seed value
    for i in range(len(data)):
        var_e ^= (data[i] + i) & 0xFFFF
    return var_e.to_bytes(2, 'big')  # 2-byte checksum, big-endian

def create_mp3_packet(sequence_num, zone_info, stream_id, mp3_chunk):
    """Create a PSA packet containing MP3 data."""
    # Part 1: Header (static "MEL")
    header = bytes.fromhex("4d454c")
    
    # Prepare all the components that will go into the payload
    # Start marker (0100)
    start_marker = bytes.fromhex("0100")
    
    # Sequence number (2 bytes, Little Endian with FF padding)
    seq = sequence_num.to_bytes(1, 'little') + b'\xff'
    
    # Command (static 070301)
    command = bytes.fromhex(DEFAULT_COMMAND)
    
    # Zone info
    zones = bytes.fromhex(zone_info)
    
    # Stream ID (4 bytes)
    stream_id_bytes = stream_id.to_bytes(4, 'little')
    
    # Constant data (appears in all original packets)
    constant_data = bytes.fromhex("05080503e8")
    
    # Build the payload (everything that comes after start_marker)
    payload = seq + command + zones + stream_id_bytes + constant_data + mp3_chunk
    
    # Calculate the length for the header
    # Length is everything AFTER the length field: start_marker + payload
    # But NOT including checksum (which is calculated last)
    length_value = len(start_marker + payload) + 7
    
    # Insert the length as a 2-byte value (big endian)
    length_bytes = length_value.to_bytes(2, 'big')
    
    # Assemble the packet without checksum
    packet_data = header + length_bytes + start_marker + payload
    
    # Calculate and append checksum
    checksum = compute_psa_checksum(packet_data)
    
    # Debug the packet length
    final_packet = packet_data + checksum
    print(f"Packet length: {len(final_packet)} bytes, Length field: {length_value} (0x{length_value:04x})")
    
    return final_packet

def send_multicast_packet(sock, packet, multicast_addr, port, send_duplicates=True):
    """Send a packet to the multicast address."""
    try:
        # Send the packet
        sock.sendto(packet, (multicast_addr, port))
        
        # Send a duplicate (common in multicast to improve reliability)
        if send_duplicates:
            sock.sendto(packet, (multicast_addr, port))
        
        return True
    except Exception as e:
        print(f"Error sending packet: {e}")
        return False

def setup_multicast_socket(ttl=DEFAULT_TTL):
    """Set up a socket for sending multicast UDP."""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, ttl)
    return sock

def check_dependencies():
    """Check if required dependencies are installed."""
    try:
        subprocess.run(['ffmpeg', '-version'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except FileNotFoundError:
        print("Error: ffmpeg is required but not found.")
        print("Please install ffmpeg and make sure it's available in your PATH.")
        sys.exit(1)

def get_audio_info(mp3_file):
    """Get audio file information using ffprobe."""
    try:
        # Get audio stream information
        cmd = [
            'ffprobe', '-v', 'quiet', '-print_format', 'json',
            '-show_streams', '-select_streams', 'a:0', mp3_file
        ]
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        if result.returncode != 0:
            print(f"Error analyzing audio file: {result.stderr}")
            return None
        
        import json
        info = json.loads(result.stdout)
        
        if 'streams' not in info or len(info['streams']) == 0:
            print("No audio stream found in the file.")
            return None
        
        stream = info['streams'][0]
        
        # Extract relevant information
        bit_rate = int(stream.get('bit_rate', '0')) // 1000 if 'bit_rate' in stream else None
        sample_rate = int(stream.get('sample_rate', '0'))
        channels = int(stream.get('channels', '0'))
        codec_name = stream.get('codec_name', 'unknown')
        
        return {
            'bit_rate': bit_rate,
            'sample_rate': sample_rate,
            'channels': channels,
            'codec_name': codec_name
        }
    
    except Exception as e:
        print(f"Error getting audio information: {e}")
        return None

def transcode_mp3(input_file, quality_preset, include_metadata=False):
    """Transcode MP3 file to match required specifications."""
    print(f"Transcoding audio to {quality_preset['description']}...")
    
    # Create temporary output file
    fd, temp_output = tempfile.mkstemp(suffix='.mp3')
    os.close(fd)
    
    try:
        # Base command
        cmd = [
            'ffmpeg',
            '-y',  # Overwrite output file
            '-i', input_file,
            '-codec:a', 'libmp3lame',  # Force LAME encoder
            '-ac', str(quality_preset['channels']),
            '-ar', str(quality_preset['sample_rate']),
            '-b:a', quality_preset['bitrate'],
            '-sample_fmt', quality_preset['sample_format'],  # Use configurable sample format
        ]
        
        # Add options to minimize metadata if requested
        if not include_metadata:
            cmd.extend(['-metadata', 'title=', '-metadata', 'artist=', 
                       '-metadata', 'album=', '-metadata', 'comment=',
                       '-map_metadata', '-1'])  # Strip all metadata
        
        cmd.append(temp_output)
        
        result = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
        
        if result.returncode != 0:
            print(f"Error transcoding audio: {result.stderr}")
            os.remove(temp_output)
            return None
        
        return temp_output
    
    except Exception as e:
        print(f"Error during transcoding: {e}")
        if os.path.exists(temp_output):
            os.remove(temp_output)
        return None

def prepare_mp3_file(mp3_file, quality, include_metadata=False):
    """Ensure the MP3 file meets the required specifications."""
    # Check if quality is valid
    quality_preset = HIGH_QUALITY if quality == "high" else LOW_QUALITY
    
    # Get audio file information
    info = get_audio_info(mp3_file)
    if not info:
        return None
    
    print(f"Audio file details: {info['codec_name']}, {info['bit_rate']}kbps, "
          f"{info['sample_rate']}Hz, {info['channels']} channel(s)")
    
    # Check if transcoding is needed
    needs_transcode = False
    
    if info['codec_name'].lower() != 'mp3':
        print("Non-MP3 format detected. Transcoding required.")
        needs_transcode = True
    elif info['bit_rate'] is None or abs(info['bit_rate'] - int(quality_preset['bitrate'][:-1])) > 5:
        print(f"Bitrate mismatch: {info['bit_rate']}kbps vs {quality_preset['bitrate']}. Transcoding required.")
        needs_transcode = True
    elif info['sample_rate'] != quality_preset['sample_rate']:
        print(f"Sample rate mismatch: {info['sample_rate']}Hz vs {quality_preset['sample_rate']}Hz. Transcoding required.")
        needs_transcode = True
    elif info['channels'] != quality_preset['channels']:
        print(f"Channel count mismatch: {info['channels']} vs {quality_preset['channels']} (mono). Transcoding required.")
        needs_transcode = True
    
    # If transcoding is needed, do it
    if needs_transcode:
        return transcode_mp3(mp3_file, quality_preset, include_metadata)
    else:
        print("Audio file already meets required specifications.")
        return mp3_file

def calculate_delay_ms(chunk_size, quality):
    """Calculate the appropriate delay between chunks for real-time streaming."""
    # Calculate bytes_per_second dynamically from the bitrate
    if quality == "high":
        # 128kbps = 16,000 bytes/second
        bytes_per_second = 16000
    else:
        # 40kbps = 5,000 bytes/second
        bytes_per_second = 5000
    
    # Calculate how long this chunk would play for in real-time
    delay_ms = (chunk_size / bytes_per_second) * 1000
    return delay_ms

def stream_mp3_to_psa(mp3_file, multicast_addr, port, zone_info, send_duplicates,
                      chunk_size, quality, include_metadata=False):
    """Stream an MP3 file to PSA system."""
    try:
        # Check dependencies
        check_dependencies()
        
        # Prepare MP3 file (transcode if needed)
        prepared_file = prepare_mp3_file(mp3_file, quality, include_metadata)
        if not prepared_file:
            print("Failed to prepare MP3 file.")
            return
        
        using_temp_file = prepared_file != mp3_file
        
        try:
            # Generate random stream ID for this transmission
            stream_id = random.randint(0, 0xFFFFFFFF)  # Changed to 32-bit to accommodate 4 bytes
            
            # Open the MP3 file
            with open(prepared_file, 'rb') as f:
                mp3_data = f.read()
            
            # Calculate number of chunks
            chunks = [mp3_data[i:i+chunk_size] for i in range(0, len(mp3_data), chunk_size)]
            total_chunks = len(chunks)
            
            # Calculate the appropriate delay for real-time streaming
            delay_ms = calculate_delay_ms(chunk_size, quality)
            
            print(f"Streaming {os.path.basename(mp3_file)} ({len(mp3_data)/1024:.2f} KB)")
            print(f"Split into {total_chunks} packets of {chunk_size} bytes each")
            print(f"Sending to multicast {multicast_addr}:{port}")
            print(f"Stream ID: {stream_id:08x}, Zone info: {zone_info}")
            print(f"Duplicate packets: {'Yes' if send_duplicates else 'No'}")
            print(f"Quality preset: {HIGH_QUALITY['description'] if quality == 'high' else LOW_QUALITY['description']}")
            print(f"Including metadata: {'Yes' if include_metadata else 'No'}")
            print(f"Real-time streaming delay: {delay_ms:.2f}ms per packet")
            
            # Setup multicast socket
            sock = setup_multicast_socket()
            
            # Process and send each chunk
            sequence_num = 0
            with tqdm(total=total_chunks, desc="Streaming progress") as pbar:
                try:
                    for i, chunk in enumerate(chunks):
                        # Create packet with current sequence number
                        packet = create_mp3_packet(sequence_num, zone_info, stream_id, chunk)
                        
                        # Send the packet
                        success = send_multicast_packet(sock, packet, multicast_addr, port, send_duplicates)
                        
                        # Increment sequence num
                        sequence_num = (sequence_num + 1) % 256
                        
                        # Update progress bar
                        pbar.update(1)
                        
                        # Delay to match real-time playback
                        if i < total_chunks - 1:  # No need to wait after the last chunk
                            time.sleep(delay_ms / 1000.0)
                except Exception as e:
                    print(f"Error during streaming: {e}")
            
            print("\nStreaming completed successfully!")
            
        finally:
            # Clean up temporary file if created
            if using_temp_file and os.path.exists(prepared_file):
                os.remove(prepared_file)
                
    except FileNotFoundError:
        print(f"Error: MP3 file {mp3_file} not found")
        sys.exit(1)
    except Exception as e:
        print(f"Error during streaming: {e}")
        sys.exit(1)

def main():
    # Set up command line arguments
    parser = argparse.ArgumentParser(description="Stream MP3 to Bodet PSA System")
    parser.add_argument("mp3_file", help="Path to MP3 file to stream")
    parser.add_argument("-a", "--addr", default=DEFAULT_MULTICAST_ADDR, 
                        help=f"Multicast address (default: {DEFAULT_MULTICAST_ADDR})")
    parser.add_argument("-p", "--port", type=int, default=DEFAULT_PORT,
                        help=f"UDP port (default: {DEFAULT_PORT})")
    parser.add_argument("-z", "--zone", default=DEFAULT_ZONE_INFO, 
                        help=f"Hex zone info (default: Zone 1)")
    parser.add_argument("-n", "--no-duplicates", action="store_true",
                        help="Don't send duplicate packets (default: send duplicates)")
    parser.add_argument("-c", "--chunk-size", type=int, default=DEFAULT_CHUNK_SIZE,
                        help=f"MP3 chunk size in bytes (default: {DEFAULT_CHUNK_SIZE})")
    parser.add_argument("-q", "--quality", choices=["high", "low"], default="high",
                        help="Audio quality preset: high (128kbps, 48kHz) or low (40kbps, 32kHz) (default: high)")
    parser.add_argument("-m", "--include-metadata", action="store_true",
                        help="Include metadata in MP3 stream (default: no metadata)")
    parser.add_argument("-s", "--sample-format", choices=["s16", "s24", "s32"], default=None,
                        help="Sample format to use for transcoding (default: s16)")
    
    args = parser.parse_args()
    
    # Stream the MP3 file
    stream_mp3_to_psa(
        args.mp3_file,
        args.addr,
        args.port,
        args.zone,
        not args.no_duplicates,
        args.chunk_size,
        args.quality,
        args.include_metadata
    )

if __name__ == "__main__":
    main()