/*
 * AVI encoder.
 * Copyright (c) 2000 Gerard Lantau.
 *
 * 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.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */
#include <stdlib.h>
#include <stdio.h>
#include <netinet/in.h>
#include <string.h>
#include "mpegenc.h"

typedef long long offset_t;

typedef struct AVIIndex {
    unsigned char tag[4];
    unsigned int flags, pos, len;
    struct AVIIndex *next;
} AVIIndex;

typedef struct {
    offset_t movi_list;
    AVIIndex *first, *last;
} AVIContext;

static offset_t start_tag(PutByteContext *pb, char *tag)
{
    put_tag(pb, tag);
    put_le32(pb, 0);
    return put_pos(pb);
}

static void end_tag(PutByteContext *pb, offset_t start)
{
    offset_t pos;

    pos = put_pos(pb);
    put_seek(pb, start - 4, SEEK_SET);
    put_le32(pb, pos - start);
    put_seek(pb, pos, SEEK_SET);
}


#define AVIF_HASINDEX		0x00000010	// Index at end of file?
#define AVIF_MUSTUSEINDEX	0x00000020
#define AVIF_ISINTERLEAVED	0x00000100
#define AVIF_TRUSTCKTYPE	0x00000800	// Use CKType to find key frames?
#define AVIF_WASCAPTUREFILE	0x00010000
#define AVIF_COPYRIGHTED	0x00020000

#define MKTAG(a,b,c,d) (a | (b << 8) | (c << 16) | (d << 24))

unsigned int codec_get_bmp_tag(int id)
{
    switch(id) {
    case CODEC_ID_H263:
        return MKTAG('I', '2', '6', '3'); 
    case CODEC_ID_MJPEG:
        return MKTAG('M', 'J', 'P', 'G'); 
    case CODEC_ID_DIVX:
        return MKTAG('d', 'i', 'v', 'x'); 
    default:
        abort();
    }
}

/* BITMAPINFOHEADER header */
void put_bmp_header(PutByteContext *pb, AVEncodeContext *enc)
{
    put_le32(pb, 40); /* size */
    put_le32(pb, enc->width);
    put_le32(pb, enc->height);
    put_le16(pb, 1); /* planes */
    put_le16(pb, 24); /* depth */
    /* compression type */
    put_le32(pb, codec_get_bmp_tag(enc->codec->id));
    put_le32(pb, enc->width * enc->height * 3);
    put_le32(pb, 0);
    put_le32(pb, 0);
    put_le32(pb, 0);
    put_le32(pb, 0);
}

/* WAVEFORMATEX header */
void put_wav_header(PutByteContext *pb, AVEncodeContext *enc)
{
    int id;

    switch(enc->codec->id) {
    case CODEC_ID_MP2:
        id = 0x55; /* MP3 format */
        break;
    default:
        abort();
    }
    put_le16(pb, id); 
    put_le16(pb, enc->channels);
    put_le32(pb, enc->rate);
    put_le32(pb, enc->bit_rate / 8);
    put_le16(pb, 1); /* block align */
    put_le16(pb, 16); /* bits per sample */
    put_le16(pb, 0); /* wav_extra_size */
}

static int avi_write_header(AVFormatContext *s)
{
    AVIContext *avi;
    PutByteContext *pb = &s->pb;
    int bitrate, n, i;
    AVEncodeContext *stream;
    offset_t list1, list2, strh, strf;

    avi = malloc(sizeof(AVIContext));
    if (!avi)
        return -1;
    memset(avi, 0, sizeof(AVIContext));
    s->priv_data = avi;

    put_tag(pb, "RIFF");
    put_le32(pb, 0); /* file length */
    put_tag(pb, "AVI ");

    /* header list */
    list1 = start_tag(pb, "LIST");
    put_tag(pb, "hdrl");

    /* avi header */
    put_tag(pb, "avih");
    put_le32(pb, 14 * 4);
    bitrate = 0;
    n = 0;
    if (s->audio_enc) {
        bitrate += s->audio_enc->bit_rate;
        n++;
    }
    if (s->video_enc) {
        bitrate += s->video_enc->bit_rate;
        n++;
    }

    put_le32(pb, 1000000 / s->video_enc->rate);
    put_le32(pb, bitrate / 8); /* XXX: not quite exact */
    put_le32(pb, 0); /* padding */
    put_le32(pb, AVIF_ISINTERLEAVED); /* flags */
    put_le32(pb, 0); /* nb frames, filled later */
    put_le32(pb, 0); /* initial frame */
    put_le32(pb, n); /* nb streams */
    put_le32(pb, 0); /* suggested buffer size */
    put_le32(pb, s->video_enc->width);
    put_le32(pb, s->video_enc->height);
    put_le32(pb, 0); /* reserved */
    put_le32(pb, 0); /* reserved */
    put_le32(pb, 0); /* reserved */
    put_le32(pb, 0); /* reserved */
    
    /* stream list */
    for(i=0;i<n;i++) {
        list2 = start_tag(pb, "LIST");
        put_tag(pb, "strl");
    
        if (i == 0)
            stream = s->video_enc;
        else
            stream = s->audio_enc;

        /* stream generic header */
        strh = start_tag(pb, "strh");
        switch(stream->codec->type) {
        case CODEC_TYPE_VIDEO:
            put_tag(pb, "vids");
            put_le32(pb, codec_get_bmp_tag(stream->codec->id));
            put_le32(pb, 0); /* flags */
            put_le16(pb, 0); /* priority */
            put_le16(pb, 0); /* language */
            put_le32(pb, 0); /* XXX: initial frame ? */
            put_le32(pb, 1000); /* scale */
            put_le32(pb, 1000 * stream->rate); /* rate */
            put_le32(pb, 0); /* start */
            put_le32(pb, 0); /* length, filled later */
            put_le32(pb, 0); /* suggested buffer size */
            put_le32(pb, 10000); /* quality */
            put_le32(pb, stream->width * stream->height * 3); /* sample size */
            break;
        case CODEC_TYPE_AUDIO:
            put_tag(pb, "auds");
            put_le32(pb, 0);
            put_le32(pb, 0); /* flags */
            put_le16(pb, 0); /* priority */
            put_le16(pb, 0); /* language */
            put_le32(pb, 1); /* XXX: initial frame ? */
            put_le32(pb, 1); /* scale */
            put_le32(pb, stream->bit_rate / 8); /* rate */
            put_le32(pb, 0); /* start */
            put_le32(pb, 0); /* length, filled later */
            put_le32(pb, 0); /* suggested buffer size */
            put_le32(pb, -1); /* quality */
            put_le32(pb, 1); /* sample size */
            break;
        }
        put_le32(pb, 0);
        put_le32(pb, 0);
        end_tag(pb, strh);

        strf = start_tag(pb, "strf");
        switch(stream->codec->type) {
        case CODEC_TYPE_VIDEO:
            put_bmp_header(pb, stream);
            break;
        case CODEC_TYPE_AUDIO:
            put_wav_header(pb, stream);
            break;
        }
        end_tag(pb, strf);
        end_tag(pb, list2);
    }

    end_tag(pb, list1);

    avi->movi_list = start_tag(pb, "LIST");
    avi->first = NULL;
    avi->last = NULL;
    put_tag(pb, "movi");

    put_flush_packet(pb);

    return 0;
}

static void new_index(AVFormatContext *s, AVEncodeContext *enc,
                      UINT8 *buf, unsigned int size)
{
    AVIContext *avi = s->priv_data;
    PutByteContext *pb = &s->pb;
    AVIIndex *idx;
    unsigned char tag[5];
    unsigned int flags;

    tag[0] = '0';
    if (enc->codec->type == CODEC_TYPE_VIDEO) {
        tag[1] = '0';
        tag[2] = 'd';
        tag[3] = 'c';
        flags = enc->key_frame ? 0x10 : 0x00;
    } else {
        tag[1] = '1';
        tag[2] = 'w';
        tag[3] = 'b';
        flags = 0;
    }

    if (!s->is_streamed) {
        idx = malloc(sizeof(AVIIndex));
        memcpy(idx->tag, tag, 4);
        idx->flags = flags;
        idx->pos = put_pos(pb) - avi->movi_list;
        idx->len = size;
        idx->next = NULL;
        if (!avi->last)
            avi->first = idx;
        else
            avi->last->next = idx;
        avi->last = idx;
    }
    
    put_buffer(pb, tag, 4);
    put_le32(pb, size);
    put_buffer(pb, buf, size);
    if (size & 1)
        put_byte(pb, 0);

    put_flush_packet(pb);
}

static int avi_write_audio(AVFormatContext *s, UINT8 *buf, int size)
{
    new_index(s, s->audio_enc, buf, size);
    return 0;
}

static int avi_write_video(AVFormatContext *s, UINT8 *buf, int size)
{
    new_index(s, s->video_enc, buf, size);
    return 0;
}

static int avi_write_trailer(AVFormatContext *s)
{
    PutByteContext *pb = &s->pb;
    AVIContext *avi = s->priv_data;
    offset_t file_size, idx_chunk;
    AVIIndex *idx;

    if (!s->is_streamed) {
        end_tag(pb, avi->movi_list);

        idx_chunk = start_tag(pb, "idx1");
        idx = avi->first;
        while (idx != NULL) {
            put_buffer(pb, idx->tag, 4);
            put_le32(pb, idx->flags);
            put_le32(pb, idx->pos);
            put_le32(pb, idx->len);
            idx = idx->next;
        }
        end_tag(pb, idx_chunk);
        
        /* update file size */
        file_size = put_pos(pb);
        put_seek(pb, 4, SEEK_SET);
        put_le32(pb, file_size);
        put_seek(pb, file_size, SEEK_SET);
    }
    put_flush_packet(pb);

    free(avi);
    return 0;
}

AVFormat avi_format = {
    "avi",
    "avi format",
    "",
    "avi",
    CODEC_ID_MP2,
    CODEC_ID_DIVX,
    avi_write_header,
    avi_write_audio,
    avi_write_video,
    avi_write_trailer,
};