#include"node.h"

#include<eebie/writer.h>
#include<eebie/reader.h>
#include<assert.h>
#include<stdio.h>
#include"node_internal.h"
#include<endian.h>

#define ID_PROJECT 0x10000000
#define ID_DURATION 0x4000
#define ID_TIME 0x4001
#define ID_KEYFRAMES 0x200000
#define ID_CHIVAL 0x80
#define ID_CHIVAL_TYPE 0x81
#define ID_LINK 0x82
#define ID_DATA 0x83
#define ID_EXTRAPOLATION_MODE 0x4002
#define ID_EXTRAPOLATION_PARAM 0x4003
#define ID_NODE 0x84
#define ID_NODE_TYPE 0x85
#define ID_LIFESPAN_START 0x4004
#define ID_LIFESPAN_END 0x4005

static void put_chival(CHiNodeGraph *ng, EBMLWriter *ebml, CHiValue val) {
	ebml_writer_push(ebml, ID_CHIVAL);
		ebml_writer_put(ebml, ID_CHIVAL_TYPE, EBML_UNSIGNED_INTEGER, (EBMLPrimitive) {.uInt = val.type});
		
		if(val.linked.to) {
			size_t index;
			for(index = 0; index < ng->count; index++) {
				if(ng->nodes[index] == val.linked.to) {
					break;
				}
			}
			assert(index < ng->count);
			
			ebml_writer_put(ebml, ID_LINK, EBML_UNSIGNED_INTEGER, (EBMLPrimitive) {.uInt = (index << 16) | (val.linked.idx)});
		}
		
		if(val.type == CUTIHI_VAL_TEXT) {
			ebml_writer_put(ebml, ID_DATA, EBML_BINARY, (EBMLPrimitive) {.binary = {.ptr = val.data.text, .length = strlen(val.data.text)}});
		} else if(val.type == CUTIHI_VAL_VEC4) {
			ebml_writer_put(ebml, ID_DATA, EBML_BINARY, (EBMLPrimitive) {.binary = {.ptr = val.data.vec4, .length = sizeof(val.data.vec4)}});
		} else if(val.type == CUTIHI_VAL_KEYED) {
			size_t index;
			for(index = 0; index < ng->keyframesList.count; index++) {
				if(ng->keyframesList.keyframes[index] == val.data.keyed) {
					break;
				}
			}
			assert(index < ng->count);
			
			ebml_writer_put(ebml, ID_DATA, EBML_UNSIGNED_INTEGER, (EBMLPrimitive) {.uInt = index});
		}
	ebml_writer_pop(ebml);
}

static size_t ebml_write_callback(EBMLWriter *ebml, const void *buf, size_t length) {
	CHiSaveWriter writer = ((void**) ebml->ud)[0];
	void *ud = ((void**) ebml->ud)[1];
	return writer(ud, buf, length);
}
static void *ebml_alloc_callback(EBMLWriter *ebml, void *buf, size_t length) {
	return realloc(buf, length);
}

CUTIVIS int CHi_NodeGraphSave(CHiNodeGraph *ng, CHiSaveWriter writer, CHiSaveCustomNodeData customNodeData, void *ud) {
	EBMLWriter ebml;
	ebml_writer_init(&ebml);
	
	void *ebmlud[2] = {writer, ud};
	
	ebml.ud = ebmlud;
	ebml.eventWrite = ebml_write_callback;
	ebml.eventAlloc = ebml_alloc_callback;
	
	ebml_writer_push(&ebml, 0x1A45DFA3);
		ebml_writer_put(&ebml, 0x4286, EBML_UNSIGNED_INTEGER, (EBMLPrimitive) {.uInt = 1});
		ebml_writer_put(&ebml, 0x4282, EBML_STRING, (EBMLPrimitive) {.string = "cuticle"});
	ebml_writer_pop(&ebml);

	ebml_writer_push(&ebml, ID_PROJECT);
		if(ng->duration != -1) {
			ebml_writer_put(&ebml, ID_DURATION, EBML_FLOAT4, (EBMLPrimitive) {.flt4 = ng->duration});
		}
		ebml_writer_put(&ebml, ID_TIME, EBML_FLOAT4, (EBMLPrimitive) {.flt4 = ng->time});
		
		for(size_t i = 0; i < ng->keyframesList.count; i++) {
			CHiKeyframes *kfs = ng->keyframesList.keyframes[i];
			
			ebml_writer_push(&ebml, ID_KEYFRAMES);
				//ebml_writer_put(&ebml, ID_CHIVAL_TYPE, EBML_UNSIGNED_INTEGER, (EBMLPrimitive) {.uInt = kfs->type});
				
				for(size_t k = 0; k < kfs->count; k++) {
					ebml_writer_put(&ebml, ID_TIME, EBML_FLOAT4, (EBMLPrimitive) {.flt4 = kfs->times[k]});
					put_chival(ng, &ebml, (CHiValue) {.type = kfs->type, .linked = {}, .data = kfs->values[k]});
				}
				
				ebml_writer_put(&ebml, ID_EXTRAPOLATION_MODE, EBML_UNSIGNED_INTEGER, (EBMLPrimitive) {.uInt = kfs->extrapolationMode});
				ebml_writer_put(&ebml, ID_EXTRAPOLATION_PARAM, EBML_BINARY, (EBMLPrimitive) {.binary = {.ptr = &kfs->extrapolationParameter, sizeof(kfs->extrapolationParameter)}});
			ebml_writer_pop(&ebml);
		}
		
		for(size_t i = 0; i < ng->count; i++) {
			CHiPubNode *node = ng->nodes[i];
			
			ebml_writer_push(&ebml, ID_NODE);
				ebml_writer_put(&ebml, ID_NODE_TYPE, EBML_BINARY, (EBMLPrimitive) {.binary = {.ptr = &node->type, sizeof(node->type)}});
				
				if(node->lifespan.start != 0) {
					ebml_writer_put(&ebml, ID_LIFESPAN_START, EBML_FLOAT4, (EBMLPrimitive) {.flt4 = node->lifespan.start});
				}
				if(node->lifespan.end != 0) {
					ebml_writer_put(&ebml, ID_LIFESPAN_END, EBML_FLOAT4, (EBMLPrimitive) {.flt4 = node->lifespan.end});
				}
				
				for(size_t sink = 0; sink < node->sinkCount; sink++) {
					put_chival(ng, &ebml, node->sinks[sink]);
				}
				
				if(customNodeData) {
					customNodeData(ud, &ebml, node);
				}
			ebml_writer_pop(&ebml);
		}
	ebml_writer_pop(&ebml);
	
	while(!ebml_writer_flush(&ebml));
	
	return 0;
}

struct LoadState {
	uint8_t *accum;
	size_t accumSize;
	
	CHiNodeGraph *ng;
	
	void *ud;
	CHiLoadCustomNodeData customNodeData;
	
	// Filled in by callbacks
	CHiValue activeVal;
	
	// Filled in by callbacks
	size_t activeNodeNextSinkIdx;
	CHiPubNode *activeNode;
};

static EBMLElementType ebml_enter_callback(EBMLReader *ebml, uint64_t id, uint64_t length) {
	struct LoadState *state = ebml->ud;
	
	if(id == ID_KEYFRAMES) {
		state->ng->keyframesList.keyframes = realloc(state->ng->keyframesList.keyframes, sizeof(*state->ng->keyframesList.keyframes) * (++state->ng->keyframesList.count));
		state->ng->keyframesList.keyframes[state->ng->keyframesList.count - 1] = calloc(1, sizeof(CHiKeyframes));
	} else if(id == ID_CHIVAL) {
		memset(&state->activeVal, 0, sizeof(state->activeVal));
	}
	
	return (id == ID_PROJECT || id == ID_KEYFRAMES || id == ID_CHIVAL || id == ID_NODE || id == 0x5000) ? EBML_TREE : EBML_BINARY;
}

static void ebml_data_callback(EBMLReader *ebml, const uint8_t *data, size_t length) {
	struct LoadState *state = ebml->ud;
	
	state->accum = realloc(state->accum, state->accumSize + length);
	memcpy(state->accum + state->accumSize, data, length);
	state->accumSize += length;
}

static void ebml_exit_callback(EBMLReader *ebml) {
	struct LoadState *state = ebml->ud;
	
	uint64_t parent = ebml->currentDepth >= 2 ? ebml->idStack[ebml->currentDepth - 2] : 0;
	uint64_t id = ebml->idStack[ebml->currentDepth - 1];
	
	if(id == ID_NODE) {
		state->activeNode = NULL;
	} else if(parent == 0x5000 && state->customNodeData) {
		assert(state->activeNode == state->ng->nodes[state->ng->count - 1]);
		
		state->customNodeData(state->ud, state->ng->count - 1, id, state->accum, state->accumSize);
	} else if(parent == ID_NODE) {
		if(id == ID_NODE_TYPE) {
			assert(state->accumSize <= 8);
			
			uint64_t type = 0;
			memcpy((char*) &type + sizeof(type) - state->accumSize, state->accum, state->accumSize);
			
			CHiPubNode *n = NULL;
			
			if(type == CUTIHI_T('CPre','view')) {
				n = CHi_Preview();
			} else if(type == CUTIHI_T('CMix','er  ')) {
				n = CHi_Mixer();
			} else if(type == CUTIHI_T('CTex','t   ')) {
				n = CHi_Text();
			} else if(type == CUTIHI_T('CTim','e   ')) {
				n = CHi_Time();
			} else if(type == CUTIHI_T('CMod','ulat')) {
				n = CHi_Modulate();
			} else if(type == CUTIHI_T('CCns','tCol')) {
				n = CHi_ConstantSample();
			} else if(type == CUTIHI_T('CEmb','ed  ')) {
				n = CHi_Embed();
			} else if(type == CUTIHI_T('CIma','ge  ')) {
				n = CHi_Image();
			} else if(type == CUTIHI_T('CWin','dow ')) {
				n = CHi_Window();
			} else if(type == CUTIHI_T('CInA','udio')) {
				n = CHi_Microphone();
			} else if(type == CUTIHI_T('CExp','Wave')) {
				n = CHi_ExportWav();
			} else if(type == CUTIHI_T('CMov','ie  ')) {
				n = CHi_Movie();
			} else if(type == CUTIHI_T('CEnc','GVP8')) {
				n = CHi_EncodeVP8();
			} else if(type == CUTIHI_T('CEnc','GVP9')) {
				n = CHi_EncodeVP9();
			} else if(type == CUTIHI_T('CExp','Webm')) {
				n = CHi_MuxWebm();
			} else if(type == CUTIHI_T('CKey','hook')) {
				n = CHi_Keyhook();
			} else if(type == CUTIHI_T('CKey','hook')) {
				n = CHi_Keyhook();
			} else if(type == CUTIHI_T('CEnc','Opus')) {
				n = CHi_EncodeOpus();
			} else if(type == CUTIHI_T('CWeb','Cam ')) {
				n = CHi_Camera();
			} else if(type == CUTIHI_T('CCmp','nScl')) {
				n = CHi_ComponentScale();
			} else if(type == CUTIHI_T('CEnc','H264')) {
				n = CHi_EncodeH264();
			} else if(type == CUTIHI_T('CStr','RTMP')) {
				n = CHi_StreamRTMP();
			} else if(type == CUTIHI_T('CEnc','AACL')) {
				n = CHi_EncodeAAC();
			} else if(type == CUTIHI_T('CExp','Mkv ')) {
				n = CHi_MuxMatroska();
			}
			
			n->ng = state->ng;
			
			if(!n) {
				CHi_NodeGraphReset(state->ng);
				assert(0 && "Unknown node type.");
			}
			
			state->ng->nodes = realloc(state->ng->nodes, sizeof(*state->ng->nodes) * (++state->ng->count));
			state->ng->nodes[state->ng->count - 1] = n;
			
			state->activeNodeNextSinkIdx = 0;
			state->activeNode = n;
		} else if(id == ID_LIFESPAN_START) {
			state->activeNode->lifespan.start = *(float*) state->accum;
		} else if(id == ID_LIFESPAN_END) {
			state->activeNode->lifespan.end = *(float*) state->accum;
		} else if(id == ID_CHIVAL) {
			if(state->activeNodeNextSinkIdx >= state->activeNode->sinkCount) {
				state->activeNode->sinks = realloc(state->activeNode->sinks, sizeof(*state->activeNode->sinks) * (state->activeNode->sinkCount = (state->activeNodeNextSinkIdx + 1)));
			}
			state->activeNode->sinks[state->activeNodeNextSinkIdx++] = state->activeVal;
		}
	} else if(parent == ID_CHIVAL) {
		if(id == ID_CHIVAL_TYPE) {
			state->activeVal.type = *(uint8_t*) state->accum;
		} else if(id == ID_LINK) {
			assert(state->accumSize <= 8);

			uint64_t data = 0;
			memcpy((char*) &data + sizeof(data) - state->accumSize, state->accum, state->accumSize);
			data = htobe64(data);
			
			uint16_t srcIdx = data & 0xFFFF;
			uint64_t nodeIdx = data >> 16;

			// linked.to is supposed to be a pointer but not all nodes are
			// available at this stage, so we'll fix it up after the EBML
			// parsing finishes
			state->activeVal.linked.to = (void*) (nodeIdx + 1);
			state->activeVal.linked.idx = srcIdx;
		} else if(id == ID_DATA) {
			if(state->activeVal.type == CUTIHI_VAL_TEXT) {
				state->activeVal.data.text = calloc(1, state->accumSize + 1);
				memcpy(state->activeVal.data.text, state->accum, state->accumSize);
			} else if(state->activeVal.type == CUTIHI_VAL_VEC4) {
				assert(state->accumSize <= 16 && state->accumSize % 4 == 0);
				
				memcpy(state->activeVal.data.vec4, state->accum, state->accumSize);
			} else if(state->activeVal.type == CUTIHI_VAL_KEYED) {
				assert(state->accumSize <= 8);
				
				uint64_t idx = 0;
				memcpy((char*) &idx + sizeof(idx) - state->accumSize, state->accum, state->accumSize);
				idx = htobe64(idx);
				
				// keyed is supposed to be a pointer but keyframes may not
				// be available at this stage, so we'll fix it up after the
				// EBML parsing finishes
				state->activeVal.data.keyed = (void*) idx;
			} else {
				assert(0 && "Cannot load CHiValue data.");
			}
		}
	} else if(parent == ID_KEYFRAMES) {
		CHiKeyframes *kfs = state->ng->keyframesList.keyframes[state->ng->keyframesList.count - 1];
		if(id == ID_TIME) {
			kfs->times = realloc(kfs->times, sizeof(*kfs->times) * (++kfs->count));
			kfs->values = realloc(kfs->values, sizeof(*kfs->values) * kfs->count);
			
			kfs->times[kfs->count - 1] = *(float*) state->accum;
		} else if(id == ID_CHIVAL) {
			assert(state->activeVal.linked.to == NULL);
			
			kfs->values[kfs->count - 1] = state->activeVal.data;
		} else if(id == ID_EXTRAPOLATION_MODE) {
			kfs->extrapolationMode = *(uint8_t*) state->accum;
		} else if(id == ID_EXTRAPOLATION_PARAM) {
			assert(state->accumSize == 16);
			
			memcpy(&kfs->extrapolationParameter, state->accum, 16);
		}
	}
	
	state->accumSize = 0;
}

CUTIVIS int CHi_NodeGraphLoad(CHiNodeGraph *ng, CHiLoadReader reader, CHiLoadCustomNodeData customNodeData, void *ud) {
	CHi_NodeGraphReset(ng);
	
	ng->time = 0;
	ng->duration = -1;
	
	EBMLReader ebml;
	ebml_reader_init(&ebml);
	
	struct LoadState state = {.ng = ng, .ud = ud, .customNodeData = customNodeData};
	ebml.ud = &state;
	
	ebml.eventEnterElement = ebml_enter_callback;
	ebml.eventDataChunk = ebml_data_callback;
	ebml.eventExitElement = ebml_exit_callback;
	
	size_t i = 0;
	uint8_t buf[4096];
	while(1) {
		size_t readBytes = reader(ud, buf + i, sizeof(buf) - i);
		
		i += readBytes;
		
		size_t fedBytes = ebml_reader_feed(&ebml, buf, i);
		
		memmove(buf, buf + fedBytes, sizeof(buf) - fedBytes);
		
		i -= fedBytes;
		
		if(readBytes == 0 && fedBytes == 0) {
			break;
		}
	}
	
	// Fix-up links/adjacencies
	for(size_t i = 0; i < ng->count; i++) {
		CHiPubNode *n = ng->nodes[i];
		for(size_t s = 0; s < n->sinkCount; s++) {
			if(n->sinks[s].linked.to) {
				n->sinks[s].linked.to = ng->nodes[(uintptr_t) n->sinks[s].linked.to - 1];
			
				adjacency_add(n->sinks[s].linked.to, n);
			}
			
			if(n->sinks[s].type == CUTIHI_VAL_KEYED) {
				n->sinks[s].data.keyed = ng->keyframesList.keyframes[(uintptr_t) n->sinks[s].data.keyed];
			}
		}
	}
	
	update_keyed_values(ng);
	
	return 0;
}