/* Copyright (c) 2011 Artjom Popov */

#import "NSData+Additions.h"

static const char encodingTable[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

UInt8 *TKDecodeBase64Data(const UInt8 *data, UInt32 dataSize, UInt32 *size)
{
    const UInt8 *end = data + dataSize;
    const UInt8 *ptr = data;
    UInt8 *decodedData = NULL;
    unsigned int decodedDataSize   = 0;
    unsigned int decodedDataLength = 0;
    int tmpIndex = 0;
    UInt8 tmp[4];
    
    if(!data || !dataSize) return NULL;
    
    while(ptr < end || tmpIndex) {
        UInt8 ch = (ptr < end) ? *ptr++ : '=';
        int tmpCount = 3;
        
        /* Character from the base64 table */
        if(ch >= 'A' && ch <= 'Z') ch = ch - 'A';
        else if(ch >= 'a' && ch <= 'z') ch = ch - 'a' + 26;
        else if(ch >= '0' && ch <= '9') ch = ch - '0' + 52;
        else if(ch == '+') ch = 62;
        else if(ch == '/') ch = 63;
        /* EOF marker */
        else if(ch == '=') {
            if(tmpIndex == 0) break;
            
            tmpCount = (tmpIndex == 1 || tmpIndex == 2) ? 1 : 2;
            tmpIndex = 3;
            ptr = end;
        }
        /* Unknown character (Probably white space) */
        else continue;
        
        tmp[tmpIndex] = ch;
        tmpIndex++;
        
        if(tmpIndex == 4) {
            tmpIndex = 0;
            
            if(decodedDataLength + 3 >= decodedDataSize) {
                // Grow factor: 32, 80, 152, 260, 422...
                decodedDataSize += decodedDataSize / 2 + 32;
                decodedData = realloc(decodedData, sizeof(UInt8) * decodedDataSize);
            }
            
            switch(tmpCount) {
                case 3:
                    decodedData[decodedDataLength + 2] = ((tmp[2] & 0x03) << 6) | (tmp[3] & 0x3F);
                case 2:
                    decodedData[decodedDataLength + 1] = ((tmp[1] & 0x0F) << 4) | ((tmp[2] & 0x3C) >> 2);
                case 1:
                    decodedData[decodedDataLength + 0] = (tmp[0] << 2) | ((tmp[1] & 0x30) >> 4);
                    break;
                default:
                    // It should never get here
                    break;
            }
            
            decodedDataLength += tmpCount;
        }
    }
    
    decodedData = realloc(decodedData, sizeof(UInt8) * decodedDataLength);
    *size = decodedDataLength;
    
    return decodedData;
}

@implementation NSData(Additions)

+ (id)dataWithBase64EncodedString:(NSString *)string;
{
	if (string == nil)
		[NSException raise:NSInvalidArgumentException format:@"", nil];
	if ([string length] == 0)
		return [NSData data];
	
	static char *decodingTable = NULL;
	if (decodingTable == NULL)
	{
		decodingTable = malloc(256);
		if (decodingTable == NULL)
			return nil;
		memset(decodingTable, CHAR_MAX, 256);
		NSUInteger i;
		for (i = 0; i < 64; i++)
			decodingTable[(short)encodingTable[i]] = i;
	}
	
	const char *characters = [string cStringUsingEncoding:NSASCIIStringEncoding];
	if (characters == NULL)     //  Not an ASCII string!
		return nil;
	char *bytes = malloc((([string length] + 3) / 4) * 3);
	if (bytes == NULL)
		return nil;
	NSUInteger length = 0;
	
	NSUInteger i = 0;
	while (YES)
	{
		char buffer[4];
		short bufferLength;
		for (bufferLength = 0; bufferLength < 4; i++)
		{
			if (characters[i] == '\0')
				break;
			
			if (isspace(characters[i]) || characters[i] == '=')
				continue;
			
			buffer[bufferLength] = decodingTable[(short)characters[i]];
			if (buffer[bufferLength++] == CHAR_MAX)      //  Illegal character!
			{
				free(bytes);
				return nil;
			}
		}
		
		if (bufferLength == 0)
			break;
		if (bufferLength == 1)      //  At least two characters are needed to produce one byte!
		{
			free(bytes);
			return nil;
		}
		
		//  Decode the characters in the buffer to bytes.
		bytes[length++] = (buffer[0] << 2) | (buffer[1] >> 4);
		if (bufferLength > 2)
			bytes[length++] = (buffer[1] << 4) | (buffer[2] >> 2);
		
		if (bufferLength > 3)
			bytes[length++] = (buffer[2] << 6) | buffer[3];
	}
	
	realloc(bytes, length);
	return [NSData dataWithBytesNoCopy:bytes length:length];
}

- (NSString *)base64Encoding;
{
	if ([self length] == 0)
		return @"";
	
    char *characters = malloc((([self length] + 2) / 3) * 4);
	if (characters == NULL)
		return nil;
	
	NSUInteger length = 0;
	
	NSUInteger i = 0;
	while (i < [self length])
	{
		char buffer[3] = {0,0,0};
		short bufferLength = 0;
		while (bufferLength < 3 && i < [self length])
			buffer[bufferLength++] = ((char *)[self bytes])[i++];
		
		//  Encode the bytes in the buffer to four characters, including padding "=" characters if necessary.
		characters[length++] = encodingTable[(buffer[0] & 0xFC) >> 2];
		characters[length++] = encodingTable[((buffer[0] & 0x03) << 4) | ((buffer[1] & 0xF0) >> 4)];
		
		if (bufferLength > 1)
			characters[length++] = encodingTable[((buffer[1] & 0x0F) << 2) | ((buffer[2] & 0xC0) >> 6)];
		else characters[length++] = '=';
		
		if (bufferLength > 2)
			characters[length++] = encodingTable[buffer[2] & 0x3F];
		else characters[length++] = '=';	
	}
	
	return [[[NSString alloc] initWithBytesNoCopy:characters length:length encoding:NSASCIIStringEncoding freeWhenDone:YES] autorelease];
}

@end
