/***************************************************************************
       cfilehasher.cpp  -  Calculate the TTH root and leaves for a file
                             -------------------
    begin                : Fri May 16 2008
    copyright            : (C) 2008 by Edward Sheldrake
    email                : ejs1920@yahoo.co.uk
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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.                                   *
 *                                                                         *
 ***************************************************************************/

#include "cfilehasher.h"

#include "core/cbytearray.h"
#include "core/cbase32.h"
#include "core/cdir.h"

#include "hash/compat.h"
#include "hash/MerkleTree.h"

#include <string.h> //for memcmp()

/** */
CTreeVerificationReport::CTreeVerificationReport( const CTreeVerificationReport & other )
{
	filename = other.filename;
	tthRoot  = other.tthRoot;
	filesize = other.filesize;
	allgood  = other.allgood;
	
	if ( other.segments )
	{
		segments = new CList<CHashedSegment>();
		
		CHashedSegment * hs = 0;
		
		/* FIXME stop casting away the const once CList is fixed / replaced */
		while ( (hs = ((CTreeVerificationReport&)other).segments->Next(hs)) != 0 )
		{
			segments->Add( new CHashedSegment(*hs) );
		}
	}
	else
	{
		segments = 0;
	}
}

/** */
CString CTreeVerificationReport::ToString() const
{
	CString s = "CTreeVerificationReport\nFileName: ";
	s += filename;
	s += "\nFileSize: ";
	s += CString::number(filesize);
	s += "\nActual TTH root: ";
	s += tthRoot;
	s += "\nAll hashes OK: ";
	if ( allgood )
	{
		s += "YES";
	}
	else
	{
		s += "NO";
	}
	s += "\nNumber of segments: ";
	s += CString::number(segments->Count());
	//        18446744073709551615 18446744073709551615 ORQ6K7F5U2QGJOUKEKPEDJ2XXM3FTVYQ2BLWQOI ORQ6K7F5U2QGJOUKEKPEDJ2XXM3FTVYQ2BLWQOI
	s += "\n\nStart                Size                 Expected TTH of block                   Actual TTH of block\n";
	
	CHashedSegment * segment = 0;
	
	while ( (segment = segments->Next(segment)) != 0 )
	{
		s += CString::number(segment->start).RightJustify(20);
		s += ' ';
		s += CString::number(segment->size).RightJustify(20);
		s += ' ';
		s += segment->expected;
		s += ' ';
		s += segment->actual;
		
		if ( segment->expected != segment->actual )
		{
			s += " *** MISMATCH ***";
		}
		
		s += "\n";
	}
	
	return s;
}

/** */
CFileHasher::CFileHasher( const CString filename, CByteArray * workmem )
{
	status      = efhsNotStarted;
	m_bStop     = false;
	m_pRootData = 0;
	m_pLeafData = 0;
	filesize    = 0;
	m_nProgress = 0;
	
	if ( file.Open( filename, IO_RAW | IO_READONLY ) )
	{
		m_pWorkMem = workmem;
		usingOwnMem = false;
		filesize = CDir().getFileSize( filename, false );
		
		if ( m_pWorkMem == 0 )
		{
			m_pWorkMem = new CByteArray( 16*1024 );
			usingOwnMem = true;
		}
		
		if ( m_pWorkMem == 0 )
		{
			status = efhsError;
		}
		else
		{
			status = efhsReady;
		}
	}
	else
	{
		m_pWorkMem = 0;
		status = efhsError;
	}
}

/** delete stuff, close file */
CFileHasher::~CFileHasher()
{
	delete m_pRootData;
	m_pRootData = 0;
	
	delete m_pLeafData;
	m_pLeafData = 0;
	
	if ( usingOwnMem && (m_pWorkMem != 0) )
	{
		delete m_pWorkMem;
		m_pWorkMem = 0;
	}
	
	if ( file.IsOpen() )
	{
		file.Close();
	}
}

/** */
unsigned long CFileHasher::HashSize()
{
	return dcpp::TigerTree::BYTES;
}

/** */
CString CFileHasher::GetHashRoot()
{
	CString s;
	
	if ( (status == efhsFinished) && m_pRootData && (m_pRootData->Size() == dcpp::TigerTree::BYTES) )
	{
		CBase32::Encode( &s, m_pRootData );
	}
	
	return s;
}

/** */
CByteArray * CFileHasher::GetHashRootRaw()
{
	CByteArray * result = 0;
	
	if ( (status == efhsFinished) && m_pRootData && (m_pRootData->Size() == dcpp::TigerTree::BYTES) )
	{
		result = new CByteArray();
		result->Append( m_pRootData->Data(), m_pRootData->Size() );
	}
	
	return result;
}

/** */
CByteArray * CFileHasher::GetLeafData()
{
	CByteArray * result = 0;
	
	if ( (status == efhsFinished) && m_pLeafData )
	{
		result = new CByteArray();
		result->Append( m_pLeafData->Data(), m_pLeafData->Size() );
	}
	
	return result;
}

/** */
void CFileHasher::StopHashing()
{
	m_bStop = true;
}

/** */
void CFileHasher::ComputeHash( const eFileHasherStatus endStatus )
{
	if ( status == efhsReady )
	{
		status = efhsWorking;
		
		dcpp::TigerTree hasher( std::max(dcpp::TigerTree::calcBlockSize(filesize,10), ((int64_t) 64*1024) ) );
		
		long len;
		
		while ( ((len = file.Read( (char*)m_pWorkMem->Data(), m_pWorkMem->Size() )) > 0) && (m_bStop == false) )
		{
			hasher.update(m_pWorkMem->Data(),len);

			m_nProgress += len;
		}
		
		file.Close();
		
		if ( m_bStop )
		{
			status = efhsError;
		}
		else
		{
			hasher.finalize();
			
			m_pRootData = new CByteArray();
			m_pRootData->Append( hasher.getRoot().data, dcpp::TigerTree::BYTES );
			
			vector<uint8_t> leafdata = hasher.getLeafData();
			m_pLeafData = new CByteArray();
			m_pLeafData->Append( &leafdata[0], leafdata.size() );
		
			status = endStatus;
		}
		
		if ( usingOwnMem )
		{
			delete m_pWorkMem;
			m_pWorkMem = 0;
		}
	}
}

/** from DownloadManager::endData in the DC++ source */
int64_t CFileHasher::GetBlockSize( const unsigned long leavesSize, const int64_t filesize )
{
	int64_t blocksize = 1024;
	while ( blocksize * leavesSize/sizeof(dcpp::TigerTree::MerkleValue) < filesize )
	{
		blocksize = blocksize * 2;
	}
	return blocksize;
}

/** */
bool CFileHasher::ValidateHashLeaves( CString tth, CByteArray * leaves, const ulonglong filesize )
{
	CByteArray dst;
	
	if ( CBase32::Decode( &dst, &tth ) != dcpp::TigerTree::BYTES )
	{
		return false;
	}
	else
	{
		return ValidateHashLeaves( &dst, leaves, filesize );
	}
}

/** */
bool CFileHasher::ValidateHashLeaves( CByteArray * root, CByteArray * leaves, const ulonglong filesize )
{
	if ( (root == 0) || (leaves == 0) )
	{
		return false;
	}
	
	if ( root->Size() != dcpp::TigerTree::BYTES )
	{
		return false;
	}
	else if ( leaves->Size() < dcpp::TigerTree::BYTES )
	{
		return false;
	}
	else if ( leaves->Size() == dcpp::TigerTree::BYTES ) // optimisation
	{
		if ( memcmp(root->Data(),leaves->Data(),dcpp::TigerTree::BYTES) == 0 )
		{
			return true;
		}
	}
	
	int64_t size = (int64_t) filesize;
	
	if ( size < 0 )
	{
		return false;
	}
	
	int64_t blocksize = GetBlockSize( leaves->Size(), size );
	
	dcpp::TigerTree hasher(size,blocksize,leaves->Data());
	
	if ( memcmp(hasher.getRoot().data,root->Data(),dcpp::TigerTree::BYTES) == 0 )
	{
		return true;
	}
	else
	{
		return false;
	}
}

/** */
CByteArray * CFileHasher::HashByteArray( CByteArray * data, unsigned long length )
{
	CByteArray * result = new CByteArray();
	
	dcpp::TigerTree hasher;
	hasher.update( data->Data(), length );
	hasher.finalize();
	result->Append( hasher.getRoot().data, dcpp::TigerTree::BYTES );
	
	return result;
}

/** */
void CFileHasherThread::Thread()
{
	Lock();
	
	ComputeHash();
	
	UnLock();
	
	/* do not try to pthread_join ourself */
	Stop(false);
}

/** */
CFileTreeVerifier::CFileTreeVerifier( const CString filename, CByteArray * leaves, CByteArray * workmem ) : CFileHasher ( filename, workmem )
{
	m_pLeaves = leaves;
	m_pReport = new CTreeVerificationReport();
	m_pReport->filename = filename;
	m_nPass = 0;
}

/** */
CFileTreeVerifier::~CFileTreeVerifier()
{
	delete m_pReport;
	m_pReport = 0;
}

/** */
CTreeVerificationReport * CFileTreeVerifier::GetReport()
{
	CTreeVerificationReport * pointer = 0;
	
	if ( GetStatus() == efhsFinished )
	{
		pointer = m_pReport;
		m_pReport = 0;
	}
	
	return pointer;
}

/** */
void CFileTreeVerifier::Thread()
{
	Lock();
	
	m_nPass = 1;
	ComputeHash( efhsWorking );
	
	if ( status == efhsWorking )
	{
		m_nPass = 2;
		m_nProgress = 0;
		bool ok = true;
		
		/* first get segment size, we need it for either method */
		int leafcount = m_pLeaves->Size() / dcpp::TigerTree::BYTES;
		unsigned long blocksize = 1024;
		while ( blocksize * leafcount < filesize )
		{
			blocksize = blocksize * 2;
		}
			
		/* Can we do this? Or could two lists of tth leaves the same length be from differently structured trees? */
		/* This should not take to long, no need to check m_bStop */
		if ( m_pLeaves->Size() == HashLeavesDirect()->Size() )
		{
			CByteArray expected_in;
			CByteArray actual_in;
			
			for ( unsigned long i = 0; i < m_pLeaves->Size(); i = i + dcpp::TigerTree::BYTES )
			{
				expected_in.SetSize(0);
				actual_in.SetSize(0);
				
				expected_in.Append( m_pLeaves->Data()+i, dcpp::TigerTree::BYTES );
				actual_in.Append( HashLeavesDirect()->Data()+i, dcpp::TigerTree::BYTES );
				
				CHashedSegment * segment = new CHashedSegment();
				
				CBase32::Encode( &(segment->expected), &expected_in );
				CBase32::Encode( &(segment->actual), &actual_in );
				
				if ( memcmp(expected_in.Data(),actual_in.Data(),dcpp::TigerTree::BYTES) != 0 )
				{
					ok = false;
				}
				
				segment->start = i * blocksize;
				segment->size = blocksize;
				if ( segment->start + segment->size > filesize )
				{
					segment->size = filesize - segment->start;
				}
				
				m_pReport->segments->Add(segment);
			}
			
			m_pReport->tthRoot  = GetHashRoot();
			m_pReport->filesize = filesize;
			m_pReport->allgood  = ok;
			status = efhsFinished;
		}
		else
		{
			/* Otherwise, we need to hash each segement of the file */
			if ( file.Open( m_pReport->filename, IO_RAW | IO_READONLY ) )
			{
				CByteArray expected_in;
				
				for ( int i = 0; i < leafcount; i++ )
				{
					expected_in.SetSize(0);
					expected_in.Append( m_pLeaves->Data()+(i*dcpp::TigerTree::BYTES), dcpp::TigerTree::BYTES );
					
					CByteArray * buffer = new CByteArray( blocksize );
					
					long read = file.Read( (char*) buffer->Data(), blocksize );
					if ( read > 0 )
					{
						CByteArray * actual_in = HashByteArray( buffer, (unsigned long)read );
						
						CHashedSegment * segment = new CHashedSegment();
						
						CBase32::Encode( &(segment->expected), &expected_in );
						CBase32::Encode( &(segment->actual), actual_in );
						
						delete actual_in;
						
						if ( segment->expected != segment->actual )
						{
							ok = false;
						}
						
						segment->start = m_nProgress;
						segment->size = read;
						
						m_pReport->segments->Add(segment);
						
						m_nProgress += read;
					}
					else
					{
						status = efhsError;
					}
	
					delete buffer;

					if ( m_bStop || (status == efhsError) )
					{
						status = efhsError;
						break;
					}
					
				}
				
				file.Close();
				if ( status == efhsWorking )
				{
					status = efhsFinished;
					m_pReport->tthRoot  = GetHashRoot();
					m_pReport->filesize = filesize;
					m_pReport->allgood  = ok;
				}
			}
			else
			{
				status = efhsError;
			}
		}
	}
	
	UnLock();
	
	/* do not try to pthread_join ourself */
	Stop(false);
}
