#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <fstream>
#include <string>
#include <stdexcept>
#include <cstdlib>
#include <locale>
#include <codecvt>
#include <windows.h>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

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

std::string readPath(const std::string& prompt) {
    std::string path;
    std::cout << prompt;
    std::getline(std::cin, path);
    if (path.size() >= 2 && path.front() == '"' && path.back() == '"') {
        path = path.substr(1, path.size() - 2);
    }
    return path;
}

void getInput() {
    inputPath = readPath("Enter source image path: ");
    outputPath = readPath("Enter output image path (.png): ");

    std::string channelInput;
    std::cout << "Channels to use (e.g. RGB, RG, B): ";
    std::getline(std::cin, channelInput);

    useR = channelInput.find('R') != std::string::npos || channelInput.find('r') != std::string::npos;
    useG = channelInput.find('G') != std::string::npos || channelInput.find('g') != std::string::npos;
    useB = channelInput.find('B') != std::string::npos || channelInput.find('b') != std::string::npos;

    if (!useR && !useG && !useB) {
        throw std::invalid_argument("At least one channel (R, G, or B) must be selected.");
    }

    std::string sourceChoice;
    std::cout << "Message from (T)ext input or (F)ile? ";
    std::getline(std::cin, sourceChoice);

    if (sourceChoice == "F" || sourceChoice == "f") {
        std::string filePath = readPath("Enter message file path: ");
        int wlen = MultiByteToWideChar(CP_ACP, 0, filePath.c_str(), -1, nullptr, 0);
        std::wstring wpath(wlen - 1, L'\0');
        MultiByteToWideChar(CP_ACP, 0, filePath.c_str(), -1, &wpath[0], wlen);
        std::ifstream file(wpath);
        if (!file.is_open()) {
            throw std::runtime_error("Failed to open message file: " + filePath);
        }
        message.assign(std::istreambuf_iterator<char>(file), std::istreambuf_iterator<char>());
    } else {
        std::cout << "Enter message: ";
        std::getline(std::cin, message);
    }

    if (message.empty()) {
        throw std::invalid_argument("Message is empty.");
    }
}

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));
    }

    int totalPixels = width * height;
    int activeChannels = useR + useG + useB;
    int availableBits = (totalPixels - 1) * activeChannels;
    int requiredBits = (static_cast<int>(message.size()) + 1) * 8;

    if (requiredBits > availableBits) {
        stbi_image_free(data);
        throw std::runtime_error(
            "Message too large. Need " + std::to_string(requiredBits) +
            " bits, but only " + std::to_string(availableBits) + " available."
        );
    }

    std::cout << "Image loaded: " << width << "x" << height << " (" << channels << " channels)\n";
    std::cout << "Capacity: " << availableBits << " bits available, " << requiredBits << " bits needed\n";
}

void writeChannelHeader() {
    data[0] = useR ? (data[0] | 1) : (data[0] & ~1);
    data[1] = useG ? (data[1] | 1) : (data[1] & ~1);
    data[2] = useB ? (data[2] | 1) : (data[2] & ~1);
}

void encodeBit(int channelOffset, int base, int& bitIndex) {
    int bytePos = bitIndex / 8;
    int bitPos = 7 - (bitIndex % 8);
    int bit = (message[bytePos] >> bitPos) & 1;
    data[base + channelOffset] = (data[base + channelOffset] & ~1) | bit;
    bitIndex++;
}

void encodeMessage() {
    message.push_back('\0');

    int totalPixels = width * height;
    int bitIndex = 0;
    int totalBits = static_cast<int>(message.size()) * 8;

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

        if (useR && bitIndex < totalBits) encodeBit(0, base, bitIndex);
        if (useG && bitIndex < totalBits) encodeBit(1, base, bitIndex);
        if (useB && bitIndex < totalBits) encodeBit(2, base, bitIndex);
    }

    std::cout << "Encoded " << bitIndex << " bits into image.\n";
}

void saveImage() {
    int result = stbi_write_png(outputPath.c_str(), width, height, channels, data, width * channels);
    stbi_image_free(data);

    if (!result) {
        throw std::runtime_error("Failed to write output image: " + outputPath);
    }

    std::cout << "Message successfully hidden in: " << outputPath << "\n";
}

int main() {
    try {
        getInput();
        loadImage();
        writeChannelHeader();
        encodeMessage();
        saveImage();
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        std::cin.get();
        return EXIT_FAILURE;
    }
    std::cin.get();
    return EXIT_SUCCESS;
}