//
//  NSDictionaryHOMAdditions.m
//  HigherOrderMessaging
//
//  Created by Ofri Wolfus on 24/02/06.
//  Copyright 2006 Ofri Wolfus. All rights reserved.
//

#import "NSDictionaryHOMAdditions.h"
#import "NSArrayHOMAdditions.h"
#import "HOMCArrayIterations.h"
#import "HOMUtilities.h"
#import "HOMIteratedArgument.h"
#include <malloc/malloc.h>
#include <objc/objc-runtime.h>


@implementation NSDictionary (NewHOMIteration)

- (id)collect:(Message *)msg {
	SEL						sel = [msg selector];
	marg_list				origArgs = [msg arguments];
	unsigned				argsSize = [msg argumentsSize];
	marg_list				args = malloc(malloc_size(origArgs));
	unsigned				argCount = hom_getNumberOfArguments(sel);
	unsigned				i, count = [self count];
	Class					iteratedArgCls = [HOMIteratedArgument class];
	char					*types = calloc(argCount, sizeof(char));
	const char				*t = [msg types];
	id						*objects = calloc(count, sizeof(id)), *keys = calloc(count, sizeof(id));
	CFMutableDictionaryRef	result = CFDictionaryCreateMutable(kCFAllocatorDefault, count,
															   &kCFTypeDictionaryKeyCallBacks,
															   &kCFTypeDictionaryValueCallBacks);
	
	// Get all our objects and keys
	CFDictionaryGetKeysAndValues((CFDictionaryRef)self, (const void **)keys, (const void **)objects);
	
	// Copy the arguments list since we're going to modify it
	memcpy(args, origArgs, malloc_size(args));
	
	// Scan the arguments of our message and find instances of HOMIteratedArgument.
	// Our types array will mark the places of these with the letter 'z' so we don't have to check again.
	// This loop also finds the largest number of arguments that needs to be iterated, and stores it in count.
	for (i = 2U; i < argCount; i++) {
		types[i] = *hom_getArgumentFromTypes(t, i);
		if (types[i] == _C_ID) {
			if ([marg_getValue(args, hom_getArgumentOffsetFromTypes(t, i), id) isKindOfClass:iteratedArgCls])
				types[i] = 'z';	// This is our special mark for an iterated argument
		}
	}
	
	// Loop through all objects in all our iterated arguments, update our args list, and resend it
	for (i = 0U; i < count; i++) {
		unsigned j;
		
		// Update the remaining arguments in our list
		for (j = 2U; j < argCount; j++) {
			if (types[j] == 'z') {
				int offset = hom_getArgumentOffsetFromTypes(t, j);
				HOMIteratedArgument *arg = marg_getValue(origArgs, offset, HOMIteratedArgument *);
				
				marg_setValue(args, offset, id, [arg nextObject]);
				
				// When there are no arguments left for this argument iterator, we mark its type as id
				// so we'll skip it in the next run.
				if ([arg isFinished])
					types[j] = _C_ID;
			}
		}
		
		// Send the updated arguments list
		//addObject(result, @selector(addObject:), objc_msgSendv(objects[i], sel, argsSize, args));
		CFDictionaryAddValue(result, keys[i], objc_msgSendv(objects[i], sel, argsSize, args));
	}
	
	// Clean up
	free(types);
	free(args);
	free(objects);
	free(keys);
	
	return [(id)result autorelease];
}

- (id)getObjectsThatResponded:(BOOL)resultVal forMessage:(Message *)firstMsg :(va_list)additionalMessages {
	Message					*msg;
	unsigned				i, count = [self count];
	CFMutableDictionaryRef	result = NULL;
	
	if (count > 0U) {
		va_list					messages;
		id						*objects = calloc(count, sizeof(id));
		id						*keys = calloc(count, sizeof(id));
		
		result = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);
		
		// Get all objects and keys
		CFDictionaryGetKeysAndValues((CFDictionaryRef)self, (const void **)keys, (const void **)objects);
		
		// Loop through all our objects, and send it all the messages passed to us.
		// All messages except the last are assumed to return id, and the last one should return a BOOL
		// (although any integer value will work).
		for (i = 0U; i < count; i++) {
			// Send the first message to our object and store the result.
			// Since id is a pointer, and a pointer can hold BOOL values,
			// we're fine even if this is the only message.
			id	obj = objc_msgSendv(objects[i], [firstMsg selector], [firstMsg argumentsSize], [firstMsg arguments]);
			
			// Send every other messages to the previous result.
			// We assume the result is an object. If it's not, bad things will happen, but that's not our fault.
			// NOTE: We don't verify that the receiver responds to the actual message like -select does
			// which give us a nice speed boost, but may also cause an exception to be thrown.
			// TODO: Cache the IMPs of -[Message selector], -[Message argumentsSize] and -[Message arguments].
			va_copy(messages, additionalMessages);
			
			while ((msg = va_arg(messages, Message *)))
				obj = objc_msgSendv(obj, [msg selector], [msg argumentsSize], [msg arguments]);
			
			va_end(messages);

			// Check if the final result equals the result value we're looking for.
			// If true, we'll add the original object (and its key) to our results dictionary.
			if ((BOOL)(int)obj == resultVal)
				CFDictionarySetValue(result, keys[i], objects[i]);
		}
		
		// Clean up
		free(objects);
		free(keys);
		va_end(additionalMessages);
		
		// Make our return value pretty
		if (CFDictionaryGetCount(result) == 0) {
			CFRelease(result);
			result = nil;
		}
	}
	
	return [(id)result autorelease];
}

- (id)selectWhere:(Message *)firstMessage, ... {
	va_list	va;
	
	DPAssert(firstMessage != nil, @"The passed message must not be nil");
	va_start(va, firstMessage);
	return [self getObjectsThatResponded:YES forMessage:firstMessage :va];
}

- (id)rejectWhere:(Message *)firstMessage, ... {
	va_list	va;
	
	DPAssert(firstMessage != nil, @"The passed message must not be nil");
	va_start(va, firstMessage);
	return [self getObjectsThatResponded:NO forMessage:firstMessage :va];
}

// Note this is getObject.. and not getObjets (a single object!)
- (id)getObjectThatResponded:(BOOL)resultVal forMessage:(Message *)firstMsg :(va_list)additionalMessages key:(id *)key {
	Message		*msg;
	unsigned	i, count = [self count];
	id			result = nil;
	
	if (count > 0U) {
		va_list					messages;
		id						*objects = calloc(count, sizeof(id));
		id						*keys = calloc(count, sizeof(id));
		
		// Get all objects and keys
		CFDictionaryGetKeysAndValues((CFDictionaryRef)self, (const void **)keys, (const void **)objects);
		
		// Loop through all our objects, and send it all the messages passed to us.
		// All messages except the last are assumed to return id, and the last one should return a BOOL
		// (although any integer value will work).
		for (i = 0U; i < count; i++) {
			// Send the first message to our object and store the result.
			// Since id is a pointer, and a pointer can hold BOOL values,
			// we're fine even if this is the only message.
			id	obj = objc_msgSendv(objects[i], [firstMsg selector], [firstMsg argumentsSize], [firstMsg arguments]);
			
			// Send every other messages to the previous result.
			// We assume the result is an object. If it's not, bad things will happen, but that's not our fault.
			// NOTE: We don't verify that the receiver responds to the actual message like -select does
			// which give us a nice speed boost, but may also cause an exception to be thrown.
			// TODO: Cache the IMPs of -[Message selector], -[Message argumentsSize] and -[Message arguments].
			va_copy(messages, additionalMessages);
			
			while ((msg = va_arg(messages, Message *)))
				obj = objc_msgSendv(obj, [msg selector], [msg argumentsSize], [msg arguments]);
			
			va_end(messages);
			
			// Check if the final result equals the result value we're looking for.
			// If true, we'll add the original object (and its key) to our results dictionary.
			if ((BOOL)(int)obj == resultVal) {
				result = objects[i];
				if (key)
					*key = keys[i];
				
				break;
			}
		}
		
		// Clean up
		free(objects);
		free(keys);
		va_end(additionalMessages);
	}
	
	return result;
}

- (id)selectSingleWhere:(Message *)firstMessage, ... {
	va_list	va;
	id		key = nil, obj = nil;
	
	DPAssert(firstMessage != nil, @"The passed message must not be nil");
	va_start(va, firstMessage);
	obj = [self getObjectThatResponded:YES forMessage:firstMessage :va key:&key];
	
	if (obj)
		return [NSDictionary dictionaryWithObject:obj forKey:key];
	
	return nil;
}

- (id)rejectSingleWhere:(Message *)firstMessage, ... {
	va_list	va;
	id		key = nil;
	
	DPAssert(firstMessage != nil, @"The passed message must not be nil");
	va_start(va, firstMessage);
	[self getObjectThatResponded:YES forMessage:firstMessage :va key:&key];
	
	if (key) {
		NSMutableDictionary *dict = [self mutableCopy];
		[dict removeObjectForKey:key];
		return [dict autorelease];
	}
	
	return nil;
}

- (id)each {
	return [[self allValues] each];
}

@end