//
//  main.m
//  VideoGeneration
//
//  Created by Joshua Moerman on 28/01/14.
//  Copyright (c) 2014 Joshua Moerman. All rights reserved.
//

@import Foundation;
@import AVFoundation;
@import CoreGraphics;
#import <complex.h>

static const size_t w = 1280;
static const size_t h = 800;

double length(double x, double y){
	return sqrt(x*x + y*y);
}

@interface Generator : NSObject{
	double* values;
}
- (void)fillBuffer:(CVPixelBufferRef)buffer;
@end

@implementation Generator

- (id)init{
	if(self = [super init]){
		values = calloc(3*w*h, sizeof(double));
		for(size_t y = 0; y < h; ++y){
			for(size_t x = 0; x < w; ++x){
				values[3*(x + y*w) + 0] = x / (double)(w-1);
				values[3*(x + y*w) + 1] = y / (double)(h-1);
				values[3*(x + y*w) + 2] = 0.5;
			}
		}
	}
	return self;
}

- (void)dealloc{
	free(values);
}

- (void)fillBuffer:(CVPixelBufferRef)image{
	for(size_t y = 0; y < h; ++y){
		for(size_t x = 0; x < w; ++x){
			for(size_t c = 0; c < 3; ++c){
				double r = 4.0 - 0.1*length(x - 0.5*w, y - 0.5*h) / length(0.5*w, 0.5*h);
				double v = values[3*(x + y*w) + c];
				values[c + 3*(x + y*w)] = r*v*(1.0 - v);
			}
		}
	}
	
	CVPixelBufferLockBaseAddress(image, 0);
	UInt32* data = CVPixelBufferGetBaseAddress(image);
	assert(data);
	for(size_t y = 0; y < h; ++y){
		for(size_t x = 0; x < w; ++x){
			uint8_t bb = 255*values[3*(x + y*w) + 0];
			uint8_t gb = 255*values[3*(x + y*w) + 1];
			uint8_t rb = 255*values[3*(x + y*w) + 2];
			uint8_t ab = 255;
			data[x + y*1280] = ab | (rb << 8) | (gb << 16) | (bb << 24);
		}
	}
	CVPixelBufferUnlockBaseAddress(image, 0);
}

@end

static void do_it(){
	[[NSFileManager defaultManager] removeItemAtURL:[NSURL fileURLWithPath:@"test.mov"] error:nil];
	
	AVAssetWriter* asset = [AVAssetWriter assetWriterWithURL:[NSURL fileURLWithPath:@"test.mov"] fileType:AVFileTypeQuickTimeMovie error:nil];
	AVAssetWriterInput* assetInput = [AVAssetWriterInput assetWriterInputWithMediaType:AVMediaTypeVideo outputSettings:@{AVVideoCodecKey: AVVideoCodecH264, AVVideoWidthKey: @1280, AVVideoHeightKey: @800}];
	NSDictionary *sourcePixelBufferAttributesDictionary = [NSDictionary dictionaryWithObjectsAndKeys: [NSNumber numberWithInt:kCVPixelFormatType_32ARGB], kCVPixelBufferPixelFormatTypeKey, nil];
	AVAssetWriterInputPixelBufferAdaptor* adaptor = [AVAssetWriterInputPixelBufferAdaptor assetWriterInputPixelBufferAdaptorWithAssetWriterInput:assetInput sourcePixelBufferAttributes:sourcePixelBufferAttributesDictionary];
	assert(asset && assetInput && adaptor && [asset canAddInput:assetInput]);
	
	[asset addInput:assetInput];
	[asset startWriting];
	[asset startSessionAtSourceTime:kCMTimeZero];
	assert(adaptor.pixelBufferPool);
	
	Generator* gen = [[Generator alloc] init];
	
	CMTime duration = CMTimeMake(1, 30);
	CMTime current = CMTimeMake(0, 30);
	for(int i = 0; i < 300; ++i){
		while(!assetInput.readyForMoreMediaData) [NSThread sleepForTimeInterval:0.005];
		
		CVPixelBufferRef image = nil;
		CVReturn ret = CVPixelBufferPoolCreatePixelBuffer(kCFAllocatorDefault, adaptor.pixelBufferPool, &image);
		assert(ret == kCVReturnSuccess && image);

		[gen fillBuffer:image];
		[adaptor appendPixelBuffer:image withPresentationTime:current];
		
		CVPixelBufferRelease(image);
		
		current = CMTimeAdd(current, duration);
	}
	
	[assetInput markAsFinished];
	[asset finishWriting];
}

int main(int argc, const char * argv[]){
	@autoreleasepool {
		do_it();
	}
	return 0;
}