#include <u.h>
#include <libc.h>
#include <bio.h>
#include <tiff.h>
#include <thread.h>	//!

int rflag = 1;
int dflag = 1;

static char *tagtab[] = {
[Tsubfiletype-0xfe]	"Tsubfiletype",
[Twidth-0xfe]		"Twidth",
[Tlength-0xfe]		"Tlength",
[Tbitspersample-0xfe]	"Tbitspersample",
[Tcompression-0xfe]	"Tcompression",
[Tphotometric-0xfe]	"Tphotometric",
[Tthresholding-0xfe]	"Tthresholding",
[Tcellwidth-0xfe]		"Tcellwidth",
[Tcelllength-0xfe]		"Tcelllength",
[Thwmodel-0xfe]		"Thwmodel",
[Tfillorder-0xfe]		"Tfillorder",
[Timagedesc-0xfe]		"Timagedesc",

[Tstripoffsets-0xfe]	"Tstripoffsets",
[Torientation-0xfe]	"Torientation",
[Tsamplespp-0xfe]	"Tsamplespp",
[Trowsperstrip-0xfe]	"Trowsperstrip",
[Tstripbytecounts-0xfe]	"Tstripbytecounts",

[Txresolution-0xfe]	"Txresolution",
[Tyresolution-0xfe]	"Tyresolution",
[Tplanarconf-0xfe]	"Tplanarconf",
[Tresolutionunit-0xfe]	"Tresolutionunit",
[Tpageno-0xfe]		"Tpageno",
[Tsoftware-0xfe]		"Tsoftware",
[Tdatetime-0xfe]		"Tdatetime",
[Tartist-0xfe]		"Tartist",
[Tcomputer-0xfe]		"Tcomputer",
[Tpredictor-0xfe]		"Tpredictor",
[Tcolormap-0xfe]		"Tcolormap",
[Textrasamples-0xfe]	"Textrasamples",
};

char*
tagname(IFDtype t)
{
	int i;

	i = t-0xfe;
	if(i<0 || i>=Tend)
		return "unknown";
	return tagtab[i];
}

typedef struct {
	IFDtype	t;
	char	*name;
	int	size;
	uint	fmtstride;
	char	*fmt;
	char	desc;
} IFDinfo;

IFDinfo  IFDtab[] = {
	{0,	"b0rked",	0,	0,"b0rked",	0,},
	{Ibyte,	"byte",	1,	1, "%02uhhx",	'c',},
	{Iascii,	"ascii",	1,	~0,"%.*s",	'a'},
	{Ishort,	"short",	2,	1, "%uhd ",	's'},	
	{Ilong,	"long",	4,	1, "%uld ",	'l'},
	{Irat,	"rat",	8,	2, "[%uld.%uld]",	'r'},
	{Isbyte,	"sbyte",	1,	1, "%hhd",	'B'},
	{Iundef,	"undef",	1,	~0, "%.*s",	'U'},
	{Isshort,	"sshort",	2,	1, "%hd ",	'S'},
	{Islong,	"slong",	4,	1, "%ld ",	'L'},
	{Israt,	"srat",	8,	1, "[%ld.%ld]",	'R'},
	{Ifloat,	"float",	4,	1, "%f ",		'f'},
	{Idbl,	"dbl",	8,	1, "%d ",		'd'},
	{0,	"b0rked",	0,	0, 0,		},
};

int
typesize(IFDtype t)
{
	return IFDtab[t].size;
}

char*
typename(IFDtype t)
{
	return IFDtab[t].name;
}

long
tiffshort(uchar *p, End o)
{
	switch(o){
	case BE:
		return p[1] | p[0]<<8;
	case LE:
		return p[0] | p[1]<<8;
	}
	assert(0);
	return -1;
}

long
Btiffshort(Biobuf *b, End o)
{
	uchar p[2];
	
	Bread(b, p, sizeof p);
	return tiffshort(p, o);
}

long
Btiffshorts(Biobuf *b, short *p, int n, End o)
{
	while(n--)
		*p++ = Btiffshort(b, o);
	return 0;
}

long
tifflong(uchar *p, End o)
{
	switch(o){
	case BE:
		return p[3] | p[2]<<8 | p[1]<<16 | p[0]<<24;
	case LE:
		return p[0] | p[1]<<8 | p[2]<<16 | p[3]<<24;
	}
	assert(0);
	return -1;
}

long
Btifflong(Biobuf *b, End o)
{
	uchar p[4];
	
	Bread(b, p, sizeof p);
	return tifflong(p, o);
}

long
Btifflongs(Biobuf *b, long *p, int n, End o)
{
	while(n--)
		*p++ = Btifflong(b, o);
	return 0;
}

int
ifdvalread(Biobuf *b, IFD* i, End o)
{
	int s;

	s = typesize(i->type);
	i->cp = mallocz((s*i->n|15)+1, 1);
	if(!i->cp)
		sysfatal("malloc");
	if(Bseek(b, i->offset, 0) == -1)
		return -1;
	switch(s){
	case 1:
		Bread(b, i->cp, i->n);
		break;
	case 2:
		Btiffshorts(b, i->sp, i->n, o);
		break;
	case 4:
		Btifflongs(b, i->lp, i->n, o);
		break;
	case 8:
		Btifflongs(b, i->lp, i->n*2, o);
		break;
	}
	return 0;
}

int
ifdvalcp(IFD* i, End o)
{
	int s, n;
	uchar *p;

	i->cp = mallocz(16, 1);
	if(!i->cp)
		sysfatal("malloc");
	p = i->roffset;
	switch(s = typesize(i->type)){
	case 1:
		memcpy(i->cp, p, 4);
		break;
	case 2:
		for(n = 0; s; n++, s -= 2)
			i->sp[n] = tiffshort(p+2*n, o);
	case 4:
	case 8:
		for(n = 0; s; n++, s -= 4)
			i->lp[n] = tifflong(p+4*n, o);
	}
//	i->offset = 0;
	return 0;
}

int
readifd(Biobuf *b, vlong offset, Tiff* t)
{
	IFD *i, *e;
	long n;

	if(Bseek(b, offset, 0) == -1)
		return -1;
	t->n = t->alloc = Btiffshort(b, t->order);
	t->ifds = mallocz(t->alloc * sizeof *t->ifds, 1);
	if(t->ifds == 0)
		return -1;

	for(i = t->ifds, e = i+t->n; i<e; i++){
		i->tag = Btiffshort(b, t->order);
		i->type = Btiffshort(b, t->order);
		i->n = Btifflong(b, t->order);
		Bread(b, i->roffset, 4);		// avoid double-endian conversion
		i->offset = tifflong(i->roffset, t->order);
	}

	// here's where we should check for other IFDs.  we want a null next ptr.
	if((n = Btifflong(b, t->order)) != 0)
		fprint(2, "another IFD after this one @%uld -- ignored.\n", n);

	for(i = t->ifds, e = i+t->n; i<e; i++)
		if(typesize(i->type)*i->n > 4)
			ifdvalread(b, i, t->order);
		else
			ifdvalcp(i, t->order);
	return 0;
}

void
tifffree(Tiff *t)
{
	IFD *i, *e;

	e = t->ifds + t->n;
	for(i = t->ifds; i<e; i++)
		free(i->sp);
	free(t->ifds);
	free(t->rawimg);
	free(t);
}

void
ifdprint(Tiff *t)
{
	IFD *i, *e;
	char *tag, *tab, *fmt, buf[9];
	int j, l, stride, tsize;

	for(i = t->ifds, e = i+t->n; i < e; i++){
		tag = tagname(i->tag);
		if(!tag){
			snprint(buf, sizeof buf, "%x", i->tag);
			tag = buf;
		}
		l = strlen(tag)+2;
		if(l<=13)
			tab = "\t\t";
		else
			tab = "\t";
		fprint(2, "%03x [%s]%s%c/%uld", i->tag, tag, tab, IFDtab[i->type].desc, i->n);
		if(typesize(i->type)>4 || i->n>1)
			fprint(2, "@%-8uld\t", i->offset);
		else
			fprint(2, "\t\t");
		
		switch(i->tag){
		case Tcompression:
			if(tiffzstr(i->sp[0]) == 0)
				break;
			fprint(2, "%s\n", tiffzstr(i->sp[0]));
			continue;
		}

		tsize = IFDtab[i->type].size;
		fmt = IFDtab[i->type].fmt;
		stride = IFDtab[i->type].fmtstride;
		for(j = 0; j < i->n ; j += stride)
			switch(tsize){
			case 1:
				if(stride == 1)
					; //fprint(2, fmt, i->cp[j]);
				else
					fprint(2, fmt, i->n, i->cp);
				break;
			case 2:
				fprint(2, fmt, i->sp[j]);
				break;
			case 4:
				fprint(2, fmt, i->lp[j]);
				break;
			case 8:
				//fprint(2, fmt, i->lp[j], i->lp[j+1]);
				fprint(2, "[%g]", (double)i->lp[j] / (double)i->lp[j+1]);
				break;
			}
		fprint(2, "\n");
	}
}

IFD*
lookifdptr(Tiff *t, Itype tag)
{
	IFD *i, *e;

	for(i = t->ifds, e = i+t->n; i<e; i++)
		if(i->tag == tag)
			return i;
	return 0;
}

long
ifdidx(IFD *i, int n)
{
	if(i)
	switch(typesize(i->type)){
	case 2:
		return i->sp[n];
	case 4:
		return i->lp[n];
	}
	sysfatal("bad type");
	return -1;
}

long
lookifd(Tiff *t, Itype tag)
{
	IFD *i;

	i = lookifdptr(t, tag);
	if(i && i->n == 1)
		return ifdidx(i, 0);
	return -1;
}

long
tiffimglen(Tiff *t)
{
	ulong x, y, bpp, i;
	IFD* cdepth;

	x = lookifd(t, Twidth);
	y = lookifd(t, Tlength);
	cdepth = lookifdptr(t, Tbitspersample);
	if(!cdepth)
		return -1;
	for(bpp = 0, i = 0; i < cdepth->n; i++)
		bpp += cdepth->sp[i];
	debug("x=%uld y=%uld bpp=%uld\n", x, y, bpp);
	if(x<1 || y<1 || bpp<1)
		return -1;
	t->bpp = bpp;
	return x*y*bpp>>3;
}

static long
sz(Tiff *t, long l, long i)
{
	int b, r, c, bpl;
	IFD *z;

	z = lookifdptr(t, Tstripoffsets);
	r = lookifd(t, Trowsperstrip);
	c = lookifd(t, Tlength);
	if(i == z->n-1)
		r = c - (z->n-1)*r;	// truncated final strip
	b = (int)((vlong)l*(vlong)r/(vlong)c);
	bpl = t->bpp*lookifd(t, Twidth);

	debug("b = %d; l = %ld, r=%d, c=%d\n", b, l, r, c);
	// rows must be an even number of bytes.
	// cheat if bpl == 0 on Z8
	if((bpl%8) == 0)
		return b;
	for(i = 1; i <= r; i++)
		if((bpl*i)%8)
			b++;
	return b;
}

int
tiffimg(Biobuf *b, Tiff *t)
{
	IFD *stripp, *stripsz;
	long l, i, n, r;

	l = tiffimglen(t);
	stripp = lookifdptr(t, Tstripoffsets);
	stripsz = lookifdptr(t, Tstripbytecounts);
	if(l == -1 || !stripp || !stripsz){
		debug("tiffimg fails(l=%ld)\n", l);
		return -1;
	}

	t->rawimge = t->rawimg = malloc(l+10);
	if(!t->rawimg)
		return -1;

	for(i = 0, n = 0; i < stripp->n; i++){
		if(Bseek(b, ifdidx(stripp, i), 0) < 0)
			return -1;
		// ifdidx(stripsz, i) is the *compressed* size
		r = tiffunz(t, lookifd(t, Tcompression), b, sz(t, l, i));
		debug("tiffunz → %ld; %ld\n", r, n);
		if(r < 0)
			return -1;
		t->rawimge += r;
		n+= r;
	}

	return 0;
}

int
tiffsetrgb(Tiff *t, int bpc)
{
	IFD *i;
	int j;

	i = lookifdptr(t, Tphotometric);
	if(!i)
		return -1;
	i->sp[0] = Prgb;

	i = lookifdptr(t, Tbitspersample);
	if(!i)
		return -1;
	i->sp = realloc(i->sp, 3*sizeof *i->sp);
	for(j = 0; j < 3; j++)
		i->sp[j] = bpc;
	i = lookifdptr(t, Tsamplespp);
	if(!i)
		return -1;
	i->sp[0] = 3;
	t->bpp = bpc*3;
	return 0;
}

int
tiffcmap(Tiff *t)
{
	uchar m[3*256], *map, *p, *q, *e;
	long i, l;
	ushort *s;
	IFD *d;

	if(lookifd(t, Tphotometric) != Prgbcmap)
		return 0;
	
	d = lookifdptr(t, Tcolormap);
	if(!d)
		return -1;
	s = (ushort*)d->sp;

	for(i = 0; i < 3*256; i++) {
		m[i] = s[i]>>8;
		if(s[i]&256 >= 128)
			m[i]++;
	}

	l = tiffimglen(t);
	map = malloc(3*l);
	if(!map)
		return -1;

//	for(i = 0; i < 256; i++)
//		if(m[i] || m[i+256] || m[i+512])
//			fprint(2, "%02x %02x %02x\n", m[i], m[i+256], m[i+512]);

	q = t->rawimg;
	e = t->rawimge;
	for(p = map; q < e; q++){
		*p++ = m[*q+0*256];
		*p++ = m[*q+1*256];
		*p++ = m[*q+2*256];
	}
	free(t->rawimg);
	t->rawimg = map;
	t->rawimge = map+3*l;
	return tiffsetrgb(t, 8);
}

int
tiffphoto(Tiff *t)
{
	ulong *u, *e;
	IFD *i;

	if((i = lookifdptr(t, Tphotometric)) == 0)
		return -1;

	switch(i->sp[0]){
	case Pblackzero:
	case Prgb:
	case Prgba:
		return 0;
	case Pwhitezero:
		e = (ulong*)t->rawimge;
		for(u = (ulong*)t->rawimg; u < e; u++)
			*u = ~*u;
		i->sp[0] = Pblackzero;
		return 0;
	case Prgbcmap:
		return tiffcmap(t);
	case Pcmyk:
	case PYCbCr:
	case PCIELab:
	default:
		return -1;
	}
}

static void
dumprawstrip(Tiff *t, Biobuf *b, int x)
{
	int o, c, i;

	o = ifdidx(lookifdptr(t, Tstripoffsets), x);
	c = ifdidx(lookifdptr(t, Tstripbytecounts), x);
	if(Bseek(b, o, 0) == -1)
		return;
	fprint(2, "%d\n", c);
	for(i = 0; i < c; i++){
		fprint(2, "%02x ", Bgetc(b));
		if((i+1)%16 == 0)
			fprint(2, "\n");
	}
}

Tiff*
tiffhdr(Biobuf *b)
{
	Tiff *t;
	char buf[3];

	t = mallocz(sizeof *t, 1);
	if(!t)
		return 0;

	buf[2] = 0;
	if(Bread(b, buf, 2) != 2)
		goto cleanup;

	if(strcmp(buf, "MM") == 0)
		t->order = BE;
	else if (strcmp(buf, "II") == 0)
		t->order = LE;
	else
		goto cleanup;

	if(42 != Btiffshort(b, t->order))
		goto cleanup;
	t->ifd0 = Btifflong(b, t->order);

	if(readifd(b, t->ifd0, t) == -1)
		goto cleanup;
	return t;
cleanup:
	tifffree(t);
	return 0;

}

Tiff*
readtiff(Biobuf *b)
{
	Tiff *t;
	int f;

	t = tiffhdr(b);
	if(!t)
		return 0;
	// dumprawstrip(t, b, 1); exits("");
	f = lookifd(t, Tfillorder);			// check for backwards bytes.
	if(f < 0 || f != 1)
	if(tiffimg(b, t) != -1)
	if(tiffexpand(t) != -1)
	if(tiffphoto(t) != -1)
		return t;
	tifffree(t);
	return 0;
}

void
tiffinfo(Tiff *t)
{
	debug("t.order = %s\n", t->order == BE ? "BE" : "LE");
	debug("t.ifd0 = %uld\n", t->ifd0);
	ifdprint(t);
}

void
tiffprint(Biobuf *b)
{
	Tiff *t;

	if(t = tiffhdr(b)){
		tiffinfo(t);
		tifffree(t);
	}
}

void
tiffdisplay(Biobuf *b)
{
	Tiff *t;

	t = readtiff(b);
	if(!t){
		fprint(2, "can't read tiff\n");
		return;
	}
	if(lookifd(t, Tsamplespp) == 3 && lookifd(t, Tphotometric) == Prgb)
		dirty(t);
	else if(lookifd(t, Tsamplespp) == 1 && lookifd(t, Tphotometric) == Pblackzero)
		dirty(t);
	tifffree(t);
}

void
tiffwrite(Biobuf *b)
{
	Tiff *t;
	char buf[20];

	t = readtiff(b);
	if(!t){
		fprint(2, "can't read tiff\n");
		return;
	}
	print("%11s %11d %11d %11ld %11ld ", tiffchantostr(t, buf),
		0, 0, lookifd(t, Twidth), lookifd(t, Tlength));
	if(write(1, t->rawimg, t->rawimge-t->rawimg) != t->rawimge-t->rawimg)
		fprint(2, "tiff: write error %r\n");
	tifffree(t);
}

int mainstacksize = 2*64*1024;
void
threadmain(int argc, char **argv)
{
	Biobuf *b;
	int r;
	void (*f)(Biobuf*);

	f = tiffdisplay;

	ARGBEGIN{
	case 'r':
		rflag ^= 1;
		break;
	case 'd':
		f = tiffdisplay;
		break;
	case 'p':
		f = tiffprint;
		break;
	case '9':
		f = tiffwrite;
		break;
	default:
		fprint(2, "usage: tiff -[rdp] [files]\n");
		threadexitsall("usage");
	}ARGEND;

	if(!*argv){
		*argv = "/fd/0";
		argv[1] = 0;
	}

	for(r = 0; *argv; argv++){
		b = Bopen(*argv, OREAD);
		if(!b){
			fprint(2, "%r\n");
			r = -1;
			continue;
		}
		f(b);
		Bterm(b);
	}
	threadexitsall(r == -1 ? "bad tiff" : 0);
}