/*
    This file is part of "Count Colours", a filter plugin for Adobe Photoshop
    Copyright (C) 2003-7 Toby Thain, toby@telegraphics.com.au

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by  
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License  
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

/* Count Colours plugin */

#include <stdio.h>

#include "world.h"

#include "PIFilter.h"

#include "entry.h"
#include "ui.h"
//#include "dbg.h"

/* use Photoshop's handle suite, since we don't know in advance how much memory will be needed */

#define NewPIHandle     pb->handleProcs->newProc
#define LockPIHandle    pb->handleProcs->lockProc
#define UnlockPIHandle  pb->handleProcs->unlockProc
#define DisposePIHandle pb->handleProcs->disposeProc

typedef char FLAG_T;
typedef unsigned long KEY_T;
typedef unsigned long long KEY64_T;

enum{ 
	CHUNK_ROWS = 128,
	BLOCKBYTES = 1L<<20, // get this much memory at a time
	SUBKEYBITS = 4, // key is broken into sub-indexes this wide; should be factor of 24: 12, 8, 6, 4...
	SUBKEYMASK = (1<<SUBKEYBITS)-1,
	TABLESIZE = sizeof(Handle)<<SUBKEYBITS,
	FLAGSSIZE = sizeof(FLAG_T)<<SUBKEYBITS
};

int bitdepth;
long toprow,total,mem,blocks;
Handle rooth = NULL,*root,indicesh = NULL;
FLAG_T *indices;

// The memory pool: a linked list of blocks allocated from Photoshop,
// each BLOCKBYTES in size.
// At the start of each block is a pointer to the previous allocation.
// When analysing 16-bit RGB, this keeps block counts in the hundreds
// rather than the 100,000s (which crashes Photoshop!)
Handle lastblock = NULL; // head of list (most recent allocation)
long blockbytes = 0; // tracks free space in the head block
Ptr lastptr; // lastptr points to free space in that block

DLLEXPORT MACPASCAL 
void PluginMain(short selector,FilterRecordPtr epb,long *data,short *result);

FLAG_T *lazyindir(FilterRecordPtr pb, Ptr *tbl, KEY_T key, int shift);
FLAG_T *lazyindir64(FilterRecordPtr pb, Ptr *tbl, KEY64_T key, int shift);
void freetree(FilterRecordPtr pb, Handle *tbl, int shift);
Ptr newfrompool(FilterRecordPtr pb,long n);
void freepool(FilterRecordPtr pb, Handle root);

void DoPrepare(FilterRecordPtr epb);
OSErr process(FilterRecordPtr pb);
OSErr DoContinue(FilterRecordPtr epb);
void DoFinish(FilterRecordPtr epb);
void RequestNext(FilterRecordPtr epb,long);

DLLEXPORT MACPASCAL 
void PluginMain(short selector,FilterRecordPtr pb,long *data,short *result){
	OSErr e = noErr;
	
	EnterCodeResource();

	switch (selector){
	case filterSelectorAbout: 
		DoAbout((AboutRecordPtr)pb); 
		break;
		
	case filterSelectorParameters: 
		break;
		
	case filterSelectorPrepare: 
		pb->maxSpace = 0; /* all allocations done via Photoshop */
		break;
		
	case filterSelectorStart: 
		// pb->depth is only available to Filters in v5.0 or later
		bitdepth = pb->depth ? pb->depth : 8;
		total = mem = blocks = 0;
		if(pb->planes == 1){
			if( (indicesh = NewPIHandle(sizeof(FLAG_T)<<bitdepth)) ){
				indices = (FLAG_T*)LockPIHandle(indicesh,false);
				memset(indices,0,sizeof(FLAG_T)<<bitdepth);
			}else
				e = memFullErr;
		}else if(pb->planes == 3){
			if( (rooth = NewPIHandle(TABLESIZE)) ){
				root = (Handle*)LockPIHandle(rooth,false);
				memset(root,0,TABLESIZE);
			}else
				e = memFullErr;
		}else{
			alert("Count Colours can process only 1 or 3 channels at once.");
			e = userCanceledErr;
		}
		toprow = pb->filterRect.top;
		RequestNext(pb,toprow);
		break;
		
	case filterSelectorContinue: 
		e = DoContinue(pb); 
		break;
		
	case filterSelectorFinish: 
		DoFinish(pb);
		
		freepool(pb,lastblock);
		if(indicesh)
			DisposePIHandle(indicesh);
		if(rooth)
			DisposePIHandle(rooth);
		
		e = userCanceledErr; // so Photoshop doesn't think the image data has changed
		break;
		
	default:
		e = filterBadParameters;
	}
		
	*result = e;
	
	ExitCodeResource();
}

void RequestNext(FilterRecordPtr pb,long toprow){
	/* Request next block of the image */

	pb->inLoPlane = pb->outLoPlane = 0;
	pb->inHiPlane = pb->outHiPlane = pb->planes-1;

	pb->inRect.left = pb->outRect.left = pb->filterRect.left;
	pb->inRect.right = pb->outRect.right = pb->filterRect.right;
	pb->inRect.top = pb->outRect.top = toprow;
	pb->inRect.bottom = toprow + CHUNK_ROWS;
	if(pb->inRect.bottom > pb->filterRect.bottom)
		pb->inRect.bottom = pb->filterRect.bottom;
	pb->outRect.bottom = pb->inRect.bottom;
}

Ptr newfrompool(FilterRecordPtr pb,long n){
	Ptr p;
	
	if(blockbytes < n){
		// won't fit in head block, need another one
		Handle h = NewPIHandle(BLOCKBYTES);
		mem += BLOCKBYTES; ++blocks;  // just accounting

		lastptr = LockPIHandle(h,false);

		// add it to chain
		(*(Handle*)lastptr) = lastblock;
		lastblock = h;

		lastptr += sizeof(Handle);
		blockbytes = BLOCKBYTES - sizeof(Handle);
	}
	p = lastptr;
	lastptr += n;
	blockbytes -= n;
	return p;
}

void freepool(FilterRecordPtr pb,Handle root){
	if(root){
		Handle *p = (Handle*)LockPIHandle(root,false),
		       h = *p;
		DisposePIHandle(root);
		freepool(pb,h);
	}
}

FLAG_T *lazyindir(FilterRecordPtr pb,Ptr *tbl,KEY_T key,int shift){
	shift -= SUBKEYBITS;
	if(!shift)  // down to leaf table?
		return (FLAG_T*)tbl + (key&SUBKEYMASK);
	else{  // no, must go deeper
		int j = (key>>shift) & SUBKEYMASK;
		if(!tbl[j]){  // missing entry?
			int n = shift==SUBKEYBITS /*new leaf table?*/ ? sizeof(FLAG_T)<<SUBKEYBITS : TABLESIZE;
			tbl[j] = newfrompool(pb,n);
			memset(tbl[j],0,n);
		}
		return lazyindir(pb, (Ptr*)tbl[j], key, shift);
	}
}

// Could use the following routine for shorter keys too,
// but let's use the 32-bit routine above unless we REALLY NEED 64-bit keys.
// Note that this routine uses the 'pool' memory management,
// because dealing with the longer key is likely to allocate large numbers of blocks.

FLAG_T *lazyindir64(FilterRecordPtr pb,Ptr *tbl,KEY64_T key,int shift){
	shift -= SUBKEYBITS;
	if(!shift)
		return (FLAG_T*)tbl + (key&SUBKEYMASK);
	else{
		int j = (key>>shift) & SUBKEYMASK;
		if(!tbl[j]){
			int n = shift==SUBKEYBITS /*new leaf table?*/ ? (1<<SUBKEYBITS)/*FLAGSSIZE*/ : TABLESIZE;
			tbl[j] = newfrompool(pb,n);
			memset(tbl[j],0,n);
		}
		return lazyindir64(pb, (Ptr*)tbl[j], key, shift);
	}
}

OSErr process(FilterRecordPtr pb){
	unsigned char *inrow,*inp;
	unsigned short *inp16;
	int ncols = pb->inRect.right - pb->inRect.left,
	    nrows = pb->inRect.bottom - pb->inRect.top,
	    i,j;

	for( inrow = pb->inData, i = 0 ; i < nrows ; ++i, inrow += pb->inRowBytes ){
		if(pb->planes == 3){
			// RGB image
			if(bitdepth == 8)
				for( inp = inrow, j = 0 ; j < ncols ; ++j, inp += pb->planes ){
					KEY_T k = (inp[0]<<16) | (inp[1]<<8) | inp[2];  // construct key
					FLAG_T *p = lazyindir(pb,(Ptr*)root,k,24);  // look it up
					if(!*p){ // not found?
						*p = 1;
						++total; // count a new colour
					}
				}
			else  // 16 bit depth
				for( inp16 = (unsigned short*)inrow, j = 0 ; j < ncols ; ++j, inp16 += pb->planes ){
					KEY64_T k = ((KEY64_T)inp16[0]<<32) | ((KEY64_T)inp16[1]<<16) | (KEY64_T)inp16[2];
					FLAG_T *p = lazyindir64(pb,(Ptr*)root,k,48);
					if(!*p){
						*p = 1;
						++total;
					}
				}
		}else{
			// single channel
			if(bitdepth == 8)
				for( inp = inrow, j = 0 ; j < ncols ; ++j )
					indices[*inp++] = 1;
			else  // 16 bit depth
				for( inp16 = (unsigned short*)inrow, j = 0 ; j < ncols ; ++j )
					indices[*inp16++] = 1;
		}
	}
	
	return noErr;
}

OSErr DoContinue(FilterRecordPtr pb){
	OSErr e;

	/* See if user has aborted the operation */
	if(pb->abortProc())
		return userCanceledErr;
	else{
		//char s[0x100];
		//sprintf(s,"mem = %ld = %ld M\nblocks = %ld\nblockbytes = %ld",mem,mem>>20,blocks,blockbytes);
		//alert(s);
		pb->progressProc(toprow - pb->filterRect.top, pb->filterRect.bottom - pb->filterRect.top);

		// got a chunk of the image; process it
	
		if( (e = process(pb)) )
			return e;
		else{
			toprow += CHUNK_ROWS;
			if(toprow < pb->filterRect.bottom)
				RequestNext(pb,toprow);
			else{
				// tell Photoshop we're done
				pb->inRect.left = pb->inRect.right = pb->inRect.top = pb->inRect.bottom = 0;
				pb->outRect = pb->maskRect = pb->inRect;
			}
			return noErr;
		}
	}
}

void DoFinish(FilterRecordPtr pb){
	int i,rgb = pb->planes == 3,
	    w = pb->filterRect.right - pb->filterRect.left,
	    h = pb->filterRect.bottom - pb->filterRect.top;
	char s[0x100];
	long n;
	
	if(rgb)
		n = total;
	else
		for( i = 1<<bitdepth, n = 0 ; i-- ; )
			n += indices[i];

	sprintf(s,"%ld %ss in %d channel%s (%d bit)\n\nSize %d x %d = %ld pixels",
		n,rgb ? "colour" : "value",pb->planes,rgb ? "s" : "",bitdepth,w,h,(long)w*h);
	alert(s);
	//sprintf(s,"mem = %ld = %ld M\nblocks = %ld\nblockbytes = %ld",mem,mem>>20,blocks,blockbytes);
	//alert(s);
}