#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <string>
#include <stdexcept>
#include <cstdlib>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

int width, height, channels;
std::string inputPath, message;
bool useR, useG, useB;
unsigned char* data;

void loadImage() {
    data = stbi_load(inputPath.c_str(), &width, &height, &channels, 0);
    if (!data) {
        throw std::runtime_error("Failed to load image: " + inputPath);
    }

    if (channels < 3) {
        stbi_image_free(data);
        throw std::runtime_error("Image must have at least 3 channels (RGB). Found: " + std::to_string(channels));
    }
}

void getInputPath() {
    std::cout << "Enter image file path: ";
    std::getline(std::cin, inputPath);
    if (inputPath.size() >= 2 && inputPath.front() == '"' && inputPath.back() == '"') {
        inputPath = inputPath.substr(1, inputPath.size() - 2);
    }
}

void readChannelHeaders() {
    useR = data[0] & 1;
    useG = data[1] & 1;
    useB = data[2] & 1;

    if (!useR && !useG && !useB) {
        stbi_image_free(data);
        throw std::runtime_error("No active channels found in header pixel. Image may not contain a hidden message.");
    }

    std::cout << "Active channels: " << (useR ? "R" : "") << (useG ? "G" : "") << (useB ? "B" : "") << "\n";
}

bool processBit(int channelOffset, int base, unsigned char& currentByte, int& bitCount) {
    int bit = data[base + channelOffset] & 1;
    currentByte = (currentByte << 1) | bit;
    bitCount++;
    if (bitCount == 8) {
        if (currentByte == '\0') return true; // null terminator found
        message.push_back(static_cast<char>(currentByte));
        currentByte = 0;
        bitCount = 0;
    }
    return false;
}

void readMessage() {
    int totalPixels = width * height;
    unsigned char currentByte = 0;
    int bitCount = 0;

    for (int pixel = 1; pixel < totalPixels; pixel++) {
        int base = pixel * channels;

        if (useR && processBit(0, base, currentByte, bitCount)) return;
        if (useG && processBit(1, base, currentByte, bitCount)) return;
        if (useB && processBit(2, base, currentByte, bitCount)) return;
    }
}

int main() {
    try {
        getInputPath();
        loadImage();
        readChannelHeaders();
        readMessage();

        stbi_image_free(data);

        if (message.empty()) {
            std::cout << "No message found (or message is empty).\n";
        } else {
            std::cout << "Hidden message (" << message.size() << " bytes):\n";
            std::cout << message << "\n";
        }

    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        std::cin.get();
        return EXIT_FAILURE;
    }
    std::cin.get();
    return EXIT_SUCCESS;
}