#include "pch.h"

#include "Texture.h"
#include <stdexcept>
#include <algorithm>
#include <cmath>

#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb_image.h>
#include <stb_image_write.h>


namespace Vivid2D::Render::Texture {

    // 静态成员初始化
    std::unordered_map<std::string, std::weak_ptr<Texture>> Texture::s_TextureCache;
    std::mutex Texture::s_CacheMutex;

    // ==================== 辅助函数实现 ====================

    GLint getGLInternalFormat(TextureFormat format) {
        switch (format) {
        case TextureFormat::RED: return GL_R8;
        case TextureFormat::RG: return GL_RG8;
        case TextureFormat::RGB: return GL_RGB8;
        case TextureFormat::RGBA: return GL_RGBA8;
        case TextureFormat::ALPHA: return GL_ALPHA;
        case TextureFormat::LUMINANCE: return GL_LUMINANCE;
        case TextureFormat::LUMINANCE_ALPHA: return GL_LUMINANCE_ALPHA;
        }
        return GL_RGBA8;
    }

    GLenum getGLFormat(TextureFormat format) {
        switch (format) {
        case TextureFormat::RED: return GL_RED;
        case TextureFormat::RG: return GL_RG;
        case TextureFormat::RGB: return GL_RGB;
        case TextureFormat::RGBA: return GL_RGBA;
        case TextureFormat::ALPHA: return GL_ALPHA;
        case TextureFormat::LUMINANCE: return GL_LUMINANCE;
        case TextureFormat::LUMINANCE_ALPHA: return GL_LUMINANCE_ALPHA;
        }
        return GL_RGBA;
    }

    GLenum getGLType(TextureType type) {
        switch (type) {
        case TextureType::UNSIGNED_BYTE: return GL_UNSIGNED_BYTE;
        case TextureType::BYTE: return GL_BYTE;
        case TextureType::UNSIGNED_SHORT: return GL_UNSIGNED_SHORT;
        case TextureType::SHORT: return GL_SHORT;
        case TextureType::UNSIGNED_INT: return GL_UNSIGNED_INT;
        case TextureType::INT: return GL_INT;
        case TextureType::FLOAT: return GL_FLOAT;
        }
        return GL_UNSIGNED_BYTE;
    }

    GLint getGLFilter(TextureFilter filter) {
        switch (filter) {
        case TextureFilter::NEAREST: return GL_NEAREST;
        case TextureFilter::LINEAR: return GL_LINEAR;
        case TextureFilter::NEAREST_MIPMAP_NEAREST: return GL_NEAREST_MIPMAP_NEAREST;
        case TextureFilter::LINEAR_MIPMAP_NEAREST: return GL_LINEAR_MIPMAP_NEAREST;
        case TextureFilter::NEAREST_MIPMAP_LINEAR: return GL_NEAREST_MIPMAP_LINEAR;
        case TextureFilter::LINEAR_MIPMAP_LINEAR: return GL_LINEAR_MIPMAP_LINEAR;
        }
        return GL_LINEAR;
    }

    GLint getGLWrap(TextureWrap wrap) {
        switch (wrap) {
        case TextureWrap::REPEAT: return GL_REPEAT;
        case TextureWrap::MIRRORED_REPEAT: return GL_MIRRORED_REPEAT;
        case TextureWrap::CLAMP_TO_EDGE: return GL_CLAMP_TO_EDGE;
        case TextureWrap::CLAMP_TO_BORDER: return GL_CLAMP_TO_BORDER;
        }
        return GL_CLAMP_TO_EDGE;
    }

    int getComponentCount(TextureFormat format) {
        switch (format) {
        case TextureFormat::RED:
        case TextureFormat::ALPHA:
        case TextureFormat::LUMINANCE: return 1;
        case TextureFormat::RG:
        case TextureFormat::LUMINANCE_ALPHA: return 2;
        case TextureFormat::RGB: return 3;
        case TextureFormat::RGBA: return 4;
        }
        return 4;
    }

    TextureFormat getTextureFormatFromComponents(int components) {
        switch (components) {
        case 1: return TextureFormat::RED;
        case 2: return TextureFormat::RG;
        case 3: return TextureFormat::RGB;
        case 4: return TextureFormat::RGBA;
        default: throw std::invalid_argument("Unsupported number of components: " + std::to_string(components));
        }
    }

    // ImageData 构造函数
    ImageData::ImageData(unsigned char* d, int w, int h, TextureFormat f, std::string path)
        : data(d, stbi_image_free), width(w), height(h), format(f), sourcePath(std::move(path)) {}


    // ==================== Texture 类实现 ====================

    Texture::Texture(std::string name, int width, int height, TextureFormat format, TextureType type)
        : m_Name(std::move(name)), m_Width(width), m_Height(height), m_Format(format), m_Type(type) {
        m_CreationTime = std::chrono::system_clock::now().time_since_epoch().count();
        m_TextureId = generateTextureId();
        createTextureObject();
        applyTextureParameters();
    }

    Texture::Texture(std::string name, int width, int height, TextureFormat format, const unsigned char* pixelData)
        : Texture(std::move(name), width, height, format, TextureType::UNSIGNED_BYTE) {
        size_t size = width * height * getComponentCount(format);
        uploadData(pixelData, size);
    }

    Texture::~Texture() {
        dispose();
    }

    Texture::Texture(Texture&& other) noexcept
        : m_TextureId(other.m_TextureId), m_Name(std::move(other.m_Name)), m_Width(other.m_Width), m_Height(other.m_Height),
        m_Format(other.m_Format), m_Type(other.m_Type), m_MinFilter(other.m_MinFilter), m_MagFilter(other.m_MagFilter),
        m_WrapS(other.m_WrapS), m_WrapT(other.m_WrapT), m_MipmapsEnabled(other.m_MipmapsEnabled),
        m_Disposed(other.m_Disposed), m_CreationTime(other.m_CreationTime), m_PixelDataCache(std::move(other.m_PixelDataCache))
    {
        other.m_TextureId = 0;
        other.m_Disposed = true;
    }

    Texture& Texture::operator=(Texture&& other) noexcept {
        if (this != &other) {
            dispose();
            m_TextureId = other.m_TextureId;
            m_Name = std::move(other.m_Name);
            m_Width = other.m_Width;
            m_Height = other.m_Height;
            m_Format = other.m_Format;
            m_Type = other.m_Type;
            m_MinFilter = other.m_MinFilter;
            m_MagFilter = other.m_MagFilter;
            m_WrapS = other.m_WrapS;
            m_WrapT = other.m_WrapT;
            m_MipmapsEnabled = other.m_MipmapsEnabled;
            m_Disposed = other.m_Disposed;
            m_CreationTime = other.m_CreationTime;
            m_PixelDataCache = std::move(other.m_PixelDataCache);

            other.m_TextureId = 0;
            other.m_Disposed = true;
        }
        return *this;
    }

    void Texture::createTextureObject() {
        bind();
        RenderSystem::texImage2D(GL_TEXTURE_2D, 0, getGLInternalFormat(m_Format), m_Width, m_Height, 0, getGLFormat(m_Format), getGLType(m_Type), nullptr);
        unbind();
        RenderSystem::checkGLError("Texture::createTextureObject");
    }

    void Texture::uploadData(const std::vector<unsigned char>& pixelData) {
        uploadData(pixelData.data(), pixelData.size());
    }

    void Texture::uploadData(const unsigned char* pixelData, size_t size) {
        if (m_Disposed) throw std::runtime_error("Cannot upload data to disposed texture: " + m_Name);
        if (!pixelData) throw std::invalid_argument("Pixel data cannot be null");

        size_t expectedSize = m_Width * m_Height * getComponentCount(m_Format);
        if (size < expectedSize) {
            throw std::invalid_argument("Pixel data buffer too small for texture dimensions.");
        }

        bind();
        RenderSystem::pixelStore(GL_UNPACK_ALIGNMENT, 1);

        glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, m_Width, m_Height, getGLFormat(m_Format), getGLType(m_Type), pixelData);

        RenderSystem::checkGLError("glTexSubImage2D");
        RenderSystem::pixelStore(GL_UNPACK_ALIGNMENT, 4);
        unbind();

        m_PixelDataCache.assign(pixelData, pixelData + expectedSize);
    }

    void Texture::generateMipmaps() {
        if (m_Disposed) throw std::runtime_error("Cannot generate mipmaps for disposed texture: " + m_Name);
        if (!isPowerOfTwo(m_Width) || !isPowerOfTwo(m_Height)) {
            // Log warning instead of throwing
            // std::cerr << "Warning: Cannot generate mipmaps for non-power-of-two texture: " << m_Name << std::endl;
            return;
        }

        bind();
        glGenerateMipmap(GL_TEXTURE_2D);
        RenderSystem::checkGLError("glGenerateMipmap");
        unbind();

        m_MipmapsEnabled = true;

        if (m_MinFilter == TextureFilter::LINEAR) {
            setMinFilter(TextureFilter::LINEAR_MIPMAP_LINEAR);
        }
        else if (m_MinFilter == TextureFilter::NEAREST) {
            setMinFilter(TextureFilter::NEAREST_MIPMAP_NEAREST);
        }
    }

    std::unique_ptr<Texture> Texture::crop(int x, int y, int w, int h, const std::string& newName) {
        if (m_Disposed) throw std::runtime_error("Cannot crop disposed texture: " + m_Name);
        if (x < 0 || y < 0 || w <= 0 || h <= 0 || x + w > m_Width || y + h > m_Height) {
            throw std::invalid_argument("Crop rectangle out of bounds");
        }
        if (m_Type != TextureType::UNSIGNED_BYTE) {
            throw std::runtime_error("Crop currently only supported for UNSIGNED_BYTE textures");
        }

        ensurePixelDataCached();
        if (!hasPixelData()) {
            throw std::runtime_error("No pixel data available for cropping texture: " + m_Name);
        }

        int comps = getComponentCount(m_Format);
        int srcStride = m_Width * comps;
        int dstStride = w * comps;
        std::vector<unsigned char> croppedData(w * h * comps);

        for (int row = 0; row < h; ++row) {
            const unsigned char* srcRow = m_PixelDataCache.data() + ((y + row) * m_Width + x) * comps;
            unsigned char* dstRow = croppedData.data() + row * dstStride;
            memcpy(dstRow, srcRow, dstStride);
        }

        auto newTex = std::make_unique<Texture>(newName, w, h, m_Format, croppedData.data());
        newTex->setMinFilter(m_MinFilter);
        newTex->setMagFilter(m_MagFilter);
        newTex->setWrapS(m_WrapS);
        newTex->setWrapT(m_WrapT);

        if (m_MipmapsEnabled && isPowerOfTwo(w) && isPowerOfTwo(h)) {
            newTex->generateMipmaps();
        }
        return newTex;
    }


    void Texture::setMinFilter(TextureFilter filter) {
        if (m_MinFilter != filter) {
            m_MinFilter = filter;
            applyTextureParameters();
        }
    }

    void Texture::setMagFilter(TextureFilter filter) {
        if (m_MagFilter != filter) {
            m_MagFilter = filter;
            applyTextureParameters();
        }
    }

    void Texture::setWrapS(TextureWrap wrap) {
        if (m_WrapS != wrap) {
            m_WrapS = wrap;
            applyTextureParameters();
        }
    }

    void Texture::setWrapT(TextureWrap wrap) {
        if (m_WrapT != wrap) {
            m_WrapT = wrap;
            applyTextureParameters();
        }
    }

    void Texture::setWrap(TextureWrap wrapS, TextureWrap wrapT) {
        setWrapS(wrapS);
        setWrapT(wrapT);
    }

    void Texture::setWrap(TextureWrap wrap) {
        setWrap(wrap, wrap);
    }

    void Texture::applyTextureParameters() {
        if (m_Disposed) return;
        bind();
        RenderSystem::texParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, getGLFilter(m_MinFilter));
        RenderSystem::texParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, getGLFilter(m_MagFilter));
        RenderSystem::texParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, getGLWrap(m_WrapS));
        RenderSystem::texParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, getGLWrap(m_WrapT));
        unbind();
    }

    void Texture::bind(int textureUnit) const {
        if (m_Disposed) throw std::runtime_error("Cannot bind disposed texture: " + m_Name);
        RenderSystem::activeTexture(GL_TEXTURE0 + textureUnit);
        RenderSystem::bindTexture(m_TextureId);
    }

    void Texture::unbind() const {
        RenderSystem::bindTexture(0);
    }

    void Texture::dispose() {
        if (!m_Disposed) {
            RenderSystem::deleteTextures(m_TextureId);
            m_TextureId = 0;
            m_Disposed = true;
            m_PixelDataCache.clear();

            std::lock_guard<std::mutex> lock(s_CacheMutex);
            s_TextureCache.erase(m_Name);
        }
    }

    bool Texture::isDisposed() const {
        return m_Disposed;
    }

    bool Texture::hasPixelData() const {
        return !m_PixelDataCache.empty();
    }

    std::vector<unsigned char> Texture::getPixelData() {
        ensurePixelDataCached();
        return m_PixelDataCache;
    }

    void Texture::setPixelData(const std::vector<unsigned char>& data) {
        m_PixelDataCache = data;
    }

    void Texture::clearPixelDataCache() {
        m_PixelDataCache.clear();
        m_PixelDataCache.shrink_to_fit();
    }

    void Texture::ensurePixelDataCached() {
        if (!hasPixelData()) {
            cachePixelDataFromGPU();
        }
    }

    size_t Texture::getPixelDataCacheSize() const {
        return m_PixelDataCache.size();
    }

    std::vector<unsigned char> Texture::extractTextureData() {
        if (m_Disposed) throw std::runtime_error("Cannot extract data from a disposed texture.");

        size_t dataSize = m_Width * m_Height * getComponentCount(m_Format);
        std::vector<unsigned char> data(dataSize);

        bind();
        RenderSystem::getTexImage(GL_TEXTURE_2D, 0, getGLFormat(m_Format), getGLType(m_Type), data.data());
        unbind();

        return data;
    }

    void Texture::cachePixelDataFromGPU() {
        if (m_Disposed) return;
        m_PixelDataCache = extractTextureData();
    }

    bool Texture::saveToFile(const std::string& filePath, const std::string& format) {
        if (m_Disposed) throw std::runtime_error("Cannot save a disposed texture.");

        ensurePixelDataCached();
        if (!hasPixelData()) {
            return false;
        }

        int comps = getComponentCount(m_Format);
        int stride = m_Width * comps;
        int success = 0;

        stbi_flip_vertically_on_write(true);

        if (format == "png") {
            success = stbi_write_png(filePath.c_str(), m_Width, m_Height, comps, m_PixelDataCache.data(), stride);
        }
        else if (format == "jpg" || format == "jpeg") {
            success = stbi_write_jpg(filePath.c_str(), m_Width, m_Height, comps, m_PixelDataCache.data(), 90);
        }
        else if (format == "bmp") {
            success = stbi_write_bmp(filePath.c_str(), m_Width, m_Height, comps, m_PixelDataCache.data());
        }
        else if (format == "tga") {
            success = stbi_write_tga(filePath.c_str(), m_Width, m_Height, comps, m_PixelDataCache.data());
        }
        else {
            throw std::invalid_argument("Unsupported image format for saving: " + format);
        }

        return success != 0;
    }

    std::unique_ptr<Texture> Texture::copy(const std::string& newName) const {
        if (m_Disposed) throw std::runtime_error("Cannot copy disposed texture");

        const_cast<Texture*>(this)->ensurePixelDataCached();
        if (!hasPixelData()) {
            throw std::runtime_error("No pixel data available to copy texture: " + m_Name);
        }

        auto newTex = std::make_unique<Texture>(newName, m_Width, m_Height, m_Format, m_PixelDataCache.data());
        newTex->setMinFilter(m_MinFilter);
        newTex->setMagFilter(m_MagFilter);
        newTex->setWrapS(m_WrapS);
        newTex->setWrapT(m_WrapT);

        if (m_MipmapsEnabled) {
            newTex->generateMipmaps();
        }
        return newTex;
    }

    long Texture::getEstimatedMemoryUsage() const {
        int bytesPerComponent = 1;
        switch (m_Type) {
        case TextureType::UNSIGNED_SHORT:
        case TextureType::SHORT:
            bytesPerComponent = 2;
            break;
        case TextureType::UNSIGNED_INT:
        case TextureType::INT:
        case TextureType::FLOAT:
            bytesPerComponent = 4;
            break;
        default:
            bytesPerComponent = 1;
        }
        long baseMemory = (long)m_Width * m_Height * getComponentCount(m_Format) * bytesPerComponent;
        if (m_MipmapsEnabled) {
            return baseMemory * 4L / 3L;
        }
        return baseMemory;
    }

    // ==================== 静态方法实现 ====================

    std::shared_ptr<Texture> Texture::createSolidColor(const std::string& name, int width, int height, uint32_t rgbaColor) {
        std::vector<uint32_t> pixels(width * height, rgbaColor);
        auto tex = std::make_shared<Texture>(name, width, height, TextureFormat::RGBA, reinterpret_cast<const unsigned char*>(pixels.data()));
        return tex;
    }

    std::shared_ptr<Texture> Texture::createCheckerboard(const std::string& name, int width, int height, int tileSize, uint32_t color1, uint32_t color2) {
        std::vector<uint32_t> pixels(width * height);
        for (int y = 0; y < height; ++y) {
            for (int x = 0; x < width; ++x) {
                bool isColor1 = ((x / tileSize) + (y / tileSize)) % 2 == 0;
                pixels[y * width + x] = isColor1 ? color1 : color2;
            }
        }
        auto tex = std::make_shared<Texture>(name, width, height, TextureFormat::RGBA, reinterpret_cast<const unsigned char*>(pixels.data()));
        return tex;
    }

    std::shared_ptr<Texture> Texture::getOrCreate(const std::string& name, int width, int height, TextureFormat format) {
        std::lock_guard<std::mutex> lock(s_CacheMutex);
        if (s_TextureCache.count(name)) {
            if (auto shared = s_TextureCache[name].lock()) {
                return shared;
            }
        }
        auto newTex = std::make_shared<Texture>(name, width, height, format);
        s_TextureCache[name] = newTex;
        return newTex;
    }

    std::unique_ptr<ImageData> Texture::loadImageFromFile(const std::string& filePath) {
        int width, height, components;
        unsigned char* data = stbi_load(filePath.c_str(), &width, &height, &components, 0);
        if (!data) {
            throw std::runtime_error("Failed to load image: " + filePath + " - " + stbi_failure_reason());
        }
        TextureFormat format = getTextureFormatFromComponents(components);
        return std::make_unique<ImageData>(data, width, height, format, filePath);
    }

    std::shared_ptr<Texture> Texture::createFromFile(const std::string& name, const std::string& filePath, TextureFilter minFilter, TextureFilter magFilter) {
        auto imageData = loadImageFromFile(filePath);
        auto texture = std::make_shared<Texture>(name, imageData->width, imageData->height, imageData->format, imageData->data.get());
        texture->setMinFilter(minFilter);
        texture->setMagFilter(magFilter);

        if (isPowerOfTwo(imageData->width) && isPowerOfTwo(imageData->height)) {
            texture->generateMipmaps();
        }

        std::lock_guard<std::mutex> lock(s_CacheMutex);
        s_TextureCache[name] = texture;
        return texture;
    }

    std::shared_ptr<Texture> Texture::createFromBytes(const std::string& name, const std::vector<unsigned char>& imageData, TextureFilter minFilter, TextureFilter magFilter) {
        int width, height, components;
        unsigned char* pixelData = stbi_load_from_memory(imageData.data(), static_cast<int>(imageData.size()), &width, &height, &components, 0);

        if (!pixelData) {
            throw std::runtime_error("Failed to load image from bytes: " + std::string(stbi_failure_reason()));
        }

        auto texture = std::make_shared<Texture>(name, width, height, getTextureFormatFromComponents(components), pixelData);
        texture->setMinFilter(minFilter);
        texture->setMagFilter(magFilter);
        stbi_image_free(pixelData);

        std::lock_guard<std::mutex> lock(s_CacheMutex);
        s_TextureCache[name] = texture;
        return texture;
    }

    std::shared_ptr<Texture> Texture::createFromARGB(const std::string& name, int width, int height, const uint32_t* pixels, bool isPremultiplied, TextureFilter minFilter, TextureFilter magFilter) {
        if (!pixels) throw std::invalid_argument("Pixel data cannot be null");

        std::vector<unsigned char> rgbaData(width * height * 4);
        for (int i = 0; i < width * height; ++i) {
            uint32_t p = pixels[i];
            int a = (p >> 24) & 0xFF;
            int r = (p >> 16) & 0xFF;
            int g = (p >> 8) & 0xFF;
            int b = p & 0xFF;

            if (isPremultiplied && a != 0) {
                float invA = 255.0f / static_cast<float>(a);
                r = std::min(255, static_cast<int>(round(r * invA)));
                g = std::min(255, static_cast<int>(round(g * invA)));
                b = std::min(255, static_cast<int>(round(b * invA)));
            }

            rgbaData[i * 4 + 0] = static_cast<unsigned char>(r);
            rgbaData[i * 4 + 1] = static_cast<unsigned char>(g);
            rgbaData[i * 4 + 2] = static_cast<unsigned char>(b);
            rgbaData[i * 4 + 3] = static_cast<unsigned char>(a);
        }

        auto texture = std::make_shared<Texture>(name, width, height, TextureFormat::RGBA, rgbaData.data());
        texture->setMinFilter(minFilter);
        texture->setMagFilter(magFilter);

        if (isPowerOfTwo(width) && isPowerOfTwo(height)) {
            texture->generateMipmaps();
        }

        std::lock_guard<std::mutex> lock(s_CacheMutex);
        s_TextureCache[name] = texture;
        return texture;
    }

    bool Texture::isSupportedImageFormat(const std::string& filePath) {
        std::string lowerPath = filePath;
        std::transform(lowerPath.begin(), lowerPath.end(), lowerPath.begin(), ::tolower);
        const char* extensions[] = { ".png", ".jpg", ".jpeg", ".bmp", ".tga", ".psd", ".gif", ".hdr", ".pic" };
        for (const auto* ext : extensions) {
            if (lowerPath.length() >= strlen(ext) && lowerPath.compare(lowerPath.length() - strlen(ext), strlen(ext), ext) == 0) {
                return true;
            }
        }
        return false;
    }

    ImageInfo Texture::getImageInfo(const std::string& filePath) {
        int width, height, components;
        if (!stbi_info(filePath.c_str(), &width, &height, &components)) {
            throw std::runtime_error("Failed to get image info: " + filePath + " - " + stbi_failure_reason());
        }
        return { width, height, components, filePath };
    }

    void Texture::cleanupAll() {
        std::lock_guard<std::mutex> lock(s_CacheMutex);
        for (auto const& [key, val] : s_TextureCache) {
            if (auto shared = val.lock()) {
                shared->dispose();
            }
        }
        s_TextureCache.clear();
    }

    void Texture::stbi_vertical_flip_on_load(bool flip) {
        stbi_set_flip_vertically_on_load(flip);
    }

    // ==================== 私有方法实现 ====================

    GLuint Texture::generateTextureId() {
        return RenderSystem::genTextures();
    }

    bool Texture::isPowerOfTwo(int value) {
        return value > 0 && (value & (value - 1)) == 0;
    }

} // namespace Vivid2D::Texture