#pragma once

#include "RenderCommands3D.h"
#include "UIRenderCommands.h"
#include "perf_log.h"
#include "render_state.h"
#include "touch_controller.h"

#include <SDL3/SDL.h>
#include <vulkan/vulkan.h>

#include <array>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>

namespace mt_host {

// The SDL3 window, Vulkan device and swapchain, and the renderer that replays a frame's recorded
// D3D8 draws (Render3DDraw) and UI commands.
class VulkanWindow {
    struct FrameSlot;
    struct GeometryId {
        std::uint64_t key = 0, signature = 0;
        bool operator==(const GeometryId&) const = default;
    };
    struct GeometryIdHash {
        std::size_t operator()(GeometryId id) const {
            return std::size_t(id.key ^ (id.signature + 0x9e3779b97f4a7c15ull + (id.key << 6) + (id.key >> 2)));
        }
    };
    struct Geometry {
        VkBuffer buffer = VK_NULL_HANDLE;
        VkDeviceMemory memory = VK_NULL_HANDLE;
        VkDeviceSize index_offset = 0;
        std::uint64_t last_used_frame = 0;
        std::uint32_t vertex_count = 0, index_count = 0;
    };
    struct GpuTexture {
        VkImage image = VK_NULL_HANDLE;
        VkDeviceMemory memory = VK_NULL_HANDLE;
        VkImageView view = VK_NULL_HANDLE;
        VkDescriptorSet descriptor = VK_NULL_HANDLE;
        VkDescriptorSet ui_descriptor = VK_NULL_HANDLE;
        std::uint64_t last_used_frame = 0;
        std::uint64_t last_decode_attempt_frame = 0;
    };
    struct PairedDescriptor {
        VkDescriptorSet set = VK_NULL_HANDLE;
        std::uint64_t last_used_frame = 0;
    };
    // An offscreen D3D render-target texture (the 40250 character shadow map, "rt:<id>:<w>x<h>"):
    // drawn by the frame's offscreen pass ahead of the back-buffer pass and sampled by that pass.
    // Between frames the colour image stays SHADER_READ_ONLY_OPTIMAL and the depth image
    // DEPTH_STENCIL_ATTACHMENT_OPTIMAL; the render pass keeps (LOAD/STORE) both, like a D3D surface.
    struct RenderTarget {
        GpuTexture texture; // colour image, view and descriptors used for sampling
        VkImage depth_image = VK_NULL_HANDLE;
        VkDeviceMemory depth_memory = VK_NULL_HANDLE;
        VkImageView depth_view = VK_NULL_HANDLE;
        VkFramebuffer framebuffer = VK_NULL_HANDLE;
        std::uint32_t width = 0, height = 0;
        bool initialized = false; // cleared to white / depth 1 by the first frame that uses it
    };
    struct PreparedDraw {
        Geometry* geometry = nullptr;
        VkPipeline pipeline = VK_NULL_HANDLE;
        VkDescriptorSet descriptor = VK_NULL_HANDLE;
        PushConstants constants{};
        FixedFunctionState fixed_state{};
        std::uint32_t state_offset = 0;
        VkViewport viewport{};
        // Offscreen render target of the draw, or null for the back buffer.
        RenderTarget* target = nullptr;
        // IDirect3DDevice8::Clear recorded in draw order; geometry is null for these.
        std::uint32_t clear_flags = 0;
        VkClearValue clear_color{};
        VkClearValue clear_depth{};
        VkRect2D clear_rect{};
    };
    struct UiBatch {
        std::uint32_t first_index = 0;
        std::uint32_t index_count = 0;
        VkPipeline pipeline = VK_NULL_HANDLE;
        VkDescriptorSet descriptor = VK_NULL_HANDLE;
        PushConstants constants{};
        bool behind_3d = false;
        std::uint32_t state_offset = 0;
    };

    static constexpr std::size_t kMaxBonesPerFrame = 65536;
    static constexpr VkDeviceSize kBoneBufferBytes = kMaxBonesPerFrame * 16 * sizeof(float);
    static constexpr std::size_t kMaxUiVerticesPerFrame = 65536;
    static constexpr std::size_t kMaxUiIndicesPerFrame = 98304;
    static constexpr VkDeviceSize kUiVertexBytes = kMaxUiVerticesPerFrame * sizeof(Vertex);
    static constexpr VkDeviceSize kUiBufferBytes = kUiVertexBytes + kMaxUiIndicesPerFrame * sizeof(std::uint32_t);
    static constexpr VkDeviceSize kStagingRingBytes = 4ull * 1024 * 1024;
    static constexpr VkDeviceSize kStagingRingMaxBytes = 64ull * 1024 * 1024;

public:
    struct Timings {
        double sync_ms = 0;
        double fence_ms = 0;  // inside sync_ms: waiting for the frame slot's previous GPU work
        double prepare_ms = 0;
        double steady_prepare_ms = 0;
        double submit_ms = 0;
        double present_ms = 0;
        double gpu_ms = 0;
        std::uint64_t gpu_samples = 0;
        // Inside prepare_ms: creating geometry buffers and the blocking texture uploads.
        double geometry_upload_ms = 0;
        double texture_upload_ms = 0;
    };

#ifdef MT_NATIVE_HAS_LIVE_CLIENT
    void enable_game_hardware_cursor();

    void sync_game_cursor();
#endif

    explicit VulkanWindow(bool vsync = true, int init_width = 960, int init_height = 640);

    ~VulkanWindow();

    // 32-bit bottom-up BMP from a B8G8R8A8 / R8G8B8A8 swapchain readback.
    void write_bmp(const std::string& path, const std::uint8_t* pixels) const;

    // Seconds since the last touch, mouse or key event; 0 while a finger or mouse button is held (a held
    // joystick or attack button sends no events). The frame-rate policy's idle test.
    double input_idle_seconds() const;

    void request_screenshot(std::string path, std::uint64_t frame);
    std::uint64_t frame_number() const { return frame_number_; }

    int msaa_samples() const { return int(samples_); }
    std::uint32_t width() const { return extent_.width; }
    std::uint32_t height() const { return extent_.height; }
    std::uint32_t logical_width() const;
    std::uint32_t logical_height() const;

    VkViewport full_viewport() const;

    // Window pixels kept clear at the left and right screen edges (rounded corners, cutouts);
    // MainActivity measures them and passes --safe-inset-px.
    void set_safe_inset_px(int px) { safe_inset_px_ = std::max(0, px); }
    // Test builds: frames per second in the top-left corner (--show-fps).
    void set_show_fps(bool show) { show_fps_ = show; }
    SDL_Window* sdl_window() const { return window_; }
    // The last render()'s blocking time: slot fence + swapchain acquire.
    double last_sync_ms() const { return last_sync_ms_; }
    // The same inset in logical UI pixels: the 40250 window layers are shifted and narrowed by it.
    int ui_safe_inset() const;

    // D3DVIEWPORT8 (in the game's logical screen pixels) scaled to the swapchain extent.
    VkViewport draw_viewport(const Render3DDraw& draw) const;

    TouchController touch_controller;

    // Touch hosts: SDL text input (and with it the system keyboard) is on only while the player
    // types into a tapped EditLine; desktop keeps the always-on text input from window creation.
    void sync_screen_keyboard();

    bool poll(bool forward_to_live_client = false);

    void reset_timings();

    void finish_gpu_timings();

    // 1 = the old serial loop (CPU waits for the previous frame's GPU work), 2 = CPU/GPU overlap.
    void set_frames_in_flight(std::uint32_t count);
    std::uint32_t frames_in_flight() const { return frames_in_flight_; }

    void render(
        const std::vector<Render3DDraw>& draws,
        const std::unordered_map<std::string, std::vector<std::uint8_t>>& capture_textures,
        std::uint32_t ui_width = 960,
        std::uint32_t ui_height = 640,
        const std::vector<UIRenderCommand>& ui_commands = {});

    std::size_t draw_count() const { return last_draw_count_; }
    std::size_t skinned_draw_count() const { return last_skinned_draw_count_; }
    std::size_t ui_batch_count() const { return last_ui_batch_count_; }
    std::size_t ui_quad_count() const { return last_ui_quad_count_; }
    std::size_t vertex_count() const { return last_vertex_count_; }
    std::size_t index_count() const { return last_index_count_; }
    std::size_t upload_count() const { return upload_count_; }
    std::size_t uploaded_bytes() const { return uploaded_bytes_; }
    std::size_t texture_upload_count() const { return texture_upload_count_; }
    std::size_t texture_uploaded_bytes() const { return texture_uploaded_bytes_; }
    const std::string& device_name() const { return device_name_; }
    const char* present_mode_name() const;
    Timings timings() const { return timings_; }
    native_perf::RendererTotals perf_totals() const;

private:
    void recreate_swapchain();

    void collect_pending_gpu_timestamp(FrameSlot& slot);

    void build_ui_batches(
        std::uint32_t ui_width,
        std::uint32_t ui_height,
        const std::vector<UIRenderCommand>& commands,
        const std::unordered_map<std::string, std::vector<std::uint8_t>>& capture_textures,
        std::vector<UiBatch>& batches,
        std::size_t& out_quad_count);

    std::uint32_t find_memory_type(std::uint32_t type_bits, VkMemoryPropertyFlags flags) const;

    void select_device();

    void create_swapchain();

    // The multisampled colour target the subpass resolves into the swapchain image. It never
    // leaves the render pass, so it is transient (lazily allocated tile memory where available).
    void create_msaa_color();

    void destroy_msaa_color();

    void create_framebuffers();

    void create_host_buffer(VkDeviceSize size, VkBufferUsageFlags usage, VkBuffer& buffer, VkDeviceMemory& memory, void** mapped);

    void create_descriptors_and_buffers();

    // The slot's fence has signalled: its oversize staging buffers are free, and the ring restarts
    // (grown to the largest frame's uploads seen so far, so texture streaming stays in the ring).
    void release_retired_buffers(FrameSlot& slot);

    void reset_staging(FrameSlot& slot);

    void ensure_bone_capacity(std::size_t required);

    void ensure_state_capacity(std::size_t required);

    void ensure_ui_capacity(std::size_t command_count);

    void create_render_pass_and_layout();

    // The render pass of offscreen render-target textures: colour (R5G6B5 like the D3D surface when
    // the device can render to it) + D32 depth, both loaded and stored, colour left shader-readable.
    void create_offscreen_pass();

    // "rt:<id>:<w>x<h>" -> its render target, created on first use (by a draw into it or by a draw
    // sampling it); null for a malformed name.
    RenderTarget* get_render_target(const std::string& name);

    void release_render_target(RenderTarget& target);

    // A new render target starts white with depth 1 (what the CPU surface's first Clear leaves), in
    // the layouts offscreen_pass_ expects. Recorded outside any render pass.
    void initialize_render_target(RenderTarget& target);

    // offscreen: for offscreen_pass_ (single-sampled render-target texture) instead of pass_.
    VkPipeline get_pipeline(std::uint8_t cull, std::uint8_t depth, std::uint8_t blend,
                            bool lines = false, std::uint8_t z_func = 4, bool offscreen = false);

    VkDescriptorSet allocate_texture_descriptor_set(VkImageView view0, VkImageView view1,
                                                     VkSampler sampler0 = VK_NULL_HANDLE,
                                                     VkSampler sampler1 = VK_NULL_HANDLE);

    VkDescriptorSet allocate_ui_texture_descriptor_set(VkImageView view0, VkImageView view1);

    GpuTexture& get_gpu_texture(
        const std::string& texture_name,
        const std::unordered_map<std::string, std::vector<std::uint8_t>>& capture_textures);

    std::uint64_t sampler_state_key(const Render3DDraw& draw, int stage) const;

    VkSampler sampler_for_draw(const Render3DDraw& draw, int stage);

    VkDescriptorSet get_texture_descriptor(
        const std::string& texture0_name,
        const std::string& mask_name,
        const std::unordered_map<std::string, std::vector<std::uint8_t>>& capture_textures,
        const Render3DDraw& draw);

    VkDescriptorSet get_ui_texture_descriptor(
        const std::string& texture0_name,
        const std::string& mask_name,
        const std::unordered_map<std::string, std::vector<std::uint8_t>>& capture_textures);

    void upload_texture(
        std::uint32_t width,
        std::uint32_t height,
        const std::uint8_t* rgba,
        GpuTexture& texture,
        bool count_upload,
        bool generate_mips = true,
        const std::vector<std::vector<std::uint8_t>>* file_mips = nullptr);

    void release_texture(GpuTexture& texture);

    void prune_textures();

    void release_geometry(Geometry& geometry);

    void upload_geometry(const Render3DDraw& draw, Geometry& geometry);

    bool vsync_ = true;
    VkPresentModeKHR present_mode_ = VK_PRESENT_MODE_FIFO_KHR;
    float timestamp_period_ns_ = 1.0f;
    float max_anisotropy_ = 1.0f;
    SDL_Window* window_ = nullptr;
    unsigned screen_keyboard_serial_ = 0;
    int safe_inset_px_ = 0;
    bool show_fps_ = false;
    int fps_ = -1;
    int fps_frames_ = 0;
    std::chrono::steady_clock::time_point fps_window_start_{};

    // Adds the scope's wall time to a Timings field.
    struct ScopedMs {
        explicit ScopedMs(double& into) : into_(into), start_(std::chrono::steady_clock::now()) {}
        ~ScopedMs() { into_ += std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start_).count(); }
        double& into_;
        std::chrono::steady_clock::time_point start_;
    };

    // Presented frames over the last half second.
    void update_fps_counter();
    VkInstance instance_ = VK_NULL_HANDLE;
    VkSurfaceKHR surface_ = VK_NULL_HANDLE;
    VkPhysicalDevice physical_ = VK_NULL_HANDLE;
    VkDevice device_ = VK_NULL_HANDLE;
    VkQueue queue_ = VK_NULL_HANDLE;
    std::uint32_t queue_family_ = 0;
    std::string device_name_;
    VkSwapchainKHR swapchain_ = VK_NULL_HANDLE;
    VkFormat swapchain_format_ = VK_FORMAT_UNDEFINED;
    VkExtent2D extent_{};
    std::vector<VkImageView> image_views_;
    std::vector<VkImage> swapchain_images_;
    void note_input(const SDL_Event& event);
    std::chrono::steady_clock::time_point last_input_ = std::chrono::steady_clock::now();
    int fingers_down_ = 0, buttons_down_ = 0;

    // --screenshot-out: read back the swapchain image of one frame into a BMP file.
    std::string screenshot_path_;
    std::uint64_t screenshot_frame_ = 0;
    VkImage depth_image_ = VK_NULL_HANDLE;
    VkDeviceMemory depth_memory_ = VK_NULL_HANDLE;
    VkImageView depth_view_ = VK_NULL_HANDLE;
    VkSampleCountFlagBits samples_ = VK_SAMPLE_COUNT_1_BIT;
    VkImage msaa_image_ = VK_NULL_HANDLE;
    VkDeviceMemory msaa_memory_ = VK_NULL_HANDLE;
    VkImageView msaa_view_ = VK_NULL_HANDLE;
    VkRenderPass pass_ = VK_NULL_HANDLE;
    std::vector<VkFramebuffer> framebuffers_;
    VkSampler sampler_ = VK_NULL_HANDLE;
    VkSampler clamp_sampler_ = VK_NULL_HANDLE;
    VkSampler ui_sampler_ = VK_NULL_HANDLE;
    std::unordered_map<std::uint64_t, VkSampler> state_samplers_;
    VkDescriptorSetLayout descriptor_layout_ = VK_NULL_HANDLE;
    VkDescriptorSetLayout bone_descriptor_layout_ = VK_NULL_HANDLE;
    VkDescriptorPool descriptor_pool_ = VK_NULL_HANDLE;
    // Everything one frame writes while the GPU may still read the previous frame's copy: with
    // kFramesInFlight slots the CPU prepares frame N while the GPU draws frame N-1. A slot is reused
    // only after its fence (the submit of frame N - frames_in_flight_) has signalled.
    struct FrameSlot {
        VkCommandBuffer command = VK_NULL_HANDLE;
        VkFence fence = VK_NULL_HANDLE;
        VkSemaphore acquire = VK_NULL_HANDLE;
        std::uint32_t query_base = 0;
        bool has_pending_query = false;
        bool recording = false;
        VkDescriptorSet bone_descriptor_set = VK_NULL_HANDLE;
        VkBuffer bone_buffer = VK_NULL_HANDLE;
        VkDeviceMemory bone_memory = VK_NULL_HANDLE;
        float* bone_mapped = nullptr;
        std::size_t bone_capacity = 0;
        VkBuffer state_buffer = VK_NULL_HANDLE;
        VkDeviceMemory state_memory = VK_NULL_HANDLE;
        std::uint8_t* state_mapped = nullptr;
        std::size_t state_capacity = 0;
        VkBuffer ui_buffer = VK_NULL_HANDLE;
        VkDeviceMemory ui_memory = VK_NULL_HANDLE;
        std::uint8_t* ui_mapped = nullptr;
        std::size_t ui_vertex_capacity = 0, ui_index_capacity = 0;
        VkDeviceSize ui_vertex_bytes = 0;
        // Texture uploads recorded into this frame's command buffer copy from here.
        VkBuffer staging_buffer = VK_NULL_HANDLE;
        VkDeviceMemory staging_memory = VK_NULL_HANDLE;
        std::uint8_t* staging_mapped = nullptr;
        VkDeviceSize staging_capacity = 0, staging_used = 0, staging_wanted = 0;
        // Uploads larger than the ring get their own staging buffer, freed when the slot comes back.
        std::vector<std::pair<VkBuffer, VkDeviceMemory>> retired_buffers;
    };
    static constexpr std::uint32_t kMaxFramesInFlight = 2;
    std::array<FrameSlot, kMaxFramesInFlight> frames_{};
    FrameSlot* f_ = &frames_[0];
    std::uint32_t frames_in_flight_ = kMaxFramesInFlight;
    std::uint32_t frame_slot_ = 0;
    // Signalled by a frame's submit, waited by the present of that swapchain image.
    std::vector<VkSemaphore> rendered_;
    VkDeviceSize state_stride_ = 0;
    VkDeviceSize staging_alignment_ = 16;
    GpuTexture fallback_texture_{};
    std::unordered_map<std::string, GpuTexture> textures_;
    std::unordered_map<std::string, PairedDescriptor> paired_descriptors_;
    std::unordered_map<std::string, RenderTarget> render_targets_;
    VkRenderPass offscreen_pass_ = VK_NULL_HANDLE;
    VkFormat offscreen_format_ = VK_FORMAT_R5G6B5_UNORM_PACK16;
    std::size_t last_offscreen_draw_count_ = 0;
    VkShaderModule vertex_module_ = VK_NULL_HANDLE;
    VkShaderModule fragment_module_ = VK_NULL_HANDLE;
    VkPipelineLayout layout_ = VK_NULL_HANDLE;
    std::unordered_map<std::uint32_t, VkPipeline> pipelines_;
    std::unordered_map<GeometryId, Geometry, GeometryIdHash> geometries_;
    std::uint64_t frame_number_ = 0;
    std::uint64_t timed_frames_ = 0;
    std::size_t upload_count_ = 0, uploaded_bytes_ = 0;
    std::size_t texture_upload_count_ = 0, texture_uploaded_bytes_ = 0;
    std::string last_texture_upload_name_;
    VkCommandPool pool_ = VK_NULL_HANDLE;
    VkQueryPool query_pool_ = VK_NULL_HANDLE;
    bool os_cursor_hidden_ = false;
    bool hardware_cursor_enabled_ = false;
    std::array<SDL_Cursor*, 15> game_cursors_{};
    SDL_Cursor* fallback_cursor_ = nullptr;
    int current_cursor_shape_ = -1;
    int last_mx_ = 0, last_my_ = 0;
    std::size_t last_draw_count_ = 0, last_skinned_draw_count_ = 0;
    std::size_t last_ui_batch_count_ = 0, last_ui_quad_count_ = 0;
    std::size_t last_vertex_count_ = 0, last_index_count_ = 0;
    Timings timings_{};
    double last_sync_ms_ = 0;
};

// One summary line of frame and renderer counters on stdout.
void print_summary(const VulkanWindow& renderer, int completed, double wall_ms, double update_ms = 0.0,
                   std::vector<double> frame_ms = {});

} // namespace mt_host