// Copyright (C) 2014-2015 ChaosForge Ltd
// http://chaosforge.org/
//
// This file is part of Nova libraries. 
// For conditions of distribution and use, see copying.txt file in root folder.

#include "nv/formats/nmd_loader.hh"
#include "nv/io/std_stream.hh"
#include "nv/stl/string.hh"

using namespace nv;

bool nv::nmd_loader::load( stream& source )
{
	// TODO: proper error handling
	reset();
	nmd_header root_header;
	source.read( &root_header, sizeof( root_header ), 1 );
	for ( uint32 i = 0; i < root_header.elements; ++i )
	{
		nmd_element_header element_header;
		source.read( &element_header, sizeof( element_header ), 1 );
		switch ( element_header.type )
		{
		case nmd_type::MESH           : load_mesh( source, element_header ); break;
		case nmd_type::ANIMATION      : load_animation( source, element_header ); break;
		case nmd_type::STRING_TABLE   : load_strings( source ); break;
		default: NV_ASSERT( false, "UNKNOWN NMD ELEMENT!" ); break;
		}
	}
	return true;
}

bool nv::nmd_loader::load_mesh( stream& source, const nmd_element_header& e )
{
	mesh_data* mesh = new mesh_data();
	for ( uint32 s = 0; s < e.children; ++s )
	{
		nmd_element_header element_header;
		source.read( &element_header, sizeof( element_header ), 1 );
		NV_ASSERT( element_header.type == nmd_type::STREAM, "STREAM expected!" );

		nmd_stream_header stream_header;
		source.read( &stream_header, sizeof( stream_header ), 1 );
		raw_data_channel* channel = raw_data_channel::create( stream_header.format, stream_header.count );
		source.read( channel->data, stream_header.format.element_size(), stream_header.count );
		mesh->add_channel( channel );
	}
	m_mesh_names.push_back( e.name );
	m_meshes.push_back( mesh );
	return true;
}

mesh_data* nv::nmd_loader::release_mesh_data( size_t index )
{
	mesh_data* result = m_meshes[ index ];
	if ( m_strings ) result->set_name( m_strings->get( m_mesh_names[ index ] ) );
	m_meshes[ index ] = nullptr;
	return result;
}

mesh_data_pack* nv::nmd_loader::release_mesh_data_pack()
{
	uint32 size = m_meshes.size();
	mesh_data* meshes = new mesh_data[ size ];
	for ( uint32 i = 0; i < size; ++i )
	{
		m_meshes[i]->move_to( meshes[i] );
		delete m_meshes[i];
	}
	m_meshes.clear();
	return new mesh_data_pack( size, meshes, release_mesh_nodes_data() );
}

void nv::nmd_loader::reset()
{
	for ( auto mesh : m_meshes ) if ( mesh ) delete mesh;
	if ( m_strings )   delete m_strings;
	if ( m_node_data ) delete m_node_data;
	m_meshes.clear();
	m_mesh_names.clear();
	m_node_names.clear();

	m_node_data  = nullptr;
	m_node_array = nullptr;
	m_strings    = nullptr;
}

nv::nmd_loader::~nmd_loader()
{
	reset();
}

bool nv::nmd_loader::load_strings( stream& source )
{
	NV_ASSERT( m_strings == nullptr, "MULTIPLE STRING ENTRIES!" );
	m_strings = new string_table( &source );
	return true;
}

bool nv::nmd_loader::load_animation( stream& source, const nmd_element_header& e )
{
	NV_ASSERT( m_node_data == nullptr, "MULTIPLE NODE ENTRIES!" );
	nmd_animation_header animation_header;
	source.read( &animation_header, sizeof( animation_header ), 1 );
	m_node_array = new mesh_node_data[ e.children ];
	for ( uint32 i = 0; i < e.children; ++i )
	{
		nmd_element_header element_header;
		source.read( &element_header, sizeof( element_header ), 1 );
		NV_ASSERT( element_header.type == nmd_type::NODE, "NODE expected!" );
		m_node_names.push_back( element_header.name );
		uint16 ch_count = element_header.children;

		nmd_node_header node_header;
		source.read( &node_header, sizeof( node_header ), 1 );
		m_node_array[i].parent_id     = node_header.parent_id;
		m_node_array[i].transform     = node_header.transform;
		m_node_array[i].data          = nullptr;
		if ( ch_count > 0 )
		{
			key_data* kdata = new key_data;
			m_node_array[i].data = kdata;
			for ( uint32 c = 0; c < ch_count; ++c )
			{
				source.read( &element_header, sizeof( element_header ), 1 );
				NV_ASSERT( element_header.type == nmd_type::KEY_CHANNEL, "CHANNEL expected!" );
				nv::nmd_stream_header cheader;
				source.read( &cheader, sizeof( cheader ), 1 );
				raw_data_channel* channel = raw_data_channel::create( cheader.format, cheader.count );
				source.read( channel->data, channel->element_size(), channel->element_count() );
				kdata->add_channel( channel );
			}
		}
	}
	m_node_data = new mesh_nodes_data( "animation", e.children, m_node_array, animation_header.frame_rate, animation_header.duration, animation_header.flat );
	return true;
}

mesh_nodes_data* nv::nmd_loader::release_mesh_nodes_data( size_t )
{
	if ( m_node_data )
	{
		if ( m_strings )
		{
			for ( uint32 i = 0; i < m_node_data->get_count(); ++i )
			{
				m_node_array[i].name = m_strings->get( m_node_names[i] );
			}
		}
		mesh_nodes_data* result = m_node_data;
		m_node_data = nullptr;
		m_node_array = nullptr;
		return result;
	}
	return nullptr;
}

// ----------------------------------------------------------------
// nmd format dump
// HACK : TEMPORARY - will go to it's own file, probably nmd_io
static void nmd_dump_mesh( const mesh_data* mesh, stream& stream_out )
{
	array_view< raw_data_channel* > data  = mesh->get_raw_channels();

	uint32 size = sizeof( nmd_element_header );
	for ( auto chan : data )
	{
		size += sizeof( nmd_element_header ) + sizeof( nmd_stream_header );
		size += chan->raw_size();
	}

	nmd_element_header eheader;
	eheader.type     = nmd_type::MESH;
	eheader.name     = 0;
	eheader.children = static_cast< uint16 >( data.size() );
	eheader.size     = size;
	stream_out.write( &eheader, sizeof( eheader ), 1 );

	for ( auto chan : data )
	{
		nmd_element_header cheader;
		eheader.name     = 0;
		cheader.type     = nmd_type::STREAM;
		cheader.children = 0;
		cheader.size     = chan->raw_size() + sizeof( nmd_stream_header );
		stream_out.write( &cheader, sizeof( cheader ), 1 );

		nmd_stream_header sheader;
		sheader.format = chan->descriptor();
		sheader.count  = chan->element_count();
		stream_out.write( &sheader, sizeof( sheader ), 1 );
		stream_out.write( chan->raw_data(), chan->element_size(), chan->element_count() );
	}
}

static void nmd_dump_nodes_data( const mesh_nodes_data* nodes, stream& stream_out, string_table_creator* strings )
{
	uint32 total = sizeof( nmd_animation_header );
	for ( uint32 i = 0; i < nodes->get_count(); ++i )
	{
		const mesh_node_data* node = nodes->get_node(i);
		total += sizeof( nmd_element_header ) + sizeof( nmd_node_header );
		if ( node->data )
			for ( uint32 c = 0; c < node->data->get_channel_count(); ++c )
			{
				total += sizeof( nmd_element_header ) + sizeof( nmd_stream_header );
				total += node->data->get_channel(c)->raw_size();
			}
	}

	nmd_element_header header;
	header.name     = 0;
	header.type     = nmd_type::ANIMATION;
	header.children = static_cast< uint16 >( nodes->get_count() );
	header.size     = total;
	stream_out.write( &header, sizeof( header ), 1 );

	nmd_animation_header aheader;
	aheader.frame_rate  = nodes->get_frame_rate();
	aheader.duration    = nodes->get_duration();
	aheader.flat        = nodes->is_flat();
	stream_out.write( &aheader, sizeof( aheader ), 1 );

	for ( uint32 i = 0; i < nodes->get_count(); ++i )
	{
		const mesh_node_data* node = nodes->get_node(i);
		uint32 chan_size  = 0;
		uint32 chan_count = ( node->data ? node->data->get_channel_count() : 0 );
		for ( uint32 c = 0; c < chan_count; ++c )
		{
			chan_size += sizeof( nmd_element_header ) + sizeof( nv::nmd_stream_header );
			chan_size += node->data->get_channel(c)->raw_size();
		}

		nmd_element_header eheader;
		eheader.type     = nmd_type::NODE;
		eheader.name     = strings->insert( node->name );
		eheader.children = static_cast< uint16 >( chan_count );
		eheader.size     = sizeof( nmd_node_header ) + chan_size;
		stream_out.write( &eheader, sizeof( eheader ), 1 );

		nmd_node_header nheader;
		nheader.parent_id = node->parent_id;
		nheader.transform = node->transform;
		stream_out.write( &nheader, sizeof( nheader ), 1 );

		for ( uint32 c = 0; c < chan_count; ++c )
		{
			const raw_data_channel* channel = node->data->get_channel(c);

			eheader.type     = nmd_type::KEY_CHANNEL;
			eheader.children = 0;
			eheader.size     = sizeof( nmd_stream_header ) + channel->raw_size();
			stream_out.write( &eheader, sizeof( eheader ), 1 );

			nmd_stream_header cheader;
			cheader.format    = channel->descriptor();
			cheader.count     = channel->element_count();
			stream_out.write( &cheader, sizeof( cheader ), 1 );
			stream_out.write( channel->raw_data(), channel->element_size(), channel->element_count() );
		}
	}
}

void nv::nmd_dump( const mesh_data_pack* model, stream& stream_out )
{
	string_table_creator strings;
	{
		nmd_header header;
		header.id       = four_cc<'n','m','f','1'>::value;
		header.elements = 1; // +1 string array
		header.elements += model->get_count();
		if ( model->get_nodes() && model->get_nodes()->get_count() > 0 ) 
			header.elements += 1;//  +1 bone array
		header.version  = 1;
		stream_out.write( &header, sizeof( header ), 1 );
	}

	for ( uint32 i = 0; i < model->get_count(); ++i )
	{
		const mesh_data* mesh = model->get_mesh(i);
		nmd_dump_mesh( mesh, stream_out );
	}

	if ( model->get_nodes() && model->get_nodes()->get_count() > 0 )
	{
		nmd_dump_nodes_data( model->get_nodes(), stream_out, &strings );
	}

	nmd_element_header sheader;
	sheader.type     = nv::nmd_type::STRING_TABLE;
	sheader.name     = 0;
	sheader.size     = strings.dump_size();
	sheader.children = 0;
	stream_out.write( &sheader, sizeof( sheader ), 1 );
	strings.dump( &stream_out );
}