#include <windows.h>
#include <string.h>
#include <math.h>
#include "avisynth.h"

extern "C"
{
    int  __stdcall HC_AVS_LOADLIB(int *);
    int  __stdcall HC_AVS_FREELIB(void);
    int  __stdcall HC_AVS_INIT(char *, int, int *, int *, int *, int *, int *, int *, int *, char *, int);
    int  __stdcall HC_AVS_GETFRAME_PLANAR(int , int , int , unsigned char *, unsigned char *, unsigned char *);
    int  __stdcall HC_AVS_GETFRAME_PACKED(int , int , unsigned char *);

    int  __stdcall HC_D2V_LOADLIB(char*, int);
    int  __stdcall HC_D2V_FREELIB(void);
    int  __stdcall HC_D2V_INIT(char *, int, int *, int *, int *, int *, int *, int *, int *);
    void __stdcall HC_D2V_GETFRAME(int , unsigned char *, unsigned char *, unsigned char *);
    void __stdcall HC_D2V_CLOSEFILE(void);
}

/////////////////////////////////////////////////////////////////////////////
////////////////////////////// A V S   S T U F F ////////////////////////////
/////////////////////////////////////////////////////////////////////////////

IScriptEnvironment *env = 0;
AVSValue val = 0;
PClip clip = 0;
HINSTANCE avsDLL = 0;

//////////////////////////////////////////////////////

int __stdcall HC_AVS_LOADLIB(int *setmem)
{
    avsDLL = LoadLibrary("avisynth.dll");
    if (!avsDLL) return 1;

    IScriptEnvironment* (* CreateScriptEnvironment)(int version)
             = (IScriptEnvironment*(*)(int)) GetProcAddress(avsDLL, "CreateScriptEnvironment");
    if (!CreateScriptEnvironment) return 2;
    env = CreateScriptEnvironment(AVISYNTH_INTERFACE_VERSION);
    if (!env) return 2;

    if (*setmem >= 0) *setmem = env->SetMemoryMax(*setmem);

    return 0;
}

//////////////////////////////////////////////////////

int __stdcall HC_AVS_FREELIB(void)
{
    val = 0;
    clip = 0;

    if (avsDLL)
    {
        int iret = FreeLibrary(avsDLL);
        if (!iret) return 7;
    }
    return 0;
}

//////////////////////////////////////////////////////

int __stdcall HC_AVS_INIT(char *fname, int lfname, int *w, int *h, int *nframes,
                          int *fpsn, int *fpsd, int *pixtype, int *info, char *errmsg, int lerrmsg)
{
    if (!avsDLL) return 1;

    AVSValue arg[1] = { fname };
    try
    {
        val = env->Invoke("Import", AVSValue(arg,1));
        clip = val.AsClip();
    }
    catch (AvisynthError err)
    {
        strncpy(errmsg, err.msg, lerrmsg);
        return 3;
    }

    if (!clip->GetVideoInfo().HasVideo()) return (4);

    if (w)       *w =       clip->GetVideoInfo().width;
    if (h)       *h =       clip->GetVideoInfo().height;
    if (nframes) *nframes = clip->GetVideoInfo().num_frames;
    if (fpsn)    *fpsn =    clip->GetVideoInfo().fps_numerator;
    if (fpsd)    *fpsd =    clip->GetVideoInfo().fps_denominator;
    if (pixtype) *pixtype = clip->GetVideoInfo().pixel_type; 

    *info = 0;
    if (clip->GetVideoInfo().IsYV12())  *info = 1;
    if (clip->GetVideoInfo().IsYUY2())  *info = 2;

    return 0;
}

//////////////////////////////////////////////////////

int __stdcall HC_AVS_GETFRAME_PLANAR(int frame, int Ypitch, int UVpitch, unsigned char *Y, unsigned char *U, unsigned char *V)
{
    PVideoFrame f;

    if (!avsDLL) return 6;
    try { f = clip->GetFrame(frame, env); } catch(...) { return 6; }

    env->BitBlt(Y,Ypitch, f->GetReadPtr(PLANAR_Y),f->GetPitch(PLANAR_Y),f->GetRowSize(PLANAR_Y),f->GetHeight(PLANAR_Y));
    env->BitBlt(U,UVpitch,f->GetReadPtr(PLANAR_U),f->GetPitch(PLANAR_U),f->GetRowSize(PLANAR_U),f->GetHeight(PLANAR_U));
    env->BitBlt(V,UVpitch,f->GetReadPtr(PLANAR_V),f->GetPitch(PLANAR_V),f->GetRowSize(PLANAR_V),f->GetHeight(PLANAR_V));

    return 0;
}

//////////////////////////////////////////////////////

int __stdcall HC_AVS_GETFRAME_PACKED(int frame, int npix, unsigned char *YUV)
{
    PVideoFrame f;

    if (!avsDLL) return 6;
    try { f = clip->GetFrame(frame, env); } catch(...) { return 6; }

    memcpy(YUV, f->GetReadPtr(PLANAR_Y), npix);

    return 0;
}

/////////////////////////////////////////////////////////////////////////////
////////////////////////////// D 2 V   S T U F F ////////////////////////////
/////////////////////////////////////////////////////////////////////////////

VideoInfo* (__cdecl *openMPEG2Source)(char*);
unsigned char* (__cdecl *getFrame)(int);
void (__cdecl *closeVideo)();
VideoInfo *vi;
HINSTANCE d2vDLL;

static int Ypl, UVpl, GetProc, LoadD2VLib;

int __stdcall HC_D2V_LOADLIB(char *fnam, int lfnam)
{
    vi=0; GetProc=0; LoadD2VLib=0;
    d2vDLL = LoadLibrary(fnam);
    if (!d2vDLL) return 128;
    LoadD2VLib=1;
    return 0;
}

//////////////////////////////////////////////////////

int __stdcall HC_D2V_FREELIB(void)
{
    if (LoadD2VLib)
    {
        int iret = FreeLibrary(d2vDLL);
        if (!iret) return 129;
    }
    return 0;
}

//////////////////////////////////////////////////////

void __stdcall HC_D2V_CLOSEFILE(void)
{
    if (vi) closeVideo();
}

//////////////////////////////////////////////////////

int __stdcall HC_D2V_INIT(char *fnam, int lfnam, int *w, int *h, int *nframes, int *fpsn, int *fpsd, int *pixtype, int *info)
{
    if ( !d2vDLL ) return 128;

    if (GetProc == 0)         // get routine addresses from DGDecode.dll
    {
        openMPEG2Source = (VideoInfo*  (__cdecl *) (char*)) GetProcAddress(d2vDLL, "openMPEG2Source");
        getFrame = (unsigned char*  (__cdecl *) (int)) GetProcAddress(d2vDLL, "getFrame");
        closeVideo = (void  (__cdecl *) ()) GetProcAddress(d2vDLL, "closeVideo");
        if ( !openMPEG2Source || !getFrame || !closeVideo ) return (130);
        GetProc=1;
    }

    vi = openMPEG2Source(fnam);
    if (!vi) {return 131;}

    if (w)       *w =       vi->width;
    if (h)       *h =       vi->height;
    if (nframes) *nframes = vi->num_frames;
    if (fpsn)    *fpsn =    vi->fps_numerator;
    if (fpsd)    *fpsd =    vi->fps_denominator;
    if (pixtype) *pixtype = vi->pixel_type;

    *info = 0;
    if (vi->IsYV12())  *info = 1;
    if (vi->IsYUY2())  *info = 2;

    if (*w == 0 || *h == 0 || *nframes == 0 || *fpsn == 0 || *fpsd == 0) {vi=0; return 132;}
    Ypl = vi->width * vi->height;
    if (vi->IsYV12()) UVpl = Ypl >> 2;
    if (vi->IsYUY2()) UVpl = Ypl >> 1;

    return 0;
}

//////////////////////////////////////////////////////

void __stdcall HC_D2V_GETFRAME(int HC_frame, unsigned char *HC_Y, unsigned char *HC_U, unsigned char *HC_V)
{
    if (!LoadD2VLib) return;
    const unsigned char* f = getFrame(HC_frame);
    memcpy(HC_Y, f, Ypl);  f += Ypl;
    memcpy(HC_U, f, UVpl); f += UVpl;
    memcpy(HC_V, f, UVpl);
}