#pragma once

#include "../extension/src/platform/EterLib/RenderCommands3D.h"
#include "../extension/src/platform/EterLib/UIRenderCommands.h"

#include <cstdint>
#include <cstring>
#include <fstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

// Diagnostic, little-endian arm64 format:
// - Version 1: matrices, positions, diffuse, indices.
// - Version 2: adds stable geometry_key and geometry_revision.
// - Version 3: adds normals, UVs, D3D8 fixed-function states, and embedded pack texture bytes.
// - Version 4: adds GPU skeletal skinning streams (bone_indices, bone_weights, bone_matrices)
//              and 2D UIRenderCommand stream + memory textures ("MTRA").
// - Version 5: adds per-draw fog color, mode, range and distance/density state.
// - Version 6: adds texture address and filtering state for both stages.
// - Version 7: adds two fixed-function lights and material color sources.
// Capture is opt-in and never runs in the normal game path.
namespace native_draw_capture {

struct Capture {
    std::uint32_t version = 7;
    std::vector<Render3DDraw> draws;
    std::unordered_map<std::string, std::vector<std::uint8_t>> textures;
    std::uint32_t ui_width = 960;
    std::uint32_t ui_height = 640;
    std::vector<UIRenderCommand> ui_commands;
};

inline std::vector<std::uint8_t> encode_raw_argb_as_mtra(
    std::uint32_t width,
    std::uint32_t height,
    const std::uint32_t* argb) {
    std::vector<std::uint8_t> out;
    if (!width || !height || !argb)
        return out;
    const std::size_t pixel_count = std::size_t(width) * std::size_t(height);
    out.resize(12 + pixel_count * 4);
    out[0] = 'M'; out[1] = 'T'; out[2] = 'R'; out[3] = 'A';
    std::memcpy(out.data() + 4, &width, 4);
    std::memcpy(out.data() + 8, &height, 4);
    std::uint8_t* dst = out.data() + 12;
    for (std::size_t i = 0; i < pixel_count; ++i) {
        const std::uint32_t c = argb[i];
        dst[i * 4 + 0] = static_cast<std::uint8_t>((c >> 16) & 0xffu);
        dst[i * 4 + 1] = static_cast<std::uint8_t>((c >> 8) & 0xffu);
        dst[i * 4 + 2] = static_cast<std::uint8_t>(c & 0xffu);
        dst[i * 4 + 3] = static_cast<std::uint8_t>((c >> 24) & 0xffu);
    }
    return out;
}

template <typename T> void write_scalar(std::ofstream& file, const T& value) {
    file.write(reinterpret_cast<const char*>(&value), sizeof(value));
}

template <typename T> void read_scalar(std::ifstream& file, T& value) {
    file.read(reinterpret_cast<char*>(&value), sizeof(value));
    if (!file) throw std::runtime_error("truncated native draw capture");
}

template <typename T>
void write_vector(std::ofstream& file, const std::vector<T>& values, std::size_t max_count = 4'000'000) {
    if (values.size() > max_count) throw std::runtime_error("native draw capture array too large");
    const auto count = static_cast<std::uint32_t>(values.size());
    write_scalar(file, count);
    if (count) file.write(reinterpret_cast<const char*>(values.data()), count * sizeof(T));
}

template <typename T>
void read_vector(std::ifstream& file, std::vector<T>& values, std::size_t max_count = 4'000'000) {
    std::uint32_t count = 0;
    read_scalar(file, count);
    if (count > max_count) throw std::runtime_error("native draw capture array too large");
    values.resize(count);
    if (count) file.read(reinterpret_cast<char*>(values.data()), count * sizeof(T));
    if (!file) throw std::runtime_error("truncated native draw capture array");
}

inline void write_string(std::ofstream& file, const std::string& value) {
    if (value.size() > 4096) throw std::runtime_error("native draw capture string too large");
    const auto size = static_cast<std::uint32_t>(value.size());
    write_scalar(file, size);
    if (size) file.write(value.data(), size);
}

inline void read_string(std::ifstream& file, std::string& value) {
    std::uint32_t size = 0;
    read_scalar(file, size);
    if (size > 4096) throw std::runtime_error("native draw capture string too large");
    value.resize(size);
    if (size) file.read(value.data(), size);
    if (!file) throw std::runtime_error("truncated native draw capture string");
}

inline void write(
    const std::string& path,
    const std::vector<Render3DDraw>& draws,
    const std::unordered_map<std::string, std::vector<std::uint8_t>>& textures = {},
    std::uint32_t ui_width = 960,
    std::uint32_t ui_height = 640,
    const std::vector<UIRenderCommand>& ui_commands = {}) {
    if (draws.size() > 10'000 || textures.size() > 8'192 || ui_commands.size() > 100'000)
        throw std::runtime_error("native draw capture has too many draws, textures, or UI commands");
    std::ofstream file(path, std::ios::binary | std::ios::trunc);
    if (!file) throw std::runtime_error("cannot create native draw capture: " + path);
    const std::uint32_t magic = 0x4d544452; // MTDR
    const std::uint32_t version = 7;
    const auto count = static_cast<std::uint32_t>(draws.size());
    write_scalar(file, magic); write_scalar(file, version); write_scalar(file, count);
    for (const auto& draw : draws) {
        file.write(reinterpret_cast<const char*>(draw.world), sizeof(draw.world));
        file.write(reinterpret_cast<const char*>(draw.view), sizeof(draw.view));
        file.write(reinterpret_cast<const char*>(draw.proj), sizeof(draw.proj));
        write_scalar(file, draw.geometry_key);
        write_scalar(file, draw.geometry_revision);
        const std::uint32_t flags =
            (draw.lines ? 1u : 0u) | (draw.pretransformed ? 2u : 0u) | (draw.light0 ? 4u : 0u);
        write_scalar(file, flags);
        write_vector(file, draw.positions);
        write_vector(file, draw.diffuse);
        write_vector(file, draw.indices);
        file.write(reinterpret_cast<const char*>(draw.viewport), sizeof(draw.viewport));
        write_string(file, draw.texture0);
        write_string(file, draw.texture1);
        write_vector(file, draw.rhw);
        write_vector(file, draw.normals);
        write_vector(file, draw.uv0);
        write_vector(file, draw.uv1);
        write_scalar(file, draw.alpha_blend);
        write_scalar(file, draw.src_blend);
        write_scalar(file, draw.dest_blend);
        write_scalar(file, draw.alpha_test);
        write_scalar(file, draw.alpha_ref);
        write_scalar(file, draw.alpha_func);
        write_scalar(file, draw.cull_mode);
        write_scalar(file, draw.z_enable);
        write_scalar(file, draw.z_write);
        write_scalar(file, draw.z_func);
        write_scalar(file, draw.lighting);
        write_scalar(file, draw.texture_factor);
        write_scalar(file, draw.fog_enable);
        file.write(reinterpret_cast<const char*>(draw.color_op), sizeof(draw.color_op));
        file.write(reinterpret_cast<const char*>(draw.color_arg1), sizeof(draw.color_arg1));
        file.write(reinterpret_cast<const char*>(draw.color_arg2), sizeof(draw.color_arg2));
        file.write(reinterpret_cast<const char*>(draw.alpha_op), sizeof(draw.alpha_op));
        file.write(reinterpret_cast<const char*>(draw.alpha_arg1), sizeof(draw.alpha_arg1));
        file.write(reinterpret_cast<const char*>(draw.alpha_arg2), sizeof(draw.alpha_arg2));
        file.write(reinterpret_cast<const char*>(draw.material_diffuse), sizeof(draw.material_diffuse));
        file.write(reinterpret_cast<const char*>(draw.material_ambient), sizeof(draw.material_ambient));
        file.write(reinterpret_cast<const char*>(draw.material_emissive), sizeof(draw.material_emissive));
        file.write(reinterpret_cast<const char*>(draw.light0_direction), sizeof(draw.light0_direction));
        file.write(reinterpret_cast<const char*>(draw.light0_diffuse), sizeof(draw.light0_diffuse));
        file.write(reinterpret_cast<const char*>(draw.light0_ambient), sizeof(draw.light0_ambient));
        write_scalar(file, draw.ambient);
        write_scalar(file, draw.fog_color);
        write_scalar(file, draw.fog_vertex_mode);
        write_scalar(file, draw.fog_table_mode);
        write_scalar(file, draw.fog_range_enable);
        write_scalar(file, draw.fog_start);
        write_scalar(file, draw.fog_end);
        write_scalar(file, draw.fog_density);
        file.write(reinterpret_cast<const char*>(draw.address_u), sizeof(draw.address_u));
        file.write(reinterpret_cast<const char*>(draw.address_v), sizeof(draw.address_v));
        file.write(reinterpret_cast<const char*>(draw.min_filter), sizeof(draw.min_filter));
        file.write(reinterpret_cast<const char*>(draw.mag_filter), sizeof(draw.mag_filter));
        file.write(reinterpret_cast<const char*>(draw.mip_filter), sizeof(draw.mip_filter));
        file.write(reinterpret_cast<const char*>(draw.lights), sizeof(draw.lights));
        write_scalar(file, draw.diffuse_material_source);
        write_scalar(file, draw.ambient_material_source);
        write_scalar(file, draw.color_vertex);
        write_vector(file, draw.bone_indices);
        write_vector(file, draw.bone_weights);
        write_vector(file, draw.bone_matrices);
    }
    const auto texture_count = static_cast<std::uint32_t>(textures.size());
    write_scalar(file, texture_count);
    for (const auto& [name, bytes] : textures) {
        write_string(file, name);
        write_vector(file, bytes, 16'777'216);
    }
    write_scalar(file, ui_width);
    write_scalar(file, ui_height);
    const auto ui_count = static_cast<std::uint32_t>(ui_commands.size());
    write_scalar(file, ui_count);
    for (const auto& cmd : ui_commands) {
        const auto kind = static_cast<std::uint32_t>(cmd.kind);
        write_scalar(file, kind);
        write_scalar(file, cmd.x1);
        write_scalar(file, cmd.y1);
        write_scalar(file, cmd.x2);
        write_scalar(file, cmd.y2);
        write_scalar(file, cmd.argb);
        write_scalar(file, cmd.end_argb);
        write_scalar(file, cmd.clip_x1);
        write_scalar(file, cmd.clip_y1);
        write_scalar(file, cmd.clip_x2);
        write_scalar(file, cmd.clip_y2);
        write_string(file, cmd.text);
        const std::uint32_t uiflags = (cmd.quad ? 1u : 0u) | (cmd.behind_3d ? 2u : 0u);
        write_scalar(file, uiflags);
        file.write(reinterpret_cast<const char*>(cmd.qx), sizeof(cmd.qx));
        file.write(reinterpret_cast<const char*>(cmd.qy), sizeof(cmd.qy));
        write_scalar(file, cmd.su);
        write_scalar(file, cmd.sv);
        write_scalar(file, cmd.eu);
        write_scalar(file, cmd.ev);
        write_scalar(file, cmd.blend);
        write_string(file, cmd.mask);
        file.write(reinterpret_cast<const char*>(cmd.mu), sizeof(cmd.mu));
        file.write(reinterpret_cast<const char*>(cmd.mv), sizeof(cmd.mv));
    }
    if (!file) throw std::runtime_error("failed to write native draw capture: " + path);
}

inline Capture read_capture(const std::string& path) {
    std::ifstream file(path, std::ios::binary);
    if (!file) throw std::runtime_error("cannot open native draw capture: " + path);
    std::uint32_t magic = 0, version = 0, count = 0;
    read_scalar(file, magic); read_scalar(file, version); read_scalar(file, count);
    if (magic != 0x4d544452 || (version < 1 || version > 7) || count > 10'000)
        throw std::runtime_error("unsupported native draw capture format");
    Capture capture;
    capture.version = version;
    capture.draws.resize(count);
    for (auto& draw : capture.draws) {
        file.read(reinterpret_cast<char*>(draw.world), sizeof(draw.world));
        file.read(reinterpret_cast<char*>(draw.view), sizeof(draw.view));
        file.read(reinterpret_cast<char*>(draw.proj), sizeof(draw.proj));
        if (!file) throw std::runtime_error("truncated native draw capture matrices");
        if (version >= 2) {
            read_scalar(file, draw.geometry_key);
            read_scalar(file, draw.geometry_revision);
        }
        std::uint32_t flags = 0;
        read_scalar(file, flags);
        draw.lines = (flags & 1u) != 0;
        draw.pretransformed = (flags & 2u) != 0;
        draw.light0 = (flags & 4u) != 0;
        read_vector(file, draw.positions);
        read_vector(file, draw.diffuse);
        read_vector(file, draw.indices);
        if (version >= 3) {
            file.read(reinterpret_cast<char*>(draw.viewport), sizeof(draw.viewport));
            read_string(file, draw.texture0);
            read_string(file, draw.texture1);
            read_vector(file, draw.rhw);
            read_vector(file, draw.normals);
            read_vector(file, draw.uv0);
            read_vector(file, draw.uv1);
            read_scalar(file, draw.alpha_blend);
            read_scalar(file, draw.src_blend);
            read_scalar(file, draw.dest_blend);
            read_scalar(file, draw.alpha_test);
            read_scalar(file, draw.alpha_ref);
            read_scalar(file, draw.alpha_func);
            read_scalar(file, draw.cull_mode);
            read_scalar(file, draw.z_enable);
            read_scalar(file, draw.z_write);
            read_scalar(file, draw.z_func);
            read_scalar(file, draw.lighting);
            read_scalar(file, draw.texture_factor);
            read_scalar(file, draw.fog_enable);
            file.read(reinterpret_cast<char*>(draw.color_op), sizeof(draw.color_op));
            file.read(reinterpret_cast<char*>(draw.color_arg1), sizeof(draw.color_arg1));
            file.read(reinterpret_cast<char*>(draw.color_arg2), sizeof(draw.color_arg2));
            file.read(reinterpret_cast<char*>(draw.alpha_op), sizeof(draw.alpha_op));
            file.read(reinterpret_cast<char*>(draw.alpha_arg1), sizeof(draw.alpha_arg1));
            file.read(reinterpret_cast<char*>(draw.alpha_arg2), sizeof(draw.alpha_arg2));
            file.read(reinterpret_cast<char*>(draw.material_diffuse), sizeof(draw.material_diffuse));
            file.read(reinterpret_cast<char*>(draw.material_ambient), sizeof(draw.material_ambient));
            file.read(reinterpret_cast<char*>(draw.material_emissive), sizeof(draw.material_emissive));
            file.read(reinterpret_cast<char*>(draw.light0_direction), sizeof(draw.light0_direction));
            file.read(reinterpret_cast<char*>(draw.light0_diffuse), sizeof(draw.light0_diffuse));
            file.read(reinterpret_cast<char*>(draw.light0_ambient), sizeof(draw.light0_ambient));
            read_scalar(file, draw.ambient);
            if (version >= 5) {
                read_scalar(file, draw.fog_color);
                read_scalar(file, draw.fog_vertex_mode);
                read_scalar(file, draw.fog_table_mode);
                read_scalar(file, draw.fog_range_enable);
                read_scalar(file, draw.fog_start);
                read_scalar(file, draw.fog_end);
                read_scalar(file, draw.fog_density);
            }
            if (version >= 6) {
                file.read(reinterpret_cast<char*>(draw.address_u), sizeof(draw.address_u));
                file.read(reinterpret_cast<char*>(draw.address_v), sizeof(draw.address_v));
                file.read(reinterpret_cast<char*>(draw.min_filter), sizeof(draw.min_filter));
                file.read(reinterpret_cast<char*>(draw.mag_filter), sizeof(draw.mag_filter));
                file.read(reinterpret_cast<char*>(draw.mip_filter), sizeof(draw.mip_filter));
            }
            if (version >= 7) {
                file.read(reinterpret_cast<char*>(draw.lights), sizeof(draw.lights));
                read_scalar(file, draw.diffuse_material_source);
                read_scalar(file, draw.ambient_material_source);
                read_scalar(file, draw.color_vertex);
            } else if (draw.light0) {
                auto& light = draw.lights[0];
                light.type = 3;
                std::memcpy(light.direction, draw.light0_direction, sizeof(light.direction));
                std::memcpy(light.diffuse, draw.light0_diffuse, sizeof(light.diffuse));
                std::memcpy(light.ambient, draw.light0_ambient, sizeof(light.ambient));
            }
            if (!file) throw std::runtime_error("truncated native draw capture state");
        } else {
            draw.z_enable = 1;
            draw.z_write = 1;
        }
        if (version >= 4) {
            read_vector(file, draw.bone_indices);
            read_vector(file, draw.bone_weights);
            read_vector(file, draw.bone_matrices);
        }
    }
    if (version >= 3) {
        std::uint32_t texture_count = 0;
        read_scalar(file, texture_count);
        if (texture_count > 8'192) throw std::runtime_error("native draw capture has too many textures");
        for (std::uint32_t i = 0; i < texture_count; ++i) {
            std::string name;
            std::vector<std::uint8_t> bytes;
            read_string(file, name);
            read_vector(file, bytes, 16'777'216);
            capture.textures.emplace(std::move(name), std::move(bytes));
        }
    }
    if (version >= 4) {
        read_scalar(file, capture.ui_width);
        read_scalar(file, capture.ui_height);
        std::uint32_t ui_count = 0;
        read_scalar(file, ui_count);
        if (ui_count > 100'000) throw std::runtime_error("native draw capture has too many UI commands");
        capture.ui_commands.resize(ui_count);
        for (auto& cmd : capture.ui_commands) {
            std::uint32_t kind = 0, uiflags = 0;
            read_scalar(file, kind);
            cmd.kind = static_cast<UIRenderCommand::Kind>(kind);
            read_scalar(file, cmd.x1);
            read_scalar(file, cmd.y1);
            read_scalar(file, cmd.x2);
            read_scalar(file, cmd.y2);
            read_scalar(file, cmd.argb);
            read_scalar(file, cmd.end_argb);
            read_scalar(file, cmd.clip_x1);
            read_scalar(file, cmd.clip_y1);
            read_scalar(file, cmd.clip_x2);
            read_scalar(file, cmd.clip_y2);
            read_string(file, cmd.text);
            read_scalar(file, uiflags);
            cmd.quad = (uiflags & 1u) != 0;
            cmd.behind_3d = (uiflags & 2u) != 0;
            file.read(reinterpret_cast<char*>(cmd.qx), sizeof(cmd.qx));
            file.read(reinterpret_cast<char*>(cmd.qy), sizeof(cmd.qy));
            read_scalar(file, cmd.su);
            read_scalar(file, cmd.sv);
            read_scalar(file, cmd.eu);
            read_scalar(file, cmd.ev);
            read_scalar(file, cmd.blend);
            read_string(file, cmd.mask);
            file.read(reinterpret_cast<char*>(cmd.mu), sizeof(cmd.mu));
            file.read(reinterpret_cast<char*>(cmd.mv), sizeof(cmd.mv));
            if (!file) throw std::runtime_error("truncated native draw capture UI command");
        }
    }
    return capture;
}

inline std::vector<Render3DDraw> read(const std::string& path) {
    return read_capture(path).draws;
}

} // namespace native_draw_capture