#include "../UberLame_src/NewFix.h"
#include "../UberLame_src/CallStack.h"
#include <stdio.h>
#include <vector>
#include <map>
#include <stdexcept>
#include "../UberLame_src/Integer.h"
#include "../UberLame_src/Dir.h"
#include "../UberLame_src/Thread.h"
#include "../UberLame_src/Timer.h"
#include "../UberLame_src/iface/PNGLoad.h"
#include "../UberLame_src/StlUtils.h"
#include "../UberLame_src/StdIOUtils.h"
#include "DPX.h"
#include "InFocus.h"

#if defined(_WIN32) || defined(_WIN64)
#define strncasecmp(a,b,n) _strnicmp((a), (b), (n))
#define strcasecmp(a,b) _stricmp((a), (b))
#endif //_WIN32 || _WIN64

template <const int n_channel_num, class _TyInputData>
void ColorspaceConvert(const _TyInputData *p_channels, size_t n_pixel_num, int n_shift, uint32_t *p_dest)
{
	enum {
		shift_Right = sizeof(_TyInputData) > 1
	};

	for(size_t i = 0; i < n_pixel_num; ++ i, p_channels += n_channel_num, ++ p_dest) {
		if(n_channel_num == 1) {
			int n_gray = (shift_Right)? *p_channels >> n_shift : *p_channels << n_shift;
			*p_dest = 0xff000000U | n_gray | (n_gray << 8) | (n_gray << 16);
		} else if(n_channel_num == 3) {
			int n_r = (shift_Right)? p_channels[0] >> n_shift : p_channels[0] << n_shift;
			int n_g = (shift_Right)? p_channels[1] >> n_shift : p_channels[1] << n_shift;
			int n_b = (shift_Right)? p_channels[2] >> n_shift : p_channels[2] << n_shift;
			*p_dest = 0xff000000U | n_b | (n_g << 8) | (n_r << 16);
		} else if(n_channel_num == 4) {
			int n_r = (shift_Right)? p_channels[0] >> n_shift : p_channels[0] << n_shift;
			int n_g = (shift_Right)? p_channels[1] >> n_shift : p_channels[1] << n_shift;
			int n_b = (shift_Right)? p_channels[2] >> n_shift : p_channels[2] << n_shift;
			int n_a = (shift_Right)? p_channels[3] >> n_shift : p_channels[3] << n_shift;
			*p_dest = (n_a << 24) | n_b | (n_g << 8) | (n_r << 16);
		}
	}
}

static void PrintShortFilename(const char *p_s_filename, int n_max_length = 80,
	const char *p_s_text_before = "", const char *p_s_text_after = "")
{
	std::pair<const char*, const char*> t_short =
		CPath::t_ShortFileName(p_s_filename, n_max_length);
	printf("%s%s%s%s", p_s_text_before, t_short.first, t_short.second, p_s_text_after);
	// print filename, omit full path if too long
}

static bool b_verbose = true, b_no_warn = false;

bool ProcessSingle(const char *p_s_filename, const char *p_s_out_dir, float &r_f_sharpness)
{
	/*if(r_t_file.b_directory || strcasecmp(r_t_file.p_s_Extension(), "dpx"))
		return true;*/
	_ASSERTE(!TFileInfo(p_s_filename).b_exists || (!TFileInfo(p_s_filename).b_directory &&
		!strcasecmp(TFileInfo::p_s_Extension(p_s_filename), "dpx"))); // if it does not exist, do not confuse the user with this
	// only .dpx files

	//printf("converting \'%s\' ...\n", p_s_filename);
	if(b_verbose)
		PrintShortFilename(p_s_filename, 80, "processing \'", "\' ...\n");
	// verbose

	InStream i;
	if(!i.Open(p_s_filename)) {
		fprintf(stderr, "error: failed to open \'%s\'\n", p_s_filename);
		return false;
	}
	dpx::Reader reader;
	reader.SetInStream(&i);
	if(!reader.ReadHeader()) {
		fprintf(stderr, "error: failed to read header of \'%s\'\n", p_s_filename);
		return false;
	}
	int n_height = reader.header.Height();
	int n_width = reader.header.Width();
	_ASSERTE(reader.header.NumberOfElements() == 1); // one element per image?
	dpx::Characteristic colori = reader.header.Colorimetric(0);
	dpx::Characteristic transf = reader.header.Transfer(0);
	int n_channel_num = reader.header.ImageElementComponentCount(0);
	int n_bytes_per_channel = reader.header.ComponentByteCount(0);
	int n_bytes_per_pixel = n_channel_num * n_bytes_per_channel;
	//printf("it is an %d x %d image, %d channels, %d B/pixel\n", n_width, n_height, n_channel_num, n_bytes_per_pixel);
	size_t n_size = n_width * n_height * n_bytes_per_pixel;
	std::vector<uint8_t> image_buffer(n_size);
	if(!reader.ReadImage(0, &image_buffer[0])) {
		fprintf(stderr, "error: image decoding failed with \'%s\'\n", p_s_filename);
		return false;
	}
	i.Close();
	// read the image

	if(!b_no_warn) {
		float f_gamma = reader.header.Gamma();
		if(fabs(f_gamma - 2.2) > 1e-3)
			fprintf(stderr, "warning: \'%s\' has gamma of %g (assumed 2.2)\n", p_s_filename, f_gamma);
#ifdef DetectFromLog
		if(colori != dpx::kLogarithmic ||
		   transf != dpx::kLogarithmic)
			fprintf(stderr, "warning: \'%s\' appears to be not in logarithmic space\n", p_s_filename);
		// the in-focus estimator is configured to work with logarithmic images
#else // DetectFromLog
		if(colori != dpx::kLinear ||
		   transf != dpx::kLinear)
			fprintf(stderr, "warning: \'%s\' appears to be not in linear space\n", p_s_filename);
		// the in-focus estimator is configured to work with linear images
#endif // DetectFromLog
	}
	// image checks (DetectFromLog may be defined in InFocus.h, line 10)

	TBmp *p_bitmap = TBmp::p_Alloc(n_width, n_height, false, false, 8/*reader.header.BitDepth(0)*/);
	if(!p_bitmap)
		throw std::bad_alloc(); // rethrow

	//int n_shift = abs(8 - reader.header.BitDepth(0)); 
	int n_shift = abs(8 - n_bytes_per_channel * 8); // it is MSB aligned
	if(n_bytes_per_channel == 2) {
		if(n_channel_num == 1)
			ColorspaceConvert<1>((uint16_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else if(n_channel_num == 3)
			ColorspaceConvert<3>((uint16_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else if(n_channel_num == 4)
			ColorspaceConvert<4>((uint16_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else
			return false;
	} else if(n_bytes_per_channel == 1) {
		if(n_channel_num == 1)
			ColorspaceConvert<1>((uint8_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else if(n_channel_num == 3)
			ColorspaceConvert<3>((uint8_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else if(n_channel_num == 4)
			ColorspaceConvert<4>((uint8_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else
			return false;
	} else if(n_bytes_per_channel == 4) {
		if(n_channel_num == 1)
			ColorspaceConvert<1>((uint32_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else if(n_channel_num == 3)
			ColorspaceConvert<3>((uint32_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else if(n_channel_num == 4)
			ColorspaceConvert<4>((uint32_t*)&image_buffer[0], n_width * n_height, n_shift, p_bitmap->p_buffer);
		else
			return false;
	} else
		return false;
	// colorspace convert

	//printf("converting to universal bmp ...\n");

	TUniversalBmpByte SourceImg;
	SourceImg.New(p_bitmap->n_width, p_bitmap->n_height, 3);
	{
		uint8_t *P = SourceImg.Pixel(0, 0);
		const uint32_t *p_src_pixel = p_bitmap->p_buffer;
		for(int x = 0; x < SourceImg.MaxX; ++ x) {
			for(int y = 0; y < SourceImg.MaxY; ++ y, P += 3, ++ p_src_pixel) {
				P[2] = *p_src_pixel & 0xff;
				P[1] = (*p_src_pixel >> 8) & 0xff;
				P[0] = (*p_src_pixel >> 16) & 0xff;
			}
		}
	}
	// convert to universal bmp

	//printf("running ...\n");

	TUniversalBmpByte DestImg;
	float f_sharpness = ComputeSharpness(DestImg, SourceImg);
	// compute sharpness

	//printf("sharpness(\'%s\'): %f\n", p_s_filename, f_sharpness);
	r_f_sharpness = f_sharpness;

	{
		const uint8_t *P = DestImg.Pixel(0, 0);
		uint32_t *p_dst_pixel = p_bitmap->p_buffer;
		for(int x = 0; x < DestImg.MaxX; ++ x) {
			for(int y = 0; y < DestImg.MaxY; ++ y, P += 3, ++ p_dst_pixel)
				*p_dst_pixel = P[2] | (P[1] << 8) | (P[0] << 16) | 0xff000000U;
		}
	}
	// convert from universal bmp

	std::string s_filename = TFileInfo::p_s_FileName(p_s_filename);
	s_filename.erase(s_filename.rfind('.'));
	s_filename = p_s_out_dir + ("/" + s_filename) + ".png";

	if(!CPngCodec::Save_PNG(s_filename.c_str(), *p_bitmap, true)) {
		fprintf(stderr, "error: failed to write \'%s\'\n", s_filename.c_str());
		return false;
	}
	// save as png

	p_bitmap->Delete();

	return true;
}

class CThreadedProcessor {
protected:
	class CWorker : public CRunable_Thread_ShallowCopy { // working thread
	protected:
		CProducerConsumerQueue<std::string> *m_p_queue; // pointer to the queue
		const char *m_p_s_dest_path; // destination path
		CThreadedProcessor *m_p_parent;

	public:
		inline CWorker(CProducerConsumerQueue<std::string> &r_queue,
			const char *p_s_dest_path, CThreadedProcessor *p_parent) // default constructor
			:m_p_queue(&r_queue), m_p_s_dest_path(p_s_dest_path), m_p_parent(p_parent)
		{}

	protected:
		virtual void Run() // thread execution function
		{
			for(;;) {
				std::string s_file;
				if(!m_p_queue->Get(s_file)) { // get a next file
					if(m_p_queue->b_Finished()) // is ther no more?
						return; // finished
					throw std::runtime_error("work item queue get failed");
				}
				float f_sharpness;
				try {
					ProcessSingle(s_file.c_str(), m_p_s_dest_path, f_sharpness); // convert the file using the serial function (is reentrant)
					m_p_parent->SetSharpness(s_file, f_sharpness);
				} catch(std::bad_alloc&) {
					fprintf(stderr, "error: not enough memory\n");
				}
			}
		}
	};

	CProducerConsumerQueue<std::string> m_queue; // queue for passing work-items (filenames)
	std::vector<CWorker> m_workers; // vector of workers
	const char *m_p_s_dest_path; // destination path
	std::map<std::string, float> m_sharpness; // sharpness per file
	CMutex m_sharpness_mutex; // mutex for the sharpness
	CTimer m_timer;

public:
	CThreadedProcessor(const char *p_s_dest_path, size_t n_thread_num = CThread::n_CPU_Num())
		:m_queue(std::max(n_thread_num * 4, size_t(64))), // space for 64 files to be queued
		m_p_s_dest_path(p_s_dest_path)
	{
		m_workers.resize(n_thread_num, CWorker(m_queue, m_p_s_dest_path, this)); // as many workers as there are CPUs
	}

	bool StartWorkers() // start all the workers (could be done in constructor)
	{
		for(size_t i = 0, n = m_workers.size(); i < n; ++ i) {
			if(!m_workers[i].Start())
				return false;
		}
		return true;
	}

	bool WaitForFinish() // wait for all the threads to finish
	{
		if(!m_queue.Signal_Finished())
			throw std::runtime_error("work item queue signal finished failed");
		for(size_t i = 0, n = m_workers.size(); i < n; ++ i) {
			if(!m_workers[i].WaitForFinish())
				return false;
			m_timer.f_Time();
			// sample the timer from time to time to avoid overflow
		}
		return true;
	}

	bool FlushSharpness()
	{
		std::string s_filename = m_p_s_dest_path;
		s_filename += "/sharpness.xml";
		FILE *p_fw;
		if(!(p_fw = fopen(s_filename.c_str(), "w")))
			return false;
		CFILE_PtrGuard guard(p_fw); // closes file automatically

		fprintf(p_fw, "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"yes\"?>\n");
		// note that the encoding of the filesystem may not be utf-8; we could perform conversion if needed,
		// now we just rely on the filenames being us-english or utf-á

		fprintf(p_fw, "<sharpness>\n");
		for(std::map<std::string, float>::const_iterator p_sh_it = m_sharpness.begin(),
		   p_end_it = m_sharpness.end(); p_sh_it != p_end_it; ++ p_sh_it) {
			fprintf(p_fw, " \t<file sharpness=\"%f\">%s</file>\n",
				(*p_sh_it).second, (*p_sh_it).first.c_str());
		}
		fprintf(p_fw, "</sharpness>\n");
		return !ferror(p_fw);
	}

	void SetSharpness(const std::string &r_s_filename, float f_sharpness)
	{
		m_sharpness_mutex.Lock();
		m_sharpness[r_s_filename] = f_sharpness; // throws
		m_sharpness_mutex.Unlock();
	}

	inline bool operator ()(const TFileInfo &r_t_file) // file callback
	{
		if(r_t_file.b_directory || strcasecmp(r_t_file.p_s_Extension(), "dpx"))
			return true;
		// only dpx files

		m_timer.f_Time();
		// sample the timer from time to time to avoid overflow

		//printf("callback \'%s\'\n", r_t_file.s_filename.c_str());
		return m_queue.Put(std::string(r_t_file.p_s_Path())); // just put it in the queue
	}

	inline bool operator ()(const char *p_s_filename) // file callback
	{
		m_timer.f_Time();
		// sample the timer from time to time to avoid overflow

		return m_queue.Put(std::string(p_s_filename)); // just put it in the queue
	}

	inline double f_Time() // returns time elapsed
	{
		return m_timer.f_Time();
	}
};

void PrintHelp()
{
	fprintf(stderr, 
		"use: inFocus <input file.dpx or folder or file list> <output folder> [-t <threads>]\n"
		"\t[-q|--quiet] [-nw|--no-warn]\n\n"
		"This program can only decode dpx files as input, which need to be in linear space\n"
		"The input files are saved in the output folder as png. There is also sharpness.xml\n"
		"saved in the output folder, with the in.focuss quantity per file, in lexicographic\n"
		"order (i.e. potentially different than specified on input).\n"
		"The input file list is a text file with a single file path per line,\n"
		"comments may be specified using hte hash \'#\' character.\n"
		"This process uses as many worker threads as there are logical cores,\n"
		"unless overriden by the -t switch\n");
}

int main(int n_arg_num, const char **p_arg_list)
{
	if(n_arg_num < 3) {
		PrintHelp();
		return -1;
	}
	int n_thread_num = CThread::n_CPU_Num(); // default
	for(int i = 3; i < n_arg_num; ++ i) {
		if(!strcmp(p_arg_list[i], "-h") || !strcmp(p_arg_list[i], "--help")) {
			PrintHelp();
			return 0;
		} else if(!strcmp(p_arg_list[i], "-q") || !strcmp(p_arg_list[i], "--quiet"))
			b_verbose = false;
		 else if(!strcmp(p_arg_list[i], "-nw") || !strcmp(p_arg_list[i], "--no-warn"))
			b_no_warn = true;
		else if(i + 1 == n_arg_num) {
			fprintf(stderr, "error: argument \'%s\': missing value or an unknown argument\n", p_arg_list[i]);
			return -1;
		} else if(!strcmp(p_arg_list[i], "-t"))
			n_thread_num = std::max(1, atoi(p_arg_list[++ i]));
		else {
			fprintf(stderr, "error: argument \'%s\': an unknown argument\n", p_arg_list[i]);
			return -1;
		}
	}
	const char *p_s_infile = p_arg_list[1];
	const char *p_s_dest_path = p_arg_list[2];
	// parse commandline

	CTimer t;
	try {
		TFileInfo t_file(p_s_infile);
		//printf("input file \'%s\': %d %d %d\n", p_s_infile, t_file.b_Valid(), t_file.b_exists, t_file.b_directory);
		// input file

		CThreadedProcessor runner(p_s_dest_path, n_thread_num);
		if(!runner.StartWorkers()) {
			fprintf(stderr, "error: threading failed\n");
			return -1;
		}
		if(t_file.b_directory) {
			if(!CDirTraversal::Traverse2(p_s_infile, runner, false)) {
				fprintf(stderr, "error: traversal failed\n");
				return -1;
			}
		} else {
			if(!strcasecmp(TFileInfo::p_s_Extension(p_s_infile), "dpx")) {
				//printf("converting a single file\n");
				if(!runner(p_s_infile)) {
					fprintf(stderr, "error: work push failed\n");
					return -1;
				}
			} else {
				FILE *p_fr;
				if(!(p_fr = fopen(p_s_infile, "r"))) {
					fprintf(stderr, "error: failed to open list file \'%s\'\n", p_s_infile);
					return -1;
				}
				CFILE_PtrGuard guard(p_fr); // takes care of closing the file properly
				std::string s_line;
				while(!feof(p_fr)) {
					if(!stl_ut::ReadLine(s_line, p_fr)) {
						fprintf(stderr, "error: failed to read list file \'%s\'\n", p_s_infile);
						return -1;
					}
					size_t n_comment;
					if((n_comment = s_line.find('#')) != std::string::npos)
						s_line.erase(n_comment);
					stl_ut::TrimSpace(s_line);
					if(s_line.empty())
						continue;
					// read the input file, line by line

					if(!runner(s_line.c_str())) {
						fprintf(stderr, "error: work push failed\n");
						return -1;
					}
					// send the file for processing
				}
			}
		}
		if(!runner.WaitForFinish()) {
			fprintf(stderr, "error: threading failed\n");
			return -1;
		}
		if(!runner.FlushSharpness()) {
			fprintf(stderr, "error: failed to write the xml file\n");
			return -1;
		}
		double f_time = runner.f_Time();
		if(b_verbose)
			printf("done. it took " PRItime "\n", PRItimeparams(f_time));
	} catch(std::exception &r_exc) {
		fprintf(stderr, "error: uncaught exception: \'%s\'\n", r_exc.what());
		return -1;
	}

	return 0;
}
