path: root/src/sigma-caster.rs

use std::env;
use std::fs::File;
use std::io::Read;
use std::net::UdpSocket;
use std::path::Path;
use std::process::Command;
use std::thread;
use std::time::{Duration, Instant};

use anyhow::{anyhow, Result};
use clap::Parser;
use indicatif::{ProgressBar, ProgressStyle};
use tempfile::NamedTempFile;

// Default stream ID that's known to work with Bodet systems (as hex string like Python)
const DEFAULT_STREAM_ID: u16 = 0x110c;

#[derive(Parser)]
#[command(name = "sigma-caster")]
#[command(about = "Stream MP3 files to Bodet PSA System")]
#[command(long_about = None)]
struct Args {
    /// Path to MP3 file to stream
    mp3_file: String,
    
    /// Multicast/unicast address
    #[arg(short = 'i', long = "ip", default_value = "239.192.55.1", help = "IP address to send to")]
    addr: String,
    
    /// UDP port
    #[arg(short, long, default_value = "1681", help = "UDP port")]
    port: u16,
    
    /// Zones (space-separated list like: 1 2 3)
    #[arg(short = 'z', long = "zones", value_delimiter = ' ', default_values = ["1"], help = "Zone numbers (space-separated)")]
    zones: Vec<u8>,
    
    /// Software volume control (1-8, where 8 is loudest)
    #[arg(short = 'v', long = "volume", default_value = "8", help = "Software volume (1-8, where 8 is loudest)")]
    volume: u8,
    
    /// Don't send duplicate packets
    #[arg(long, help = "Don't send duplicate packets")]
    no_duplicates: bool,
    
    /// MP3 chunk size in bytes
    #[arg(short, long, default_value = "1000", help = "MP3 chunk size in bytes")]
    chunk_size: usize,
    
    /// Audio quality preset
    #[arg(short, long, value_enum, default_value = "high", help = "Audio quality preset")]
    quality: Quality,
    
    /// Fixed stream ID (default: 0x110c)
    #[arg(long, help = "Fixed stream ID (default: 0x110c)")]
    stream_id: Option<u16>,
    
    /// Show debug information
    #[arg(long, help = "Show debug information")]
    debug: bool,
}

#[derive(clap::ValueEnum, Clone, Debug)]
enum Quality {
    High,
    Low,
}

struct QualityPreset {
    bitrate: &'static str,
    sample_rate: u32,
    description: &'static str,
}

impl Quality {
    fn preset(&self) -> QualityPreset {
        match self {
            Quality::High => QualityPreset {
                bitrate: "256k",
                sample_rate: 48000,
                description: "High Quality (256kbps, 48kHz, Mono)",
            },
            Quality::Low => QualityPreset {
                bitrate: "64k",
                sample_rate: 32000,
                description: "Low Quality (64kbps, 32kHz, Mono)",
            },
        }
    }
}

fn compute_psa_checksum(data: &[u8]) -> [u8; 2] {
    let mut var_e: u16 = 0x0000;
    for (i, &b) in data.iter().enumerate() {
        var_e ^= (b as u16).wrapping_add(i as u16) & 0xFFFF;
    }
    var_e.to_be_bytes()
}

fn create_mp3_packet(sequence_num: u8, zone_info: &[u8], stream_id: u16, mp3_chunk: &[u8]) -> Result<Vec<u8>> {
    let mut packet = Vec::new();
    
    // MEL header
    packet.extend_from_slice(b"MEL");
    
    // Build payload components
    let start_marker = [0x01, 0x00];
    let seq = [sequence_num, 0xff];
    let command = [0x07, 0x03];
    let stream_id_bytes = stream_id.to_le_bytes();
    let constant_data = [0x05, 0x08, 0x05, 0x03, 0xe8];
    
    // Calculate payload length - matches Python: start_marker + payload + 7
    let payload_len = seq.len() + command.len() + zone_info.len() + 
                     stream_id_bytes.len() + constant_data.len() + mp3_chunk.len();
    let total_len = start_marker.len() + payload_len + 7;
    
    // Length field (big endian)
    packet.extend_from_slice(&(total_len as u16).to_be_bytes());
    
    // Start marker
    packet.extend_from_slice(&start_marker);
    
    // Payload
    packet.extend_from_slice(&seq);
    packet.extend_from_slice(&command);
    packet.extend_from_slice(zone_info);
    packet.extend_from_slice(&stream_id_bytes);
    packet.extend_from_slice(&constant_data);
    packet.extend_from_slice(mp3_chunk);
    
    // Checksum
    let checksum = compute_psa_checksum(&packet);
    packet.extend_from_slice(&checksum);
    
    Ok(packet)
}

fn check_lame_available() -> Result<()> {
    Command::new("lame")
        .arg("--version")
        .output()
        .map_err(|_| anyhow!("LAME encoder not found. Please install LAME and ensure it's in PATH."))?;
    Ok(())
}

fn apply_volume_to_mp3(input_file: &str, output_file: &str, volume: u8, preset: &QualityPreset) -> Result<()> {
    // Convert volume (1-8) to LAME scale factor
    let volume_scale = (volume as f32) / 8.0;
    
    let mut cmd = Command::new("lame");
    cmd.arg("--quiet")
        .arg("-m").arg("m") // mono
        .arg("--resample").arg(&(preset.sample_rate / 1000).to_string())
        .arg("-b").arg(preset.bitrate.trim_end_matches('k'))
        .arg("--cbr")
        .arg("--noreplaygain");
    
    // Apply volume scaling if needed
    if volume < 8 {
        cmd.arg("--scale").arg(&format!("{:.3}", volume_scale));
    }
    
    cmd.arg("--tt").arg("") // empty title
        .arg("--tc").arg("") // empty comment
        .arg("--nohist")
        .arg(input_file)
        .arg(output_file);
    
    let output = cmd.output()?;
    
    if !output.status.success() {
        return Err(anyhow!("LAME transcoding with volume failed: {}", 
                          String::from_utf8_lossy(&output.stderr)));
    }
    
    Ok(())
}

fn transcode_mp3(input_file: &str, preset: &QualityPreset, volume: u8) -> Result<NamedTempFile> {
    println!("Transcoding audio to {} (Volume: {}/8)...", preset.description, volume);
    
    let temp_file = NamedTempFile::new()?;
    
    apply_volume_to_mp3(input_file, temp_file.path().to_str().unwrap(), volume, preset)?;
    
    Ok(temp_file)
}

fn chunk_data(data: &[u8], chunk_size: usize) -> Vec<&[u8]> {
    data.chunks(chunk_size).collect()
}

fn calculate_delay_ms(chunk_size: usize, bitrate_kbps: u32) -> u64 {
    let bytes_per_second = bitrate_kbps * 1000 / 8;
    (chunk_size as u64 * 1000) / bytes_per_second as u64
}

fn build_zone_bytes(zones: &[u8]) -> [u8; 13] {
    let mut zone_bytes = [0u8; 13]; // 13 bytes like Python (12 + extra 00)
    for &zone in zones {
        if zone == 0 || zone > 100 {
            continue;
        }
        let idx = (zone - 1) / 8;
        let bit = (zone - 1) % 8;
        zone_bytes[idx as usize] |= 1 << bit;
    }
    // Last byte stays 0x00 as per Python DEFAULT_ZONE_INFO pattern
    zone_bytes
}

fn main() -> Result<()> {
    // Show help if no arguments provided
    if env::args().len() == 1 {
        eprintln!("sigma-caster - Stream MP3 files to Bodet PSA System

Usage: sigma-caster [OPTIONS] <MP3_FILE>

Arguments:
  <MP3_FILE>  Path to MP3 file to stream

Options:
  -i, --ip <IP>              IP address to send to [default: 239.192.55.1]
  -p, --port <PORT>          UDP port [default: 1681]
  -z, --zones <ZONES>...     Zone numbers (space-separated) [default: 1]
  -v, --volume <VOLUME>      Software volume (1-8, where 8 is loudest) [default: 8]
      --no-duplicates        Don't send duplicate packets
  -c, --chunk-size <SIZE>    MP3 chunk size in bytes [default: 1000]
  -q, --quality <QUALITY>    Audio quality preset [default: high] [possible values: high, low]
      --stream-id <ID>       Fixed stream ID (default: 0x110c)
      --debug                Show debug information
  -h, --help                 Print help

Examples:
  sigma-caster audio.mp3
      # Stream audio.mp3 to zone 1 at full volume
  sigma-caster -i 172.16.20.109 -z 1 2 3 -v 6 audio.mp3
      # Stream to zones 1,2,3 at volume 6/8 on specific IP
  sigma-caster --debug -q low audio.mp3
      # Stream with low quality and show debug info");
        std::process::exit(1);
    }

    let args = Args::parse();
    
    // Validate volume range
    if !(1..=8).contains(&args.volume) {
        return Err(anyhow!("Volume must be between 1 and 8"));
    }
    
    // Check dependencies
    check_lame_available()?;
    
    // Validate input file
    if !Path::new(&args.mp3_file).exists() {
        return Err(anyhow!("MP3 file not found: {}", args.mp3_file));
    }
    
    // Build zone info from zone numbers
    let zone_bytes = build_zone_bytes(&args.zones);
    
    // Get quality preset
    let preset = args.quality.preset();
    
    // Transcode MP3 with volume control
    let temp_file = transcode_mp3(&args.mp3_file, &preset, args.volume)?;
    
    // Read transcoded MP3 data
    let mut mp3_data = Vec::new();
    File::open(temp_file.path())?.read_to_end(&mut mp3_data)?;
    
    // Use default stream ID if not specified
    let stream_id = args.stream_id.unwrap_or(DEFAULT_STREAM_ID);
    
    // Chunk the data
    let chunks = chunk_data(&mp3_data, args.chunk_size);
    let total_chunks = chunks.len();
    
    // Calculate timing
    let bitrate_kbps = preset.bitrate.trim_end_matches('k').parse::<u32>()?;
    let delay_ms = calculate_delay_ms(args.chunk_size, bitrate_kbps);
    
    println!("Streaming {} ({:.2} KB)", 
             Path::new(&args.mp3_file).file_name().unwrap().to_string_lossy(),
             mp3_data.len() as f64 / 1024.0);
    println!("Split into {} packets of {} bytes each", total_chunks, args.chunk_size);
    println!("Sending to {}:{}", args.addr, args.port);
    println!("Stream ID: {:04x}, Quality: {}", stream_id, preset.description);
    println!("Zones: {:?}, Volume: {}/8", args.zones, args.volume);
    println!("Duplicate packets: {}", if args.no_duplicates { "No" } else { "Yes" });
    println!("Delay per packet: {}ms", delay_ms);
    
    // Setup UDP socket
    let socket = UdpSocket::bind("0.0.0.0:0")?;
    let target_addr = format!("{}:{}", args.addr, args.port);
    
    // Progress bar
    let progress = ProgressBar::new(total_chunks as u64);
    progress.set_style(ProgressStyle::default_bar()
        .template("[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}")
        .unwrap());
    
    // Stream chunks
    let start_time = Instant::now();
    let mut sequence_num = 0u8;
    
    for (i, chunk) in chunks.iter().enumerate() {
        // Create packet
        let packet = create_mp3_packet(sequence_num, &zone_bytes, stream_id, chunk)?;
        
        if args.debug && i == 0 {
            let hex_str = packet.iter().map(|b| format!("{:02x}", b)).collect::<String>();
            println!("DEBUG: First packet hex: {}", hex_str);
        }
        
        // Send packet
        socket.send_to(&packet, &target_addr)?;
        
        // Send duplicate if requested
        if !args.no_duplicates {
            socket.send_to(&packet, &target_addr)?;
        }
        
        sequence_num = sequence_num.wrapping_add(1);
        progress.inc(1);
        
        // Wait for next packet (except for last one)
        if i < total_chunks - 1 {
            thread::sleep(Duration::from_millis(delay_ms));
        }
    }
    
    progress.finish_with_message("Streaming completed!");
    println!("Total time: {:.2}s", start_time.elapsed().as_secs_f64());
    
    Ok(())
}