#include "AviMediaReadHandler.h"
#include "AviMediaReadStream.h"
#include <fourcc.h>
#include <stdio.h>

#define __MODULE__ "MediaReadHandler"
#define Debug if(0)

// reset profiler status - just for debuging
//extern "C" void kprof_reset_stats(void);

#include <iostream>
using namespace std;

AviMediaReadHandler::AviMediaReadHandler(const char* pszFile) throw(FatalError)
    :m_Input(pszFile)
{
    m_Input.seek(0);
    checkDword(FOURCC_RIFF);
    /*int riff_chunk=*/
    m_Input.readDword();
    checkDword(formtypeAVI);

    int movie_chunk = -1;
    bool index_chunk_found = false;
    unsigned int chunk_type;
    size_t chunk_size;
    unsigned int type, nextpos;
    while (!m_Input.eof())
    {
	chunk_type = m_Input.readDword();
	chunk_size = m_Input.readDword();
//	if(chunk_size<0)throw FATAL("Damaged AVI file");
	nextpos = m_Input.pos() + chunk_size + (chunk_size & 1);
	switch(chunk_type)
	{
	case FOURCC_LIST:
	    if (chunk_size < 4)
		throw FATAL("Damaged AVI file");
	    type = m_Input.readDword();
	    switch(type)
	    {
	    case listtypeAVIMOVIE:
		movie_chunk=m_Input.pos();
	        m_Input.seek(nextpos);
//		readMovieChunk();
		break;
	    case listtypeAVIHEADER:
		continue;
	    case listtypeSTREAMHEADER:
		readStreamHeader();
		break;
	    default:
		break;
	    };
	    break;
	case ckidAVIMAINHDR:
	    readMainHeader(chunk_size);
	    break;
	case ckidAVINEWINDEX:
	    index_chunk_found=true;
	    readIndexChunk(chunk_size, (movie_chunk < 0) ? 0 : movie_chunk);
	    break;
    	default:
	    break;
	};
	m_Input.seek(nextpos);
    }
/*
    checkInt(FOURCC_LIST);
    int header_chunk=m_Input.readInt();
    checkInt(listtypeAVIHEADER);
    header_chunk-=4;
//    riff_chunk-=(header_chunk+8);

    header_chunk-=readMainHeader();
    if(header_chunk<0)throw FATAL("Damaged AVI file");

    int nextpos=m_Input.pos()+header_chunk;
    readStreams(header_chunk);
    m_Input.seek(nextpos);

    int movie_chunk=readMovieChunk();
    if(m_Input.eof() || (readIndexChunk(movie_chunk)<0))
	reconstructIndexChunk(movie_chunk);
*/
    if (movie_chunk < 0)
	throw FATAL("Movie chunk not found");
    if (!index_chunk_found)
	reconstructIndexChunk(movie_chunk);

    //vector<AviMediaReadStream>::iterator it;
    //    for(it=m_Streams.begin(); it!=m_Streams.end(); it++)
    for (unsigned i = 0; i < m_Streams.size(); i++)
    {
	printf("Stream %d: 0x%x (%.4s) : 0x%x (%.4s)   %d chunks\n",
	       i, m_Streams[i].m_Header.fccType,
	       (const char*)&m_Streams[i].m_Header.fccType,
	       m_Streams[i].m_Header.fccHandler,
	       (const char*)&m_Streams[i].m_Header.fccHandler,
	       m_Streams[i].m_Index.size());
	m_Input.addStream(m_Streams[i].m_iId, m_Streams[i].m_Index);
    }
    m_Input.async();
}

AviMediaReadHandler::~AviMediaReadHandler()
{
}

//HRESULT AviMediaReadHandler::GetHeader(MainAVIHeader* pheader) const
HRESULT AviMediaReadHandler::GetHeader(void* pheader, size_t size) const
{
    if (!pheader || (size < sizeof(MainAVIHeader)))
        return -1;
    memset(pheader, 0, size);
    memcpy(pheader, &m_Header, sizeof(MainAVIHeader));
    return 0;
}

void AviMediaReadHandler::checkDword(unsigned int value)
{
    if (m_Input.readDword() != value)
	throw FATAL("Not an AVI file!");
}

int AviMediaReadHandler::readMainHeader(size_t data_size)
{
//    checkInt(ckidAVIMAINHDR);//avih
//    unsigned int data_size=m_Input.readInt();
    size_t read_size = sizeof(MainAVIHeader);
    if (data_size > read_size)
	cerr << "WARNING: unexpected main AVI header size" << endl;
    else
	read_size = data_size;
    m_Input.read(&m_Header, read_size);
    if (data_size > read_size)
	m_Input.seekCur(data_size - read_size);
    if (data_size & 1)
	m_Input.seekCur(1);
    return data_size + 8;
}

int AviMediaReadHandler::readStreamHeader()
{
    checkDword(ckidSTREAMHEADER);//strh
    size_t data_size = m_Input.readDword();
    size_t read_size = sizeof(AVIStreamHeader);

    if (data_size>read_size)
	cerr<<"WARNING: unexpected AVI stream header size"<<endl;
    else
	read_size=data_size;

    m_Streams.push_back(AviMediaReadStream(this));

    AviMediaReadStream& s = m_Streams.back();

    m_Input.read(&s.m_Header, read_size);
    if (data_size > read_size)
	m_Input.seekCur(data_size-read_size);
    if (data_size % 2)
	m_Input.seekCur(1);

    checkDword(ckidSTREAMFORMAT);//strf
    data_size = m_Input.readDword();
    s.m_pcFormat = new char[data_size];
    m_Input.read(s.m_pcFormat, data_size);
    s.m_uiFormatSize = data_size;
    s.m_iId = m_Streams.size()-1;
    if (data_size % 2)
	m_Input.seekCur(1);

    return 0;
}
/*
void AviMediaReadHandler::readStreams(int streams_size)
{
    int read_bytes=0;
    while(read_bytes<streams_size)
    {
	if(m_Input.readInt()!=FOURCC_LIST)
	    return;
	int stream_size=m_Input.readInt();
	if(m_Input.readInt()!=listtypeSTREAMHEADER)//strl
	    return;
	read_bytes+=12;
	stream_size-=4;
	int nextpos=m_Input.pos()+stream_size;
	readStreamHeader();
	m_Input.seek(nextpos);
	read_bytes+=stream_size;
    }
}
int AviMediaReadHandler::readMovieChunk()
{
    int moviem_uiChunk;
    while(!m_Input.eof())
    {
	if(m_Input.readInt()!=FOURCC_LIST)
	{
	    int chunk_size=m_Input.readInt();
	    if(chunk_size<0)throw FATAL("Damaged AVI file");
	    if(chunk_size%2) chunk_size++;
	    m_Input.seekCur(chunk_size);
	    continue;
	}
	int chunk_size=m_Input.readInt();
	if(chunk_size<4)throw FATAL("Damaged AVI file");
	if(chunk_size%2)chunk_size++;
	checkInt(listtypeAVIMOVIE);
	moviem_uiChunk=m_Input.pos();
	m_Input.seekCur(chunk_size-4);
	return moviem_uiChunk;
    }
    throw FATAL("Movie chunk not found");
}
*/
int AviMediaReadHandler::readIndexChunk(int index_size, int movie_chunk_offset)
{
//    if(m_Input.readInt()!=ckidAVINEWINDEX)
//	return -1;

    bool zero_based_offsets = false;
//    movie_chunk_offset-=4;

//    int index_size=m_Input.readInt();
//    if(index_size<0)return -1;
    int64_t positions[m_Streams.size()];

    for (unsigned i = 0; i < m_Streams.size(); i++)
	positions[i] = 0;

    AVIINDEXENTRY entry;

    for(; (index_size >= (int) sizeof(AVIINDEXENTRY)) && !m_Input.eof();
	index_size -= sizeof(AVIINDEXENTRY))
    {
	AVIINDEXENTRY2 entry2;
	int len = m_Input.read(&entry, sizeof(entry));
	if (len < (int) sizeof(entry))
	{
	    cerr<<"WARNING: incomplete chunk entry, len="<<len<<endl;
	    break;
	}

	unsigned int id = StreamFromFOURCC(entry.ckid);
	if (id >= m_Streams.size())
	    continue;

	AviMediaReadStream& stream = m_Streams[id];

	entry2.dwFlags = entry.dwFlags;
	if (entry.dwChunkOffset < movie_chunk_offset)
	    zero_based_offsets = true;

	if (zero_based_offsets)
	    entry2.qwChunkOffset = entry.dwChunkOffset + movie_chunk_offset - 4;
	else
	    entry2.qwChunkOffset = entry.dwChunkOffset;

	entry2.dwChunkLength = entry.dwChunkLength;
	entry2.qwTimestamp = positions[id];
	if (stream.m_Header.dwSampleSize)
	    // audio stream, chunks contain variable number of frames
	    positions[id] += entry2.dwChunkLength / stream.m_Header.dwSampleSize;
	else
	    // video stream, one frame / chunk
	    positions[id]++;
	stream.m_Index.push_back(entry2);
    }

    //kprof_reset_stats();

    return 0;
}

void AviMediaReadHandler::reconstructIndexChunk(int offset)
{
    m_Input.seek(offset);
    int64_t positions[m_Streams.size()];
    bool isdivx[m_Streams.size()];

    cout << "Reconstructing index for this file" << endl;
    // fixing KeyFrames for DivX movies
    // they contain flag bit so we will check it
    for (unsigned i = 0; i < m_Streams.size(); i++)
    {
	positions[i] = 0;
	isdivx[i] = false;

	if (m_Streams[i].m_Header.fccType == streamtypeVIDEO)
	    switch (m_Streams[i].m_Header.fccHandler)
	    {
	    case fccDIV3:
	    case fccdiv3:
	    case fccDIV4:
	    case fccdiv4:
	    case fccDIV5:
	    case fccdiv5:
	    case fccDIV6:
	    case fccdiv6:
	    case fccMP43:
	    case fccmp43:
	    case fccMP42:
	    case fccmp42:
	    case fccAP41:
		isdivx[i] = true;
	    }
    }

    // set when trying to reconstruct even very broken .avi files
    bool seriouslyBroken = false;
    while (!m_Input.eof())
    {
//	cout << hex << m_Input.pos() << dec << endl;
	unsigned int ckid = m_Input.readDword();
	streamid_t id = StreamFromFOURCC(ckid);
	if (id >= m_Streams.size())
	{
	    if (m_Input.eof())
		break;
	    m_Input.seekCur(-2);
            if (!seriouslyBroken)
		cout << "Trying to reconstruct very broken avi stream (could take a lot of time)" << endl;
	    seriouslyBroken = true;
	    continue;
	}
	AviMediaReadStream& stream = m_Streams[id];
	AVIINDEXENTRY2 entry2;
	entry2.dwFlags = AVIIF_KEYFRAME;
	entry2.qwChunkOffset = m_Input.pos() - 4;
	size_t size = m_Input.readDword();
        // max size for frame will be 250MB
	if (size > ((seriouslyBroken) ? 100000U : 0xfffffffU))
	// || size == 0)
	// size = 0 is legal
	    continue;
	entry2.dwChunkLength = size;
	entry2.qwTimestamp = positions[id];
	if (stream.m_Header.dwSampleSize) //audio stream, chunks contain variable number of frames
	    positions[id] += entry2.dwChunkLength / stream.m_Header.dwSampleSize;
	else //video stream, one frame / chunk
	    positions[id]++;

	if (size & 1)
	    size++;

	if (isdivx[id] && size)
	{
	    int b = m_Input.readByte();
	    if (b & 0x40)
		entry2.dwFlags = 0;
            size--;
	}

	//cout << "Size " << entry2.dwChunkLength << " offset " << entry2.qwChunkOffset << endl;
	stream.m_Index.push_back(entry2);
	m_Input.seekCur(size);
    }
}

IMediaReadStream* AviMediaReadHandler::GetStream(fourcc_t fccType, streamid_t lParam)
{
    std::vector<AviMediaReadStream>::iterator it;
    streamid_t match = 0;
    for (it = m_Streams.begin(); it != m_Streams.end(); it++)
	if (it->m_Header.fccType == fccType)
	{
	    if (match == lParam)
		return &(*it); //safe for gcc3.0
	    match++;
	}
    return 0;
}

IMediaReadHandler *CreateAviMediaReadHandler(const char *pszFile) throw(FatalError)
{
    return new AviMediaReadHandler(pszFile);
}
