/*
ScreenshotQuickTake v20201129
By DarkRanger <darkranger.no-ip.org>
License: MIT License
*/
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <windows.h>
#include <winuser.h>
#include <wingdi.h>

#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#define STBI_ONLY_BMP

#include "stb_image.h"
#include "stb_image_write.h"

#define DEFAULT_FILENAME_LENGTH 30

char *get_time(void) {
    time_t rawtime;
    struct tm *timeinfo;
    static char current_time[DEFAULT_FILENAME_LENGTH];
    
    time(&rawtime);
    timeinfo = localtime(&rawtime);
    strftime(current_time, DEFAULT_FILENAME_LENGTH, "%Y_%m%d_%H%M%S", timeinfo);
    
    return current_time;
}

int save_to_png(BITMAPFILEHEADER bmfHeader, BITMAPINFOHEADER bi, char *lpbitmap, DWORD dwBmpSize, int width, int height) {
    unsigned char *buffer;
    unsigned char *image;
    int components;
    char file_name[DEFAULT_FILENAME_LENGTH] = "screenshot.png";
    
    //sprintf(file_name, "%s.png", get_time());
    
    buffer = malloc(sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + dwBmpSize);
    memcpy(buffer, &bmfHeader, sizeof(BITMAPFILEHEADER));
    memcpy(buffer + sizeof(BITMAPFILEHEADER), &bi, sizeof(BITMAPINFOHEADER));
    memcpy(buffer + sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER), lpbitmap, dwBmpSize);
    
    //image = stbi_load("screenshot.bmp", &width, &height, &components, 0);
    
    image = stbi_load_from_memory(buffer, 
          sizeof(BITMAPINFOHEADER) + sizeof(BITMAPFILEHEADER) + dwBmpSize, 
          &width, &height, &components, 0);
    stbi_write_png(file_name, width, height, components, image, 0);
    stbi_image_free(image);
    free(buffer);
    
    return 0;
}

int save_to_bmp(BITMAPFILEHEADER bmfHeader, BITMAPINFOHEADER bi, char *lpbitmap, DWORD dwBmpSize) {
    LPCTSTR file_name = "screenshot.bmp";
    DWORD dwBytesWritten;
    
    dwBytesWritten = 0;
    HANDLE hFile = CreateFile(file_name, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    WriteFile(hFile, (LPSTR)&bmfHeader, sizeof(BITMAPFILEHEADER), &dwBytesWritten, NULL);
    WriteFile(hFile, (LPSTR)&bi, sizeof(BITMAPINFOHEADER), &dwBytesWritten, NULL);
    WriteFile(hFile, (LPSTR)lpbitmap, dwBmpSize, &dwBytesWritten, NULL);
    CloseHandle(hFile);
    
    return 0;
}

int screenshot_quick_take(void) {
    HDC hdcScreen = NULL;
    HDC hdcMemDC = NULL;
    HBITMAP hbmScreen = NULL;
    BITMAP bmpScreen;
    BITMAPFILEHEADER bmfHeader;
    BITMAPINFOHEADER bi;
    DWORD dwBmpSize;
    HANDLE hDIB;
    char *lpbitmap;
    DWORD dwSizeofDIB;
    int width = 0, height = 0;
    
    hdcScreen = GetDC(NULL);
    width = GetDeviceCaps(hdcScreen, DESKTOPHORZRES);
    height = GetDeviceCaps(hdcScreen, DESKTOPVERTRES);
    hdcMemDC = CreateCompatibleDC(hdcScreen);
    hbmScreen = CreateCompatibleBitmap(hdcScreen, width, height);
    SelectObject(hdcMemDC, hbmScreen);
    BitBlt(hdcMemDC, 0, 0, width, height, hdcScreen, 0, 0, SRCCOPY | CAPTUREBLT);
    GetObject(hbmScreen, sizeof(BITMAP), &bmpScreen);
     
    bi.biSize = sizeof(BITMAPINFOHEADER);
    bi.biWidth = bmpScreen.bmWidth;
    bi.biHeight = bmpScreen.bmHeight;
    bi.biPlanes = 1;
    bi.biBitCount = 24;
    bi.biCompression = BI_RGB;
    bi.biSizeImage = 0;
    bi.biXPelsPerMeter = 0;
    bi.biYPelsPerMeter = 0;
    bi.biClrUsed = 0;
    bi.biClrImportant = 0;
    
    dwBmpSize = ((bmpScreen.bmWidth * bi.biBitCount + 31) / 32) * 4 * bmpScreen.bmHeight;
    hDIB = GlobalAlloc(GHND, dwBmpSize);
    lpbitmap = (char *)GlobalLock(hDIB);
    
    GetDIBits(hdcScreen, hbmScreen, 0, (UINT)bmpScreen.bmHeight, lpbitmap, (BITMAPINFO *)&bi, DIB_RGB_COLORS);
    dwSizeofDIB = dwBmpSize + sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
    bmfHeader.bfOffBits = (DWORD)sizeof(BITMAPFILEHEADER) + (DWORD)sizeof(BITMAPINFOHEADER); 
    bmfHeader.bfSize = dwSizeofDIB; 
    bmfHeader.bfType = 0x4D42;
    
    //save_to_bmp(bmfHeader, bi, lpbitmap, dwBmpSize);
    save_to_png(bmfHeader, bi, lpbitmap, dwBmpSize, width, height);
    
    GlobalUnlock(hDIB);
    GlobalFree(hDIB);
    
    DeleteObject(hbmScreen);
    DeleteDC(hdcMemDC);
    ReleaseDC(NULL, hdcScreen);
    
    printf("screenshot saved.\n");
    
    return 0;
}

int main(int argc, char *argv[]) {
    int second = 0;
    
    if (argc == 1) {
        screenshot_quick_take();
    }
    else {
        second = atoi(argv[1]);

        if (second > 0) {
            while (1) {
                Sleep(second * 1000);
                screenshot_quick_take();
            }
        }
        else {
            screenshot_quick_take();
        }
    }
    
    return 0;
}

