//---------------------------------------------------------------------------
//  Screenshot - Copyright © 2001 Michael Fötsch.
//
//  Source code for my tutorial "Taking Screenshots with Direct3D 8"
//  at www.mr-gamemaker.com.
//
//  Visit my homepage: http://www.geocities.com/foetsch
//  Contact me: foetsch@yahoo.com
//
//  You can download the Borland-compatible LIB files for DirectX 8
//  from http://www.geocities.com/foetsch.
//---------------------------------------------------------------------------

#include "Screenshot.h"
#include <stdexcept>
#include <string>
using namespace std;


#define SAFERELEASE(x) if (x) { x->Release(); x = NULL; x; }
#define SAFEDELETEARRAY(x) if (x) { delete[] x; x = NULL; x; }

// path to screenshot files, relative to the folder that contains the
// EXE file (used if the lpszFileName parameter of the Screenshot function
// is NULL)
// If, for example, you want to save the screenshots to the sub-folder
// "shots", set SCREENSHOT_REL_PATH to "shots\\".
#ifndef SCREENSHOT_REL_PATH
#define SCREENSHOT_REL_PATH ""
#endif

// bitmap file signature
#ifdef __BORLANDC__
#define BITMAP_FILE_SIGNATURE 'BM'
#else
#define BITMAP_FILE_SIGNATURE 'MB'
#endif


// GetScreenshotFileName
//  [out] FileName: This string is filled with a unique filename that
//      can be used to store the screenshot. The filename is assembled
//      as follows: 1. The function takes the folder that contains the EXE.
//      2. The function appends the string in SCREENSHOT_REL_PATH. 3. The
//      function searches for the first unused filename of the form
//      "shotXXX.bmp", where XXX is a number from 000 to 999, and appends
//      it to the path.
//
void GetScreenshotFileName(string& FileName)
{
    // retrieve the filename of the EXE file
    string ModuleFileName;
    ModuleFileName.reserve(MAX_PATH);
    GetModuleFileName(
        NULL, const_cast<char*>(ModuleFileName.data()), MAX_PATH);
    // extract the path info from the filename
    FileName = ModuleFileName.substr(0, ModuleFileName.find_last_of(":\\"));
    // append the sub-folder path
    FileName += SCREENSHOT_REL_PATH;

    // search for first unused filename
    char Buf[MAX_PATH];
    WIN32_FIND_DATA ffd;
    HANDLE h;
    for (int i = 0; i < 1000; i++)
    {
        // prepare search mask for FindFirstFile
        wsprintf(Buf, (FileName + "shot%03i.bmp").c_str(), i);
        // check whether this file already exists
        h = FindFirstFile(Buf, &ffd);
        // if the file exists, close the search handle and continue
        if (h != INVALID_HANDLE_VALUE)
        {   FindClose(h); }
        // if the file does not exist, exit from the loop
        else
        {   break; }
    }

    // set FileName to the first unused filename
    FileName = Buf;
}


// THROW_FWRITE_EXCEPTION
//  [in] FileName: Name of the file to which the write operation failed.
//
void THROW_FWRITE_EXCEPTION(const char* FileName)
{
    char Msg[MAX_PATH + 256];
    wsprintf(
        Msg,
        "Function \"Screenshot\": "
         "Could not write all data to file \"%s\".",
        FileName
    );
    throw runtime_error(Msg);
}


// Screenshot
//  For a parameter description, see the header file.
//
HRESULT Screenshot(
    LPDIRECT3DDEVICE8 lpDevice,
    const char* lpszFileName
)
{
    LPDIRECT3DSURFACE8 lpSurface = NULL;    // copy of the front buffer
    LPDIRECT3D8 lpD3D = NULL;   // IDirect3D8 interface that created lpDevice
    LPBYTE Bits = NULL;         // bitmap bits of the DIB
    FILE* f = NULL;             // BMP file handle

    // check whether parameters are valid
    if (!lpDevice)
    {
        throw invalid_argument("Function \"Screenshot\": "
            "LPDIRECT3DDEVICE8 lpDevice is NULL.");
    }

try
{
    string FileName;            // name of the BMP file

    // if a filename is given, use it
    if (lpszFileName)
    {   FileName = lpszFileName; }
    // if no filename is given, have a filename generated
    else
    {   GetScreenshotFileName(FileName); }

    // We need to know for which adapter the Direct3D device was created,
    // so that we can query the right adapter's screen dimensions.
    // Therefore, query the device's creation parameters
    D3DDEVICE_CREATION_PARAMETERS dcp;
    dcp.AdapterOrdinal = D3DADAPTER_DEFAULT;
    lpDevice->GetCreationParameters(&dcp);

    D3DDISPLAYMODE dm;
    dm.Width = dm.Height = 0;

    // retrieve pointer to IDirect3D8 interface,
    // which provides the GetAdapterDisplayMode method
    lpDevice->GetDirect3D(&lpD3D);
    if (lpD3D)
    {
        // query the screen dimensions of the current adapter
        lpD3D->GetAdapterDisplayMode(dcp.AdapterOrdinal, &dm);
        SAFERELEASE(lpD3D);
    }

    // create a 32-bit ARGB system-memory surface that is going to receive
    // a copy of the front buffer's contents
    if (FAILED(lpDevice->CreateImageSurface(
        dm.Width, dm.Height,
        D3DFMT_A8R8G8B8,
        &lpSurface
    )))
    {
        throw runtime_error("Function \"Screenshot\": "
            "IDirect3DDevice8::CreateImageSurface failed.");
    }

    // have the GetFrontBuffer method copy the contents of the front
    // buffer to our system-memory surface
    if (FAILED(lpDevice->GetFrontBuffer(lpSurface)))
    {
        throw runtime_error("Function \"Screenshot\": "
            "IDirect3DDevice8::GetFrontBuffer failed.");
    }

    // prepare the bitmap info header
    BITMAPINFOHEADER bmih;
    bmih.biSize = sizeof(bmih);
    bmih.biWidth = dm.Width;
    bmih.biHeight = dm.Height;
    bmih.biPlanes = 1;
    bmih.biBitCount = 24;
    bmih.biCompression = BI_RGB;
    bmih.biSizeImage = dm.Width * dm.Height * 3;
    bmih.biXPelsPerMeter = 0;
    bmih.biYPelsPerMeter = 0;
    bmih.biClrUsed = 0;
    bmih.biClrImportant = 0;

    // reserve memory for the DIB's bitmap bits
    // (The extra byte is needed because the bitmap is 24-bit but we're
    // going to write 32 bits (a DWORD) at a time. When we write the
    // last pixel, we'll exceed the array's limit if we don't reserve an
    // extra byte.)
    Bits = new BYTE[bmih.biSizeImage + 1];
    if (!Bits)
    {
        throw runtime_error("Function \"Screenshot\": "
            "Could not allocate memory for bitmap bits.");
    }

    // lock the surface for reading
    D3DLOCKED_RECT LockedRect;
    if (FAILED(lpSurface->LockRect(&LockedRect, NULL, D3DLOCK_READONLY)))
    {
        throw runtime_error("Function \"Screenshot\": "
            "IDirect3DSurface8::LockRect failed.");
    }

    // flip the bitmap vertically (because that's how DIBs are stored)
    // and convert it from 32-bits to 24-bits (some bitmap viewers can't
    // handle 32-bit bitmaps, although it's a valid format)

    LPDWORD lpSrc;
    LPBYTE lpDest = Bits;

    // read pixels beginning with the bottom scan line
    for (int y = dm.Height - 1; y >= 0; y--)
    {
        // calculate address of the current source scan line
        lpSrc = reinterpret_cast<LPDWORD>(LockedRect.pBits) + y * dm.Width;
        for (int x = 0; x < dm.Width; x++)
        {
            // store the source pixel in the bitmap bits array
            *reinterpret_cast<LPDWORD>(lpDest) = *lpSrc;
            lpSrc++;        // increment source pointer by 1 DWORD
            lpDest += 3;    // increment destination pointer by 3 bytes
        }
    }

    // we can unlock and release the surface
    lpSurface->UnlockRect();
    SAFERELEASE(lpSurface);

    // prepare the bitmap file header
    BITMAPFILEHEADER bmfh;
    bmfh.bfType = BITMAP_FILE_SIGNATURE;
    bmfh.bfSize = sizeof(bmfh) + sizeof(bmih) + bmih.biSizeImage;
    bmfh.bfReserved1 = bmfh.bfReserved2 = 0;
    bmfh.bfOffBits = sizeof(bmfh) + sizeof(bmih);

    // create the BMP file
    f = fopen(FileName.c_str(), "wb");
    if (!f)
    {
        char Msg[MAX_PATH + 256];
        wsprintf(
            Msg,
            "Function \"Screenshot\": "
             "The file \"%s\" could not be created.",
            FileName.c_str()
        );
        throw runtime_error(Msg);
    }

    // dump the file header
    if (!fwrite(reinterpret_cast<void*>(&bmfh), sizeof(bmfh), 1, f))
    {   THROW_FWRITE_EXCEPTION(FileName.c_str()); }
    // dump the info header
    if (!fwrite(reinterpret_cast<void*>(&bmih), sizeof(bmih), 1, f))
    {   THROW_FWRITE_EXCEPTION(FileName.c_str()); }
    // dump the bitmap bits
    if (fwrite(reinterpret_cast<void*>(Bits), sizeof(BYTE),
        bmih.biSizeImage, f) < bmih.biSizeImage)
    {   THROW_FWRITE_EXCEPTION(FileName.c_str()); }

    // close the file
    fclose(f);

    // free the memory for the bitmap bits
    SAFEDELETEARRAY(Bits);

    return S_OK;

}
catch (...)
{
    // on error, release all local resources
    if (lpSurface) { lpSurface->UnlockRect(); }
    SAFERELEASE(lpSurface);
    SAFERELEASE(lpD3D);
    SAFEDELETEARRAY(Bits);
    if (f) { fclose(f); }

#ifdef SCREENSHOT_DONT_THROW
    return E_FAIL;
#else
    throw;
#endif
}
}