/*
 * Copyright (c) 2006 Ariff Abdullah <ariff@FreeBSD.org>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 */
#include <sys/soundcard.h>
#include <sys/ioctl.h>
#include <fcntl.h>
#include <inttypes.h>
#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sysexits.h>
#include <unistd.h>
#include <signal.h>

#define WAVE_HDR_SIZE 44
#define RIFF_MAGIC    "RIFF"
#define WAVE_MAGIC    "WAVE"
#define FMT_MAGIC     "fmt "
#define DATA_MAGIC    "data"

#define AFMT_8BIT (AFMT_MU_LAW|AFMT_A_LAW|AFMT_U8|AFMT_S8)
#define AFMT_16BIT (AFMT_S16_LE|AFMT_U16_LE|AFMT_S16_BE|AFMT_U16_BE)
#define AFMT_24BIT (AFMT_S24_LE|AFMT_U24_LE|AFMT_S24_BE|AFMT_U24_BE)
#define AFMT_32BIT (AFMT_S32_LE|AFMT_U32_LE|AFMT_S32_BE|AFMT_U32_BE)
#define AFMT_SIGNED (AFMT_S8|AFMT_S16_LE|AFMT_S16_BE|AFMT_S24_LE|AFMT_S24_BE|\
			AFMT_S32_LE|AFMT_S32_BE)
#define AFMT_BIGENDIAN (AFMT_S16_BE|AFMT_U16_BE|AFMT_S24_BE|AFMT_U24_BE|\
				AFMT_S32_BE|AFMT_U32_BE)

#define VGAIN	10

struct wave_header {
	char riff[4];
	uint32_t rifflen;
	char wave[4];
	char fmt[4];
	uint32_t flen;
	uint16_t tag;
	uint16_t channels;
	uint32_t rate;
	uint32_t byterate;
	uint16_t align;
	uint16_t bit;
	char data[4];
	uint32_t size;
};

static FILE *fp = NULL;
struct wave_header nhdr = {
	{'R', 'I', 'F', 'F'},
	0, {'W', 'A', 'V', 'E'}, {'f', 'm', 't', ' '},
	0, 0, 0, 0, 0, 0, 0, {'d', 'a', 't', 'a'},
	0
};
static struct wave_header *whdr = NULL;

static uint8_t *buf = NULL;

static void
usage(void)
{
        fprintf(stderr, 
		"pcmrec [options...]\n"
		"         -w             output wave header\n"
		"                        default: raw headerless\n"
		"         -B arg         IO blocksize\n"
		"                        default: 512\n"
		"         -c arg         1 (mono) 2 (stereo)\n"
		"                        Default: 2\n"
		"         -d arg         dsp device\n"
		"                        default: /dev/dsp\n"
		"         -f arg         output format\n"
		"                        default: s16le\n"
		"         -o arg         output file\n"
		"                        default: output to stdout\n"
		"         -r arg         sampling rate\n"
		"                        default: 44100hz\n"
		"         -v arg         apply volume gain (-%d to %d)\n"
		"                        default: 0\n\n"
		"Valid Format: u8 s8 mulaw alaw\n"
		"              s16le u16le s16be u16be\n"
		"              s24le u24le s24be u24be\n"
		"              s32le u32le s32be u32be\n", VGAIN, VGAIN);
        exit(EX_USAGE);
}

static sig_t
trapint(int sig)
{
	sig = sig;
	fprintf(stderr, "\nfinished..\n");
	if (fp && fp != stdout) {
		if (whdr) {
			whdr->size = ftell(fp) - WAVE_HDR_SIZE;
			whdr->rifflen = whdr->size - 8 + WAVE_HDR_SIZE;
			if (fseek(fp, 0, SEEK_SET) != 0) {
				perror("fseek()");
				fclose(fp);
				exit(EX_IOERR);
			}
			if (fwrite((uint8_t *)whdr, 1,
					WAVE_HDR_SIZE, fp) != WAVE_HDR_SIZE) {
				perror("fwrite()");
				fclose(fp);
				exit(EX_IOERR);
			}
		}
		fflush(NULL);
		fclose(fp);
	}
	if (buf) {
		free(buf);
		buf = NULL;
	}
	exit(EX_OK);
}

static void
record(const char*	sample_device, 
       uint32_t 	sample_rate, 
       uint32_t 	sample_channels, 
       uint32_t 	sample_encoding,
       uint32_t		bits,
       uint32_t		blksz,
       int		vgain,
       int		rawout,
       char 		*outfile,
       char		*fmtstr)
{
	int		fd, frag;
	uint32_t	bps;
	uint32_t	total = 0, argrate;
	int32_t		sample;
	int		i, sign, be, max;

	fd = open(sample_device, O_RDONLY);
	if (fd < 0) {
		perror(sample_device);
		exit(EX_OSFILE);
	}
	if (ioctl(fd, SNDCTL_DSP_SETFMT, &sample_encoding) < 0) {
		perror("SNDCTL_DSP_CHANNELS");
		exit(EX_DATAERR);
	}
	if (ioctl(fd, SNDCTL_DSP_CHANNELS, &sample_channels) < 0) {
		perror("SNDCTL_DSP_CHANNELS");
		exit(EX_DATAERR);
	}

	argrate = sample_rate;
	if (ioctl(fd, SNDCTL_DSP_SPEED, &argrate) < 0 || argrate != sample_rate) {
		if (argrate != sample_rate)
			warnx("SNDCTL_DSP_SPEED: Failed to set speed (%u -> %u)",
				sample_rate, argrate);
		else
			perror("SNDCTL_DSP_SPEED");
		exit(EX_DATAERR);
	}

	bps = (bits / 8) * sample_channels;
	blksz = bps * blksz * (sample_rate / 44100.0);
	if (blksz < bps)
		blksz = bps;
	else if (blksz > 131072)
		blksz = 131072;
	blksz -= blksz % bps;

	frag = 0;
	while (blksz >> frag)
		frag++;
	frag |= 0xffff0000;

	/*if (ioctl(fd, SNDCTL_DSP_SETFRAGMENT, &frag) < 0) {
		perror("SNDCTL_DSP_SETFRAGMENT");
		exit(EX_DATAERR);
	}*/

	if (outfile) {
		int ofd = open(outfile, O_WRONLY|O_CREAT|O_EXCL, 0644);
		if (ofd < 0) {
			perror("open()");
			exit(EX_IOERR);
		}
		fp = fdopen(ofd, "w");
		if (!fp) {
			perror("fdopen()");
			exit(EX_IOERR);
		}
	} else
		fp = stdout;

	be = (sample_encoding & AFMT_BIGENDIAN) ? 1 : 0;
	sign = (sample_encoding & AFMT_SIGNED) ? 0x00 : 0x80;

	fprintf(stderr, "Bps/Block     : %u / %u\n", bps, blksz);
	fprintf(stderr, "Format        : %s 0x%08x\n", fmtstr, sample_encoding);
	fprintf(stderr, "Sampling Rate : %uhz\n", sample_rate);
	fprintf(stderr, "Output        : %s\n", (outfile) ? outfile : "<stdout>");
	fprintf(stderr, "Volume Gain   : %d\n", vgain);
	signal(SIGINT, (sig_t)trapint);
	if (!rawout) {
		nhdr.size = 0xffffffffU - WAVE_HDR_SIZE + 7;
		nhdr.bit = bits;
		nhdr.rate = sample_rate;
		nhdr.channels = sample_channels;
		nhdr.align = (nhdr.bit * (2 - (nhdr.channels % 2))) >> 3;
		nhdr.byterate = nhdr.rate * nhdr.align;
		nhdr.flen = 16;
		nhdr.rifflen = nhdr.size - 8 + WAVE_HDR_SIZE;
		nhdr.tag = 1;
		if (fp == stdout) {
			if (fwrite((uint8_t *)&nhdr, 1, WAVE_HDR_SIZE, fp) != WAVE_HDR_SIZE) {
				perror("write()");
				exit(EX_IOERR);
			}
			whdr = NULL;
		} else {
			if (fseek(fp, WAVE_HDR_SIZE, SEEK_SET) != 0) {
				perror("fseek()");
				exit(EX_IOERR);
			}
			whdr = &nhdr;
		}
	}
	buf = (uint8_t *)malloc(sizeof(uint8_t) * blksz);
	fprintf(stderr, "Size          : %11u bytes\r", total);
	if (vgain > 0)
		vgain *= vgain;
	vgain += VGAIN;
	for(;;) {
		ssize_t wd, rd = read(fd, buf, blksz);
		if (rd < 1) {
			perror("read");
			fprintf(stderr, "read error.");
			break;
		}
		if (vgain != VGAIN && !(sample_encoding & (AFMT_A_LAW|AFMT_MU_LAW))) {
			max = rd - (rd % (bps / sample_channels));
			for (i = 0; i < max; i += (bps / sample_channels)) {
				switch (bps / sample_channels) {
				case 4:
#if 0
					if (be) {
						sample = (buf[i] ^ sign) << 16 |
							buf[i + 1] << 8 |
							buf[i + 2];
					} else {
						sample = buf[i + 1] |
							buf[i + 2] << 8 |
							(buf[i + 3] ^ sign) << 16;
					}
					if (sample & 0x800000)
						sample ^= ~0xffffff;
					sample = (sample * vgain) / VGAIN;
					if (sample > 0x7fffff)
						sample = 0x7fffff;
					else if (sample < -0x800000)
						sample = -0x800000;
					if (be) {
						buf[i] = (sample >> 16) ^ sign;
						buf[i + 1] = sample >> 8;
						buf[i + 2] = sample;
						buf[i + 3] = 0;
					} else {
						buf[i] = 0;
						buf[i + 1] = sample;
						buf[i + 2] = sample >> 8;
						buf[i + 3] = (sample >> 16) ^ sign;
					}
#else
					if (be) {
						sample = (buf[i] ^ sign) << 24 |
							buf[i + 1] << 16 |
							buf[i + 2] << 8 |
							buf[i + 3];
					} else {
						sample = buf[i] |
							buf[i + 1] << 8 |
							buf[i + 2] << 16 |
							(buf[i + 3] ^ sign) << 24;
					}
					sample = ((sample >> 8) * vgain) / VGAIN;
					if (sample > 0x7fffff)
						sample = 0x7fffffff;
					else if (sample < -0x800000)
						sample = -0x80000000;
					else
						sample <<= 8;
					if (be) {
						buf[i] = (sample >> 24) ^ sign;
						buf[i + 1] = sample >> 16;
						buf[i + 2] = sample >> 8;
						buf[i + 3] = sample;
					} else {
						buf[i] = sample;
						buf[i + 1] = sample >> 8;
						buf[i + 2] = sample >> 16;
						buf[i + 3] = (sample >> 24) ^ sign;
					}
#endif
					break;
				case 3:
					if (be) {
						sample = (buf[i] ^ sign) << 16 |
							buf[i + 1] << 8 |
							buf[i + 2];
					} else {
						sample = buf[i] |
							buf[i + 1] << 8 |
							(buf[i + 2] ^ sign) << 16;
					}
					if (sample & 0x800000)
						sample ^= ~0xffffff;
					sample = (sample * vgain) / VGAIN;
					if (sample > 0x7fffff)
						sample = 0x7fffff;
					else if (sample < -0x800000)
						sample = -0x800000;
					if (be) {
						buf[i] = (sample >> 16) ^ sign;
						buf[i + 1] = sample >> 8;
						buf[i + 2] = sample;
					} else {
						buf[i] = sample;
						buf[i + 1] = sample >> 8;
						buf[i + 2] = (sample >> 16) ^ sign;
					}
					break;
				case 2:
					if (be) {
						sample = (buf[i] ^ sign) << 8 |
							buf[i + 1];
					} else {
						sample = buf[i] |
							(buf[i + 1] ^ sign) << 8;
					}
					if (sample & 0x8000)
						sample ^= ~0xffff;
					sample = (sample * vgain) / VGAIN;
					if (sample > 0x7fff)
						sample = 0x7fff;
					else if (sample < -0x8000)
						sample = -0x8000;
					if (be) {
						buf[i] = (sample >> 8) ^ sign;
						buf[i + 1] = sample;
					} else {
						buf[i] = sample;
						buf[i + 1] = (sample >> 8) ^ sign;
					}
					break;
				case 1:
					sample = buf[i] ^ sign;
					if (sample & 0x80)
						sample ^= ~0xff;
					sample = (sample * vgain) / VGAIN;
					if (sample > 0x7f)
						sample = 0x7f;
					else if (sample < -0x80)
						sample = -0x80;
					buf[i] = sample ^ sign;
					break;
				}
			}
		}
		wd = fwrite(buf, 1, rd, fp);
		if (wd == 0) {
			perror("write");
			break;
		}
		total += wd;
		fprintf(stderr, "Size          : %11u bytes\r", total);
	}
	(void)trapint(SIGINT);
}

int main(int argc, char * const argv[]) 
{
	const char* sample_device = "/dev/dsp";
	uint32_t sample_rate = 44100;
	uint32_t sample_channels = 2;
	uint32_t sample_encoding = AFMT_S16_LE;
	uint32_t bits = 16;
	int32_t vgain = 0;
	char *outfile = NULL;
	char xfmtstr[] = "s16le";
	char *fmtstr = xfmtstr;
	int blksz, rawout = 1;

	int ch;

	blksz = 512;
	
	while ((ch = getopt(argc, argv, "c:d:f:hB:o:r:v:w")) != -1) {
		switch(ch) {
		case 'c':
			sample_channels = atoi(optarg);
			break;
		case 'd':
			sample_device = optarg;
			break;
		case 'f':
			fmtstr = optarg;
			break;
		case 'B':
			blksz = atoi(optarg);
			break;
		case 'o':
			outfile = optarg;
			break;
		case 'r':
			sample_rate = atoi(optarg);
			break;
		case 'v':
			vgain = atoi(optarg);
			break;
		case 'w':
			rawout = 0;
			break;
		case 'h':
		case '?':
		default:
			usage();
		}
	}

	argc -= optind;
	argv += optind;

	if (argc != 0)
		usage();

	if (sample_rate < 1 || sample_rate > 0x7ffffff) {
		fprintf(stderr, 
			"rate must lie in range 1 - %d hz\n", 0x7ffffff);
		usage();
	}

	if (vgain < -VGAIN || vgain > VGAIN) {
		fprintf(stderr, "volume gain must lie in range -%d to %d\n",
			VGAIN, VGAIN);
		usage();
	}

	if (fmtstr) {
		if (strcmp(fmtstr, "u8") == 0) {
			sample_encoding = AFMT_U8;
			bits = 8;
		} else if (strcmp(fmtstr, "s8") == 0) {
			sample_encoding = AFMT_S8;
			bits = 8;
		} else if (strcmp(fmtstr, "mulaw") == 0) {
			if (!rawout)
				fprintf(stderr, "WARNING: Using raw output!\n");
			rawout = 1;
			sample_encoding = AFMT_MU_LAW;
			bits = 8;
		} else if (strcmp(fmtstr, "alaw") == 0) {
			if (!rawout)
				fprintf(stderr, "WARNING: Using raw output!\n");
			rawout = 1;
			sample_encoding = AFMT_A_LAW;
			bits = 8;
		} else if (strcmp(fmtstr, "s16le") == 0) {
			sample_encoding = AFMT_S16_LE;
			bits = 16;
		} else if (strcmp(fmtstr, "u16le") == 0) {
			sample_encoding = AFMT_U16_LE;
			bits = 16;
		} else if (strcmp(fmtstr, "s16be") == 0) {
			sample_encoding = AFMT_S16_BE;
			bits = 16;
		} else if (strcmp(fmtstr, "u16be") == 0) {
			sample_encoding = AFMT_U16_BE;
			bits = 16;
		} else if (strcmp(fmtstr, "s24le") == 0) {
			sample_encoding = AFMT_S24_LE;
			bits = 24;
		} else if (strcmp(fmtstr, "u24le") == 0) {
			sample_encoding = AFMT_U24_LE;
			bits = 24;
		} else if (strcmp(fmtstr, "s24be") == 0) {
			sample_encoding = AFMT_S24_BE;
			bits = 24;
		} else if (strcmp(fmtstr, "u24be") == 0) {
			sample_encoding = AFMT_U24_BE;
			bits = 24;
		} else if (strcmp(fmtstr, "s32le") == 0) {
			sample_encoding = AFMT_S32_LE;
			bits = 32;
		} else if (strcmp(fmtstr, "u32le") == 0) {
			sample_encoding = AFMT_U32_LE;
			bits = 32;
		} else if (strcmp(fmtstr, "s32be") == 0) {
			sample_encoding = AFMT_S32_BE;
			bits = 32;
		} else if (strcmp(fmtstr, "u32be") == 0) {
			sample_encoding = AFMT_U32_BE;
			bits = 32;
		} else {
			fprintf(stderr, "Unknown format '%s'\n", fmtstr);
			usage();
		}
	}

	record(sample_device, sample_rate, sample_channels, sample_encoding,
			bits, blksz, vgain, rawout, outfile, fmtstr);

	return 0;
}