/********************************************************

	Classes for writing AVIs
	Copyright 2000 Eugene Kuznetsov  (divx@euro.ru)

*********************************************************/
#ifndef _AVIWRITE_H
#define _AVIWRITE_H

#include "config.h"

#include <avifmt.h>
#include <default.h>
#include <unistd.h>
#ifdef HAVE_MALLOC_H
#include <malloc.h>
#else
#include <stdlib.h>
#endif
#include <string.h>
#include <stdio.h>
#include <pthread.h>
#include <avifile.h>    
#include <videoencoder.h>
#include "FileBuffer.h"
#include <vector>
using namespace std;

#include <except.h>
#define __MODULE__ "AviWriteFile"

class AviWriteFile;

class AviWriteStream : public IAviWriteStream
{
friend class AviWriteFile;

protected:
    AviWriteFile* m_file;
    FileBuffer* m_fd;
//    int m_fd;
    StreamType m_type;
    
    char* m_format;
    Unsigned m_forsize;
    int m_ckid;
    AviWriteStream(AviWriteFile* file, int ckid,
	enum AviStream::StreamType, int handler, int frame_rate, int flags=0);
    void SetVideoData(int quality, int width, int height)
    {
    	m_header.dwQuality=quality;
        m_header.rcFrame.right=width;
	m_header.rcFrame.bottom=height;
    }
//    void SetAudioData(
public:
    AviWriteStream(AviWriteFile* file, int ckid,
	enum AviStream::StreamType type,
	const char* format, Unsigned format_size,
	int handler, int frame_rate,
	int samplesize=0, int quality=0, int flags=0);
    
    ~AviWriteStream();
    //
    // these two should be called before data insertion begins
    //
    //
    // for video streams, specify time in microsecs per frame in frame_rate
    // for audio streams - count of bytes per second
    //   
    StreamType GetType() const
    {
    	return m_type;
    }	    
    HRESULT AddChunk(const char* chunk, Unsigned size, Unsigned flags=0);

    Unsigned GetLength()//0 on error
    {
	return m_header.dwLength;
    }	    
};            

class AviAudioWriteStream : public IAviAudioWriteStream, public AviWriteStream
{
    IAudioEncoder* m_pIEnc;
    WAVEFORMATEX srcfmt;
    int _status;
    int _bitrate;
public:
    AviAudioWriteStream(AviWriteFile* file, int ckid, int fourcc, WAVEFORMATEX* fmt,
	int bitrate, int flags=0)
	:AviWriteStream(file, ckid, AviStream::Audio, fourcc, bitrate, flags),
	 _status(0), _bitrate(bitrate), m_pIEnc(0)
//    AviWriteStream(AviWriteFile* file, int ckid,
//	enum AviStream::StreamType type,
//	const char* format, Unsigned format_size,
//	int handler, int frame_rate,
//	int samplesize=0, int quality=0, int flags=0);
    {
	if(fmt)srcfmt=*fmt;
    }
    virtual ~AviAudioWriteStream()
    {
	delete m_pIEnc;
    }
    virtual HRESULT Start()
    {
	if(_status)return 0;

	m_pIEnc=IAudioEncoder::Create(AviWriteStream::m_header.fccHandler, &srcfmt);
	m_pIEnc->SetBitrate(_bitrate);
	WAVEFORMATEX* fmt;
	m_forsize=m_pIEnc->GetFormatSize();
	m_format=new char[m_forsize];
	m_pIEnc->GetFormat(m_format, m_forsize);
	fmt=(WAVEFORMATEX*)m_format;
	AviWriteStream::m_header.dwSampleSize=fmt->nBlockAlign;
	AviWriteStream::m_header.dwRate=_bitrate;
	AviWriteStream::m_header.dwScale=fmt->nBlockAlign;
	
	m_pIEnc->Start();
	_status=1;
    }
    virtual HRESULT AddData(void* data, int size)
    {
	if(!_status)return -1;
	int outsize=2*size/srcfmt.nBlockAlign+7200;
	char* buf=new char[outsize];
	int written=0;
	HRESULT hr;
	if(data)
	{
	    hr=m_pIEnc->Convert((char*)data, size/srcfmt.nBlockAlign, buf, outsize, (int*)0, &written);
//	    cerr<<"Converted "<<size<<" bytes to "<<written<<" bytes"<<endl;
	    if(hr==0)
		hr=AviWriteStream::AddChunk(buf, written, 16);
	}
	else hr=AviWriteStream::AddChunk(NULL, 0);
	delete buf;
	return hr; 
    }
    virtual HRESULT Stop()
    {
	if(!_status)return -1;
	char buf[7200];
	int written=0;
	m_pIEnc->Close(buf, 7200, &written);
	if(written>0)
	    AviWriteStream::AddChunk(buf, written);
	delete m_pIEnc;
	m_pIEnc=0;
	_status=0;
	return 0;	
    }
    //0..10
    virtual HRESULT SetQuality(int quality)
    {
	if(m_pIEnc)return m_pIEnc->SetQuality(quality);
	return -1;
    }	
    Unsigned GetLength()
    {
	return AviWriteStream::GetLength();
    }	
};

class AviVideoWriteStream : public IAviVideoWriteStream, public AviWriteStream
{
    IVideoEncoder* m_pIEnc;
    int m_vstatus;
    char* m_local_buffer;
public:
    AviVideoWriteStream(AviWriteFile* file, int ckid, int fourcc, BITMAPINFOHEADER* srchdr,
	int frame_rate, int flags=0)
	:AviWriteStream(file, ckid, Video, fourcc, frame_rate, flags ), m_vstatus(0),
	m_local_buffer(0)
    {
	m_pIEnc=IVideoEncoder::Create(fourcc, *srchdr);
    }
    ~AviVideoWriteStream()
    {
	if(m_pIEnc)delete m_pIEnc;
	if(m_local_buffer)delete m_local_buffer;
    }
//    void Create(int fourcc, BITMAPINFOHEADER* srchdr)
//    {
//	m_encoder.Init(fourcc, (const char*)srchdr);
//	m_vstatus=1;
//	m_fourcc=fourcc;
//	SetType(Video);
//        return;
//    }   
//    virtual void SetHeader(int frame_rate,
//            int flags=0)
//    {
//	if(m_vstatus!=1)throw FATAL("Bad status");
//	AviWriteStream::SetHeader(m_fourcc, frame_rate, 0, m_encoder.GetQuality(), flags);
//    }
    virtual HRESULT Start()
    {
	if(m_vstatus)return -1;
	BITMAPINFOHEADER bh=m_pIEnc->QueryOutputFormat();
	m_forsize=sizeof(BITMAPINFOHEADER);
	m_format=new char[m_forsize];
	memcpy(m_format, &bh, m_forsize);
	m_pIEnc->Start();
	m_local_buffer=new char[m_pIEnc->QueryOutputSize()];
	SetVideoData(m_pIEnc->GetQuality(), bh.biWidth, bh.biHeight);
	m_vstatus=1;
	return 0;
    }	
    virtual HRESULT Stop()
    {
	m_pIEnc->Stop();
	return 0;
    }	
    HRESULT AddFrame(CImage* chunk)
    {
	if(m_vstatus!=1)return -1;
	int size;
	int is_keyframe;
	HRESULT hr;
	if(chunk)
	{	    
	    hr=m_pIEnc->EncodeFrame(chunk, m_local_buffer, &is_keyframe, &size);
	    if(hr==0)
		hr=AviWriteStream::AddChunk(m_local_buffer, size, is_keyframe);
	}
	else hr=AviWriteStream::AddChunk(NULL, 0);
	return hr; 
    }
    virtual HRESULT SetQuality(int quality)
    {
	return m_pIEnc->SetQuality(quality);
    }	
    HRESULT SetKeyFrame(int frequency)
    {
	return m_pIEnc->SetKeyFrame(frequency);
    }	
    Unsigned GetLength()
    {
	return AviWriteStream::GetLength();
    }	
};

//
//maintains global file info
//
class AviSegWriteFile;
class AviWriteFile: public IAviWriteFile
{
friend class AviWriteStream;
friend class AviSegWriteFile;
private:
    vector<AviWriteStream*> m_streams;
    int m_strcnt;
    MainAVIHeader m_header;
    int m_status;
    vector<AVIINDEXENTRY> m_index;
protected:
    FileBuffer* m_fd;
//    int m_fd;
    void _destruct();
public:
    AviWriteFile(const char* name, int flags=0, int mask=00777);
    ~AviWriteFile();
    void WriteHeaders();
//    void Create(const char* name, int flags=0, int mask=00777);
    IAviAudioWriteStream* AddAudioStream(int fourcc, 
	WAVEFORMATEX* fmt, int bitrate, int flags=0);
    IAviVideoWriteStream* AddVideoStream(int fourcc, 
	BITMAPINFOHEADER* srchdr, int frame_rate, int flags);
				    
    IAviWriteStream* AddStream(enum AviStream::StreamType type,
	const char* format, Unsigned format_size,
	int handler, int frame_rate,
	int samplesize=0, int quality=0, int flags=0);
    HRESULT Reserve(Unsigned size);
    HRESULT WriteChunk(int fourcc, char* data, int size);
    long long FileSize();

protected:
    //
    //called when data is added to stream
    //
    void AddChunk(offset_t offset, Unsigned size, Unsigned id, Unsigned flags=0);    
};

#endif