/*
 * Basic user interface for ffmpeg system
 * Copyright (c) 2000,2001 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 <linux/videodev.h>
#include <linux/soundcard.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <errno.h>
#include <sys/time.h>
#include <getopt.h>

#include "mpegenc.h"

static AVFormat *file_format;
static int frame_width  = 160;
static int frame_height = 128;
static int frame_rate = 25;
static int bit_rate = 200000;
static int video_disable = 0;

static const char *video_filename, *audio_filename;
static float recording_time = 10.0;
static int nb_frames;
static int gop_size = 12;
static int intra_only = 0;
static int audio_freq = 44100;
static int audio_bit_rate = 64000;
static int audio_disable = 0;
static int audio_channels = 1;
static int high_quality = 0;

static long long time_start;

#define INFMT_YUV 1
#define INFMT_PGM 2

static inline int pnm_space(int c)  
{
    return (c==' ' || c=='\n' || c=='\r' || c=='\t');
}

static void pnm_get(FILE *f, char *str, int buf_size) 
{
    char *s;
    int c;
    
    do  {
        c=fgetc(f);
        if (c=='#')  {
            do  {
                c=fgetc(f);
            } while (c!='\n');
            c=fgetc(f);
        }
    } while (pnm_space(c));
    
    s=str;
    do  {
        *s++=c;
        c=fgetc(f);
    } while (!pnm_space(c));
    *s=0;
}

static int file_read_picture(AVEncodeContext *s, 
                             UINT8 *picture[3],
                             int width, int height,
                             int picture_number)
{
    FILE *f;
    char buf[1024], buf1[100];
    static int init = 0;
    static int format = 0;
    static UINT8 *pict[3];
    int i;

    if (!init) {
        pict[0] = malloc(width * height);
        pict[1] = malloc(width * height / 4);
        pict[2] = malloc(width * height / 4);
        init = 1;
    }
    
    picture[0] = pict[0];
    picture[1] = pict[1];
    picture[2] = pict[2];

    /* auto detect format */
    if (format == 0) {
        sprintf(buf, "%s%d.Y", video_filename, picture_number);
        f=fopen(buf, "r");
        if (f) {
            fclose(f);
            format = INFMT_YUV;
            goto ok;
        }
        sprintf(buf, "%s%d.pgm", video_filename, picture_number);
        f=fopen(buf, "r");
        if (f) {
            fclose(f);
            format = INFMT_PGM;
            goto ok;
        }
        fprintf(stderr, "Bad file name or unknown input video format\n");
        return -1;
    }
 ok:    
    switch(format) {
    case INFMT_YUV:
        sprintf(buf, "%s%d.Y", video_filename, picture_number);
        f=fopen(buf, "r");
        if (!f) {
            return -1;
        }
        
        fread(picture[0], 1, width * height, f);
        fclose(f);
        
        sprintf(buf, "%s%d.U", video_filename, picture_number);
        f=fopen(buf, "r");
        if (!f) {
            perror(buf);
            return -1;
        }
        fread(picture[1], 1, width * height / 4, f);
        fclose(f);
        
        sprintf(buf, "%s%d.V", video_filename, picture_number);
        f=fopen(buf, "r");
        if (!f) {
            perror(buf);
            return -1;
        }
        fread(picture[2], 1, width * height / 4, f);
        fclose(f);
        break;
    case INFMT_PGM:
        sprintf(buf, "%s%d.pgm", video_filename, picture_number);
        f=fopen(buf, "r");
        if (!f) {
            return -1;
        }
        pnm_get(f, buf1, sizeof(buf1));
        if (strcmp(buf1, "P5")) {
            fprintf(stderr, "%s: bad pgm file\n", buf);
            fclose(f);
            return -1;
        }
        pnm_get(f, buf1, sizeof(buf1));
        pnm_get(f, buf1, sizeof(buf1));
        pnm_get(f, buf1, sizeof(buf1));
        
        fread(picture[0], 1, width * height, f);
        for(i=0;i<height/2;i++) {
            fread(picture[1] + i * width/2, 1, width/2, f);
            fread(picture[2] + i * width/2, 1, width/2, f);
        }
        fclose(f);
        break;
    }
    return 0;
}

static void display_stats(AVEncodeContext *video_ctx, 
                          AVEncodeContext *audio_ctx,
                          int batch_mode, int the_end)
{
    if (video_ctx && 
        ((video_ctx->frame_number % video_ctx->rate) == 0 || 
        the_end)) {
        float ti;
        
        if (batch_mode) {
            ti = (float)video_ctx->frame_number / video_ctx->rate;
        } else {
            ti = (gettime() - time_start) / 1000000.0;
            if (ti < 0.1)
                    ti = 0.1;
        }
        
        fprintf(stderr,
                "frame=%5d size=%8dkB time=%0.1f fps=%4.1f bitrate=%6.1fkbits/s q=%2d\r", 
                video_ctx->frame_number,
                data_out_size / 1024,
                ti,
                video_ctx->frame_number / ti,
                (data_out_size * 8 / ti / 1000),
                video_ctx->quality);
        if (the_end) {
            fprintf(stderr,"\n");
        }
        fflush(stderr);
    }
#if 0
    if (the_end && batch_mode && audio_ctx) {
        duration = (gettime() - ti) / 1000000.0;
        factor = 0;
        if (ti > 0) {
            factor = (float)nb_samples / s->sample_rate / duration;
        }
        fprintf(stderr, "%0.1f seconds compressed in %0.1f seconds (speed factor: %0.1f)\n",
                (float)nb_samples / s->sample_rate, 
                duration,
                factor);
    }
#endif
}

void raw_write_data(void *opaque, 
                    unsigned char *buf, int size)
{
    FILE *outfile = opaque;
    fwrite(buf, 1, size, outfile);
    data_out_size += size;
}

int raw_seek(void *opaque, long long offset, int whence)
{
    FILE *outfile = opaque;
    fseek(outfile, offset, whence);
    return 0;
}

static void av_encode(AVFormatContext *ctx,
                      const char *video_filename,
                      const char *audio_filename)
{
    UINT8 audio_buffer[4096];
    UINT8 video_buffer[128*1024];
    char buf[256];
    short *samples;
    int ret;
    int audio_fd;
    FILE *infile;
    int sample_count;
    int batch_mode;
    AVEncodeContext *audio_enc, *video_enc;
    int frame_size, frame_bytes;
    AVEncoder *audio_encoder, *video_encoder;
    UINT8 *picture[3];

    /* audio */
    audio_enc = ctx->audio_enc;
    sample_count = 0;
    infile = NULL;
    frame_size = 0;
    samples = NULL;
    audio_fd = -1;
    frame_bytes = 0;
    batch_mode = 0;
    if (audio_filename ||
        video_filename)
        batch_mode = 1;
    
    if (audio_enc) {
        if (batch_mode) {
            if (!audio_filename) {
                fprintf(stderr, "Must give audio input file\n");
                exit(1);
            }
            infile = fopen(audio_filename, "r");
            if (!infile) {
                fprintf(stderr, "Could not open '%s'\n", audio_filename);
                exit(1);
            }
            audio_fd = -1;
        } else {
            audio_fd = audio_open(audio_enc->rate, audio_enc->channels);
            if (audio_fd < 0) {
                fprintf(stderr, "Could not open audio device\n");
                exit(1);
            }
        }
        
        audio_encoder = avencoder_find(ctx->format->audio_codec);
        if (avencoder_open(audio_enc, audio_encoder) < 0) {
            fprintf(stderr, "Audio encoder: incorrect audio frequency or bitrate\n");
            exit(1);
        }
        avencoder_string(buf, sizeof(buf), audio_enc);
        fprintf(stderr, "  %s\n", buf);
        
        frame_size = audio_enc->frame_size;
        
        frame_bytes = frame_size * 2 * audio_enc->channels;
        samples = malloc(frame_bytes);
    }

    /* video */
    video_enc = ctx->video_enc;
    if (video_enc) {
        if (batch_mode) {
            if (!video_filename) {
                fprintf(stderr, "Must give video input file\n");
                exit(1);
            }
        } else {
            ret = v4l_init(video_enc->rate, video_enc->width, video_enc->height);
            if (ret < 0) {
                fprintf(stderr,"Could not init video 4 linux capture\n");
                exit(1);
            }
        }
        
        video_encoder = avencoder_find(ctx->format->video_codec);
        if (avencoder_open(video_enc, video_encoder) < 0) {
            fprintf(stderr, "Error while initializing video codec\n");
            exit(1);
        }

        avencoder_string(buf, sizeof(buf), video_enc);
        fprintf(stderr, "  %s\n", buf);
    }
    
    ctx->format->write_header(ctx);
    time_start = gettime();

    for(;;) {
        /* read & compression audio frames */
        if (audio_enc) {
            if (!batch_mode) {
                for(;;) {
                    ret = read(audio_fd, samples, frame_bytes);
                    if (ret != frame_bytes)
                        break;
                    ret = avencoder_encode(audio_enc,
                                           audio_buffer, sizeof(audio_buffer), samples);
                    ctx->format->write_audio_frame(ctx, audio_buffer, ret);
                }
            } else {
                if (video_enc)
                    sample_count += audio_enc->rate / video_enc->rate;
                else
                    sample_count += frame_size;
                while (sample_count > frame_size) {
                    if (fread(samples, 1, frame_bytes, infile) == 0)
                        goto the_end;
                    
                    ret = avencoder_encode(audio_enc,
                                           audio_buffer, sizeof(audio_buffer), samples);
                    ctx->format->write_audio_frame(ctx, audio_buffer, ret);

                    sample_count -= frame_size;
                }
            }
        }

        if (video_enc) {
            /* read video image */
            if (batch_mode) {
                ret = file_read_picture (video_enc, picture, 
                                         video_enc->width, video_enc->height, 
                                         video_enc->frame_number);
            } else {
                ret = v4l_read_picture (picture, 
                                        video_enc->width, video_enc->height, 
                                        video_enc->frame_number);
            }
            if (ret < 0)
                break;
            ret = avencoder_encode(video_enc, video_buffer, sizeof(video_buffer), picture);
            ctx->format->write_video_picture(ctx, video_buffer, ret);
        }
        
        display_stats(video_enc, NULL, batch_mode, 0);
        if (video_enc && video_enc->frame_number >= nb_frames)
            break;
    }
 the_end:
    display_stats(video_enc, NULL, batch_mode, 1);

    if (video_enc)
        avencoder_close(video_enc);

    if (audio_enc)
        avencoder_close(audio_enc);
    
    ctx->format->write_trailer(ctx);

    if (!infile) {
        close(audio_fd);
    } else {
        fclose(infile);
    }
}

typedef struct {