/* This program is licensed under the GNU General Public License, version 2,
 * a copy of which is included with this program.
 *
 *
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <ogg/ogg.h>
#include <vorbis/codec.h>
#include <time.h>
#include <errno.h>
#include <string.h>

#include "oggfix.h"

#include <locale.h>

#ifdef _WIN32
#define FORMAT_INT64 "%I64d"
#else
#define FORMAT_INT64 "%Ld"
#endif

static vcut_packet *save_packet(ogg_packet *packet)
{
	vcut_packet *p = malloc(sizeof(vcut_packet));

	p->length = packet->bytes;
	p->packet = malloc(p->length);
	memcpy(p->packet, packet->packet, p->length);

	return p;
}

static void free_packet(vcut_packet *p)
{
	if(p)
	{
		if(p->packet)
			free(p->packet);
		free(p);
	}
}

static long get_blocksize(vcut_state *s, vorbis_info *vi, ogg_packet *op)
{
	int this = vorbis_packet_blocksize(vi, op);
	int ret = (this+s->prevW)/4;

	s->prevW = this;
	return ret;
}

static int update_sync(vcut_state *s, FILE *f)
{
	unsigned char *buffer = ogg_sync_buffer(s->sync_in, 4096);
	int bytes = fread(buffer,1,4096,f);
	ogg_sync_wrote(s->sync_in, bytes);
	return bytes;
}

/* Returns 0 for success, or -1 on failure. */
static int write_pages_to_file(ogg_stream_state *stream, 
		FILE *file, int flush)
{
	ogg_page page;

	if(flush)
	{
		while(ogg_stream_flush(stream, &page))
		{
			if(fwrite(page.header,1,page.header_len, file) != page.header_len)
				return -1;
			if(fwrite(page.body,1,page.body_len, file) != page.body_len)
				return -1;
		}
	}
	else
	{
		while(ogg_stream_pageout(stream, &page))
		{
			if(fwrite(page.header,1,page.header_len, file) != page.header_len)
				return -1;
			if(fwrite(page.body,1,page.body_len, file) != page.body_len)
				return -1;
		}
	}

	return 0;
}



static int process_stream(vcut_state *s, ogg_stream_state *stream, 
		FILE *in, FILE *f)
{
	ogg_packet packet;
	ogg_page page;
	int eos=0;
	int result;
	ogg_int64_t page_granpos, current_granpos=s->initialgranpos;
	ogg_int64_t packetnum=0; /* Should this start from 0 or 3 ? */

	while(!eos)
	{
		while(!eos)
		{
			result=ogg_sync_pageout(s->sync_in, &page);
			if(result==0) break;
			else if(result==-1)
				fprintf(stderr, "Recoverable bitstream error(fixed)\n");
			else
			{
				page_granpos = ogg_page_granulepos(&page) - s->cutpoint;
				if(ogg_page_eos(&page))eos=1;
				ogg_stream_pagein(s->stream_in, &page);
				while(1)
				{
					result = ogg_stream_packetout(s->stream_in, &packet);
					if(result==0) break;
					else if(result==-1) fprintf(stderr, "Bitstream error(fixed)\n");
					else
					{
						int bs = get_blocksize(s, s->vi, &packet);
						current_granpos += bs;
						if(current_granpos < 0)
						{
							current_granpos = page_granpos;
						} 
						/*
							recalculate the ganulepos an packetno values.
							this fixes the recorded stream.
						*/
						packet.granulepos = current_granpos;
					  packet.packetno = packetnum++;
					  
						ogg_stream_packetin(stream, &packet);
						if(write_pages_to_file(stream,f, 0))
							return -1;
					}
				}
			}
		}
		if(!eos)
		{
			if(update_sync(s, in)==0)
			{
				eos=1;
			}
		}
	}

	return 0;
}			

static void submit_headers_to_stream(ogg_stream_state *stream, vcut_state *s) 
{
	int i;
	for(i=0;i<3;i++)
	{
		ogg_packet p;
		p.bytes = s->headers[i]->length;
		p.packet = s->headers[i]->packet;
		p.b_o_s = ((i==0)?1:0);
		p.e_o_s = 0;
		p.granulepos=0;

		ogg_stream_packetin(stream, &p);
	}
}
									
/* Pull out and save the 3 header packets from the input file.
 */
static int process_headers(vcut_state *s)
{
	vorbis_comment vc;
	ogg_page page;
	ogg_packet packet;
	int bytes;
	int i;
	unsigned char *buffer;

	ogg_sync_init(s->sync_in);
	
	vorbis_info_init(s->vi);
	vorbis_comment_init(&vc);

	buffer = ogg_sync_buffer(s->sync_in, 4096);
	bytes = fread(buffer, 1, 4096, s->in);
	ogg_sync_wrote(s->sync_in, bytes);

	if(ogg_sync_pageout(s->sync_in, &page)!=1){
		fprintf(stderr, "Input not ogg.\n");
		return -1;
	}

	s->serial = ogg_page_serialno(&page);

	ogg_stream_init(s->stream_in, s->serial);

	if(ogg_stream_pagein(s->stream_in, &page) <0)
	{
		fprintf(stderr,"Error in first page\n");
		return -1;
	}

	if(ogg_stream_packetout(s->stream_in, &packet)!=1){
		fprintf(stderr,"error in first packet\n");
		return -1;
	}

	if(vorbis_synthesis_headerin(s->vi, &vc, &packet)<0)
	{
		fprintf(stderr, "Error in primary header: not vorbis?\n");
		return -1;
	}

	s->headers[0] = save_packet(&packet);

	i=0;
	while(i<2)
	{
		while(i<2) {
			int res = ogg_sync_pageout(s->sync_in, &page);
			if(res==0)break;
			if(res==1)
			{
				ogg_stream_pagein(s->stream_in, &page);
				while(i<2)
				{
					res = ogg_stream_packetout(s->stream_in, &packet);
					if(res==0)break;
					if(res<0){
						fprintf(stderr, "Secondary header corrupt\n");
						return -1;
					}
					s->headers[i+1] = save_packet(&packet);
					vorbis_synthesis_headerin(s->vi,&vc,&packet);
					i++;
				}
			}
		}
		buffer=ogg_sync_buffer(s->sync_in, 4096);
		bytes=fread(buffer,1,4096,s->in);
		if(bytes==0 && i<2)
		{
			fprintf(stderr,"EOF in headers\n");
			return -1;
		}
		ogg_sync_wrote(s->sync_in, bytes);
	}

	vorbis_comment_clear(&vc);

	return 0;
}


int oggfix(FILE *in, FILE *out)
{
	ogg_int64_t cutpoint=0;
	int ret=0;
	vcut_state *state;


	state = vcut_new();

	vcut_set_files(state, in, out);
	
	if(vcut_process(state))
	{
		fprintf(stderr, "Processing failed\n");
		ret = 1;
	}

	vcut_free(state);

	return ret;
}

int vcut_process(vcut_state *s)
{
	ogg_stream_state  stream_out;
	
	/* Read headers in, and save them */
	if(process_headers(s))
	{
		fprintf(stderr, "Error reading headers\n");
		return -1;
	}

	/* ok, headers are all in, and saved */
	vorbis_synthesis_init(s->vd,s->vi);
	vorbis_block_init(s->vd,s->vb);

	ogg_stream_init(&stream_out,s->serial); /* first file gets original */


	submit_headers_to_stream(&stream_out, s);
	if(write_pages_to_file(&stream_out, s->out, 1))
		return -1;

	if(process_stream(s, &stream_out, s->in, s->out))
	{
		fprintf(stderr, "Error writing second output file\n");
		return -1;
	}
	ogg_stream_clear(&stream_out);

	return 0;

}

vcut_state *vcut_new(void)
{
	vcut_state *s = malloc(sizeof(vcut_state));
	memset(s,0,sizeof(vcut_state));

	s->sync_in = malloc(sizeof(ogg_sync_state));
	s->stream_in = malloc(sizeof(ogg_stream_state));
	s->vd = malloc(sizeof(vorbis_dsp_state));
	s->vi = malloc(sizeof(vorbis_info));
	s->vb = malloc(sizeof(vorbis_block));

	s->headers = malloc(sizeof(vcut_packet)*3);
	memset(s->headers, 0, sizeof(vcut_packet)*3);
	s->packets = malloc(sizeof(vcut_packet)*2);
	memset(s->packets, 0, sizeof(vcut_packet)*2);

	return s;
}

/* Full cleanup of internal state and vorbis/ogg structures */
void vcut_free(vcut_state *s)
{
	if(s)
	{
		if(s->packets)
		{
			if(s->packets[0])
				free_packet(s->packets[0]);
			if(s->packets[1])
				free_packet(s->packets[1]);
			free(s->packets);
		}

		if(s->headers)
		{
			int i;
			for(i=0; i < 3; i++)
				if(s->headers[i])
					free_packet(s->headers[i]);
			free(s->headers);
		}

		if(s->vb)
		{
			vorbis_block_clear(s->vb);
			free(s->vb);
		}
		if(s->vd)
		{
			vorbis_dsp_clear(s->vd);
			free(s->vd);
		}
		if(s->vi)
		{
			vorbis_info_clear(s->vi);
			free(s->vi);
		}
		if(s->stream_in)
		{
			ogg_stream_clear(s->stream_in);
			free(s->stream_in);
		}
		if(s->sync_in)
		{
			ogg_sync_clear(s->sync_in);
			free(s->sync_in);
		}

		free(s);
	}
}

void vcut_set_files(vcut_state *s, FILE *in, FILE *out)
{
	s->in = in;
	s->out = out;
	s->cutpoint = 0;
}