Concatenate AEAs

aea-concat.py

#!/usr/bin/env python3

import os
import sys
import struct
import argparse
from pathlib import Path

AEA_MAGIC = bytes([0x00, 0x08, 0x00, 0x00])
AEA_HEADER_SIZE = 2048
AEA_TITLE_OFFSET = 4
AEA_TITLE_SIZE = 256
AEA_FRAME_COUNT_OFFSET = 260
AEA_CHANNEL_COUNT_OFFSET = 264
SOUND_UNIT_SIZE = 212

class AeaFile:
    """AEA file format handler for ATRAC1 audio files"""
    
    @staticmethod
    def create_header(title="", frame_count=0, channel_count=1):
        """Creates an AEA file header with the specified metadata"""
        header = bytearray(AEA_HEADER_SIZE)
        
        # Magic number
        header[0:4] = AEA_MAGIC
        
        # Title (null-terminated string)
        title_bytes = title.encode('utf-8')
        title_length = min(len(title_bytes), AEA_TITLE_SIZE - 1)
        header[AEA_TITLE_OFFSET:AEA_TITLE_OFFSET + title_length] = title_bytes[:title_length]
        
        # Frame count (little endian)
        struct.pack_into('<I', header, AEA_FRAME_COUNT_OFFSET, frame_count)
        
        # Channel count
        header[AEA_CHANNEL_COUNT_OFFSET] = channel_count
        
        return bytes(header)
    
    @staticmethod
    def parse_header(header):
        """Parses an AEA file header to extract metadata"""
        if len(header) != AEA_HEADER_SIZE:
            raise ValueError(f"Header must be {AEA_HEADER_SIZE} bytes")
        
        # Check magic number
        if header[0:4] != AEA_MAGIC:
            raise ValueError("Invalid AEA file")
        
        # Extract title
        title_end = header.find(0, AEA_TITLE_OFFSET)
        if title_end == -1:
            title_length = AEA_TITLE_SIZE
        else:
            title_length = title_end - AEA_TITLE_OFFSET
        
        title = header[AEA_TITLE_OFFSET:AEA_TITLE_OFFSET + title_length].decode('utf-8', errors='ignore')
        
        # Extract metadata (little endian)
        frame_count = struct.unpack_from('<I', header, AEA_FRAME_COUNT_OFFSET)[0]
        channel_count = header[AEA_CHANNEL_COUNT_OFFSET]
        
        return {
            'title': title,
            'frame_count': frame_count,
            'channel_count': channel_count
        }

class AeaReader:
    """AEA file reader with iteration support for streaming frame processing"""
    
    def __init__(self, file_path):
        self.file_path = file_path
        self.metadata = None
    
    def load_metadata(self):
        """Load and parse AEA file metadata from header"""
        with open(self.file_path, 'rb') as f:
            header = f.read(AEA_HEADER_SIZE)
            self.metadata = AeaFile.parse_header(header)
    
    def iter_frames(self):
        """Generator for streaming frame data"""
        with open(self.file_path, 'rb') as f:
            f.seek(AEA_HEADER_SIZE)
            while True:
                frame_data = f.read(SOUND_UNIT_SIZE)
                if len(frame_data) != SOUND_UNIT_SIZE:
                    break
                yield frame_data

def concatenate_aea_files(input_files, output_file, title="concatenated"):
    """Main concatenation function"""
    print(f"Concatenating {len(input_files)} AEA files...")
    
    readers = []
    common_channel_count = None
    total_frames = 0
    
    # Load metadata and validate channel consistency
    for file_path in input_files:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
        
        reader = AeaReader(file_path)
        reader.load_metadata()
        readers.append(reader)
        
        if common_channel_count is None:
            common_channel_count = reader.metadata['channel_count']
        elif reader.metadata['channel_count'] != common_channel_count:
            raise ValueError(f"Channel count mismatch: {file_path} has {reader.metadata['channel_count']} channels, expected {common_channel_count}")
        
        total_frames += reader.metadata['frame_count']
        print(f"  {Path(file_path).name}: {reader.metadata['frame_count']} frames, {reader.metadata['channel_count']} channels")
    
    print(f"Total frames: {total_frames}, Channels: {common_channel_count}")
    
    # Create new header
    new_header = AeaFile.create_header(title, total_frames, common_channel_count)
    
    # Create output file
    with open(output_file, 'wb') as output_f:
        # Write header
        output_f.write(new_header)
        
        # Concatenate all frames
        frame_count = 0
        for reader in readers:
            if frame_count > 0:
                print()
            print(f"Processing {Path(reader.file_path).name}...")
            for frame_data in reader.iter_frames():
                output_f.write(frame_data)
                frame_count += 1
                if frame_count % 1000 == 0:
                    print(f"\r  Frames processed: {frame_count}/{total_frames}", end='', flush=True)
        
        print(f"\r  Frames processed: {frame_count}/{total_frames}")
    
    print(f"Successfully created {output_file}")

def main():
    parser = argparse.ArgumentParser(
        description="Concatenates multiple AEA files into a single AEA file.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Requirements:
  - All input files must have the same channel count (mono or stereo)
  - At least 2 input files and 1 output file must be specified
  - Input files must be valid AEA format files

Examples:
  python aea-concat.py song1.aea song2.aea combined.aea
  python aea-concat.py --title "Album Mix" track*.aea album.aea
        """
    )
    
    parser.add_argument('files', nargs='+', help='Input AEA files followed by output AEA file')
    parser.add_argument('--title', default='concatenated', help='Set the title for the output file (default: "concatenated")')
    
    args = parser.parse_args()
    
    if len(args.files) < 3:
        parser.error("At least 2 input files and 1 output file required")
    
    output_file = args.files[-1]
    input_files = args.files[:-1]
    
    try:
        concatenate_aea_files(input_files, output_file, args.title)
    except Exception as error:
        print(f"Error: {error}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()