#include "renderer_api/recording_renderer.h"

#include <new>
#include <utility>

namespace pp::renderer {

namespace {

[[nodiscard]] const char* non_null_name(const char* name) noexcept
{
    return name == nullptr ? "" : name;
}

void push_command(
    std::vector<RecordedRenderCommand>* commands,
    RecordedRenderCommand command)
{
    if (commands != nullptr) {
        commands->push_back(command);
    }
}

template <typename Resource, typename Interface, typename... Args>
[[nodiscard]] pp::foundation::Result<std::unique_ptr<Interface>> make_recording_resource(
    Args&&... args) noexcept
{
    auto resource = std::unique_ptr<Resource>(new (std::nothrow) Resource(std::forward<Args>(args)...));
    if (!resource) {
        return pp::foundation::Result<std::unique_ptr<Interface>>::failure(
            pp::foundation::Status::out_of_range("renderer resource allocation failed"));
    }

    std::unique_ptr<Interface> erased = std::move(resource);
    return pp::foundation::Result<std::unique_ptr<Interface>>::success(std::move(erased));
}

}

RecordingTexture2D::RecordingTexture2D(TextureDesc desc) noexcept
    : desc_(desc)
{
}

TextureDesc RecordingTexture2D::desc() const noexcept
{
    return desc_;
}

RecordingRenderTarget::RecordingRenderTarget(TextureDesc color_desc) noexcept
    : color_desc_(color_desc)
{
}

TextureDesc RecordingRenderTarget::color_desc() const noexcept
{
    return color_desc_;
}

RecordingShaderProgram::RecordingShaderProgram(const char* debug_name) noexcept
    : debug_name_(non_null_name(debug_name))
{
}

const char* RecordingShaderProgram::debug_name() const noexcept
{
    return debug_name_;
}

RecordingMesh::RecordingMesh(MeshDesc desc) noexcept
    : desc_(desc)
{
}

MeshDesc RecordingMesh::desc() const noexcept
{
    return desc_;
}

RecordingReadbackBuffer::RecordingReadbackBuffer(std::uint64_t size_bytes) noexcept
    : size_bytes_(size_bytes)
{
}

std::uint64_t RecordingReadbackBuffer::size_bytes() const noexcept
{
    return size_bytes_;
}

RecordingCommandContext::RecordingCommandContext(std::vector<RecordedRenderCommand>& commands) noexcept
    : commands_(&commands)
{
}

pp::foundation::Status RecordingCommandContext::begin_render_pass(
    IRenderTarget& target,
    RenderPassDesc desc) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass is already active");
    }

    active_target_ = target.color_desc();
    if (!has_texture_usage(active_target_.usage, TextureUsage::render_target)) {
        return pp::foundation::Status::invalid_argument("render target texture must allow render_target usage");
    }

    const auto size_status = texture_byte_size(active_target_);
    if (!size_status.ok()) {
        return size_status.status();
    }

    const auto render_pass_status = validate_render_pass_desc(desc);
    if (!render_pass_status.ok()) {
        return render_pass_status;
    }

    in_render_pass_ = true;
    shader_bound_ = false;
    mesh_bound_ = false;
    active_mesh_ = MeshDesc {};
    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::begin_render_pass,
        .target_desc = active_target_,
        .clear_color_enabled = desc.clear_color_enabled,
        .clear_color = desc.clear_color,
        .clear_depth_enabled = desc.clear_depth_enabled,
        .clear_depth = desc.clear_depth,
        .clear_stencil_enabled = desc.clear_stencil_enabled,
        .clear_stencil = desc.clear_stencil,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::set_viewport(Viewport viewport) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    const auto status = validate_viewport(viewport, active_target_.extent);
    if (!status.ok()) {
        return status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::set_viewport,
        .viewport = viewport,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::set_scissor(ScissorRect scissor) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    const auto status = validate_scissor(scissor, active_target_.extent);
    if (!status.ok()) {
        return status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::set_scissor,
        .scissor = scissor,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::set_blend_state(BlendState state) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    const auto status = validate_blend_state(state);
    if (!status.ok()) {
        return status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::set_blend_state,
        .blend_state = state,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::set_depth_state(DepthState state) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    const auto status = validate_depth_state(state);
    if (!status.ok()) {
        return status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::set_depth_state,
        .depth_state = state,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::bind_shader(IShaderProgram& shader) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    shader_bound_ = true;
    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::bind_shader,
        .name = shader.debug_name(),
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::set_shader_uniform(
    const char* name,
    std::span<const std::byte> bytes) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }
    if (!shader_bound_) {
        return pp::foundation::Status::invalid_argument("shader must be bound before setting uniforms");
    }

    const auto status = validate_shader_uniform_write(name, bytes);
    if (!status.ok()) {
        return status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::set_shader_uniform,
        .uniform_bytes = static_cast<std::uint64_t>(bytes.size()),
        .name = name,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::bind_texture(
    std::uint32_t slot,
    ITexture2D& texture) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    const auto slot_status = validate_texture_slot(slot);
    if (!slot_status.ok()) {
        return slot_status;
    }

    const auto desc = texture.desc();
    if (!has_texture_usage(desc.usage, TextureUsage::sampled)) {
        return pp::foundation::Status::invalid_argument("bound texture must allow sampled usage");
    }

    const auto size_status = texture_byte_size(desc);
    if (!size_status.ok()) {
        return size_status.status();
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::bind_texture,
        .texture_desc = desc,
        .texture_slot = slot,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::bind_sampler(
    std::uint32_t slot,
    SamplerDesc sampler) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    const auto slot_status = validate_texture_slot(slot);
    if (!slot_status.ok()) {
        return slot_status;
    }

    const auto sampler_status = validate_sampler_desc(sampler);
    if (!sampler_status.ok()) {
        return sampler_status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::bind_sampler,
        .sampler_desc = sampler,
        .sampler_slot = slot,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::bind_mesh(IMesh& mesh) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }

    const auto desc = mesh.desc();
    const auto status = validate_mesh_desc(desc);
    if (!status.ok()) {
        return status;
    }

    mesh_bound_ = true;
    active_mesh_ = desc;
    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::bind_mesh,
        .mesh_desc = desc,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::draw(DrawDesc desc) noexcept
{
    if (!in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render pass has not begun");
    }
    if (!shader_bound_) {
        return pp::foundation::Status::invalid_argument("shader must be bound before draw");
    }
    if (!mesh_bound_) {
        return pp::foundation::Status::invalid_argument("mesh must be bound before draw");
    }

    const auto draw_status = validate_draw_desc(active_mesh_, desc);
    if (!draw_status.ok()) {
        return draw_status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::draw,
        .mesh_desc = active_mesh_,
        .draw_desc = desc,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::read_texture(
    ITexture2D& texture,
    ReadbackRegion region,
    IReadbackBuffer& destination) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("readback must be outside a render pass");
    }

    const auto desc = texture.desc();
    if (!has_texture_usage(desc.usage, TextureUsage::readback_source)) {
        return pp::foundation::Status::invalid_argument("readback texture must allow readback_source usage");
    }

    const auto bytes = readback_byte_size(desc, region);
    if (!bytes) {
        return bytes.status();
    }

    if (destination.size_bytes() < bytes.value()) {
        return pp::foundation::Status::out_of_range("readback buffer is too small");
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::read_texture,
        .texture_desc = desc,
        .readback_region = region,
        .readback_bytes = bytes.value(),
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::upload_texture(
    ITexture2D& texture,
    ReadbackRegion region,
    std::span<const std::byte> rgba_or_channel_bytes) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("texture upload must be outside a render pass");
    }

    const auto desc = texture.desc();
    if (!has_texture_usage(desc.usage, TextureUsage::upload_destination)) {
        return pp::foundation::Status::invalid_argument("texture upload destination must allow upload_destination usage");
    }

    const auto bytes = readback_byte_size(desc, region);
    if (!bytes) {
        return bytes.status();
    }

    if (rgba_or_channel_bytes.size() != bytes.value()) {
        return pp::foundation::Status::invalid_argument("texture upload byte size does not match the region");
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::upload_texture,
        .texture_desc = desc,
        .readback_region = region,
        .upload_bytes = bytes.value(),
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::generate_mipmaps(ITexture2D& texture) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("mipmap generation must be outside a render pass");
    }

    const auto desc = texture.desc();
    const auto desc_status = validate_mipmap_generation_desc(desc);
    if (!desc_status.ok()) {
        return desc_status;
    }

    const auto bytes = texture_byte_size(desc);
    if (!bytes.ok()) {
        return bytes.status();
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::generate_mipmaps,
        .texture_desc = desc,
        .generated_mip_levels = desc.mip_levels,
        .generated_mip_bytes = bytes.value(),
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::transition_texture(
    ITexture2D& texture,
    TextureState before,
    TextureState after) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("texture transition must be outside a render pass");
    }

    const auto desc = texture.desc();
    const auto desc_status = validate_texture_transition_desc(desc, before, after);
    if (!desc_status.ok()) {
        return desc_status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::transition_texture,
        .texture_desc = desc,
        .before_state = before,
        .after_state = after,
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::copy_texture(
    ITexture2D& source,
    ReadbackRegion source_region,
    ITexture2D& destination,
    ReadbackRegion destination_region) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("texture copy must be outside a render pass");
    }

    const auto source_desc = source.desc();
    const auto destination_desc = destination.desc();
    const auto desc_status = validate_texture_copy_descs(
        source_desc,
        source_region,
        destination_desc,
        destination_region);
    if (!desc_status.ok()) {
        return desc_status;
    }

    const auto source_bytes = readback_byte_size(source_desc, source_region);
    if (!source_bytes.ok()) {
        return source_bytes.status();
    }

    const auto destination_bytes = readback_byte_size(destination_desc, destination_region);
    if (!destination_bytes.ok()) {
        return destination_bytes.status();
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::copy_texture,
        .source_desc = source_desc,
        .destination_desc = destination_desc,
        .source_region = source_region,
        .destination_region = destination_region,
        .copy_source_bytes = source_bytes.value(),
        .copy_destination_bytes = destination_bytes.value(),
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::capture_frame(
    IRenderTarget& target,
    IReadbackBuffer& destination) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("frame capture must be outside a render pass");
    }

    const auto desc = target.color_desc();
    const auto bytes = frame_capture_byte_size(desc);
    if (!bytes) {
        return bytes.status();
    }

    if (destination.size_bytes() < bytes.value()) {
        return pp::foundation::Status::out_of_range("frame capture buffer is too small");
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::capture_frame,
        .target_desc = desc,
        .capture_bytes = bytes.value(),
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingCommandContext::blit_render_target(
    IRenderTarget& source,
    ReadbackRegion source_region,
    IRenderTarget& destination,
    ReadbackRegion destination_region,
    BlitFilter filter) noexcept
{
    if (in_render_pass_) {
        return pp::foundation::Status::invalid_argument("render target blit must be outside a render pass");
    }

    const auto source_desc = source.color_desc();
    const auto destination_desc = destination.color_desc();
    const auto desc_status = validate_blit_descs(source_desc, destination_desc);
    if (!desc_status.ok()) {
        return desc_status;
    }

    const auto filter_status = validate_blit_filter(filter);
    if (!filter_status.ok()) {
        return filter_status;
    }

    const auto source_bytes = readback_byte_size(source_desc, source_region);
    if (!source_bytes) {
        return source_bytes.status();
    }

    const auto destination_bytes = readback_byte_size(destination_desc, destination_region);
    if (!destination_bytes) {
        return destination_bytes.status();
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::blit_render_target,
        .source_desc = source_desc,
        .destination_desc = destination_desc,
        .source_region = source_region,
        .destination_region = destination_region,
        .blit_filter = filter,
        .blit_source_bytes = source_bytes.value(),
        .blit_destination_bytes = destination_bytes.value(),
    });
    return pp::foundation::Status::success();
}

void RecordingCommandContext::end_render_pass() noexcept
{
    if (!in_render_pass_) {
        return;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::end_render_pass,
    });
    in_render_pass_ = false;
    shader_bound_ = false;
    mesh_bound_ = false;
    active_mesh_ = MeshDesc {};
}

bool RecordingCommandContext::in_render_pass() const noexcept
{
    return in_render_pass_;
}

void RecordingCommandContext::reset() noexcept
{
    active_target_ = TextureDesc {};
    active_mesh_ = MeshDesc {};
    in_render_pass_ = false;
    shader_bound_ = false;
    mesh_bound_ = false;
}

RecordingRenderTrace::RecordingRenderTrace(std::vector<RecordedRenderCommand>& commands) noexcept
    : commands_(&commands)
{
}

pp::foundation::Status RecordingRenderTrace::marker(const char* component, const char* name) noexcept
{
    const auto status = validate_trace_label(component, name);
    if (!status.ok()) {
        return status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::trace_marker,
        .component = non_null_name(component),
        .name = non_null_name(name),
    });
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingRenderTrace::begin_scope(const char* component, const char* name) noexcept
{
    const auto status = validate_trace_label(component, name);
    if (!status.ok()) {
        return status;
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::trace_begin_scope,
        .component = non_null_name(component),
        .name = non_null_name(name),
    });
    ++scope_depth_;
    return pp::foundation::Status::success();
}

pp::foundation::Status RecordingRenderTrace::end_scope() noexcept
{
    if (scope_depth_ == 0U) {
        return pp::foundation::Status::invalid_argument("trace scope has not begun");
    }

    push_command(commands_, RecordedRenderCommand {
        .kind = RecordedRenderCommandKind::trace_end_scope,
    });
    --scope_depth_;
    return pp::foundation::Status::success();
}

void RecordingRenderTrace::reset() noexcept
{
    scope_depth_ = 0U;
}

RecordingRenderDevice::RecordingRenderDevice() noexcept
    : context_(commands_)
    , trace_(commands_)
{
}

const char* RecordingRenderDevice::backend_name() const noexcept
{
    return "recording";
}

RenderDeviceFeatures RecordingRenderDevice::features() const noexcept
{
    return RenderDeviceFeatures {
        .explicit_texture_transitions = true,
        .texture_copy = true,
        .render_target_blit = true,
        .frame_capture = true,
    };
}

pp::foundation::Result<std::unique_ptr<ITexture2D>> RecordingRenderDevice::create_texture(
    TextureDesc desc) noexcept
{
    const auto desc_status = validate_texture_desc(desc);
    if (!desc_status.ok()) {
        return pp::foundation::Result<std::unique_ptr<ITexture2D>>::failure(desc_status);
    }

    const auto bytes = texture_byte_size(desc);
    if (!bytes.ok()) {
        return pp::foundation::Result<std::unique_ptr<ITexture2D>>::failure(bytes.status());
    }

    return make_recording_resource<RecordingTexture2D, ITexture2D>(desc);
}

pp::foundation::Result<std::unique_ptr<IRenderTarget>> RecordingRenderDevice::create_render_target(
    TextureDesc color_desc) noexcept
{
    if (!has_texture_usage(color_desc.usage, TextureUsage::render_target)) {
        return pp::foundation::Result<std::unique_ptr<IRenderTarget>>::failure(
            pp::foundation::Status::invalid_argument("render target texture must allow render_target usage"));
    }

    const auto desc_status = validate_texture_desc(color_desc);
    if (!desc_status.ok()) {
        return pp::foundation::Result<std::unique_ptr<IRenderTarget>>::failure(desc_status);
    }

    const auto bytes = texture_byte_size(color_desc);
    if (!bytes.ok()) {
        return pp::foundation::Result<std::unique_ptr<IRenderTarget>>::failure(bytes.status());
    }

    return make_recording_resource<RecordingRenderTarget, IRenderTarget>(color_desc);
}

pp::foundation::Result<std::unique_ptr<IShaderProgram>> RecordingRenderDevice::create_shader_program(
    ShaderProgramDesc desc) noexcept
{
    const auto status = validate_shader_program_desc(desc);
    if (!status.ok()) {
        return pp::foundation::Result<std::unique_ptr<IShaderProgram>>::failure(status);
    }

    return make_recording_resource<RecordingShaderProgram, IShaderProgram>(desc.debug_name);
}

pp::foundation::Result<std::unique_ptr<IMesh>> RecordingRenderDevice::create_mesh(
    MeshDesc desc) noexcept
{
    const auto status = validate_mesh_desc(desc);
    if (!status.ok()) {
        return pp::foundation::Result<std::unique_ptr<IMesh>>::failure(status);
    }

    return make_recording_resource<RecordingMesh, IMesh>(desc);
}

pp::foundation::Result<std::unique_ptr<IReadbackBuffer>> RecordingRenderDevice::create_readback_buffer(
    std::uint64_t size_bytes) noexcept
{
    if (size_bytes == 0U) {
        return pp::foundation::Result<std::unique_ptr<IReadbackBuffer>>::failure(
            pp::foundation::Status::invalid_argument("readback buffer size must be greater than zero"));
    }

    if (size_bytes > max_texture_bytes) {
        return pp::foundation::Result<std::unique_ptr<IReadbackBuffer>>::failure(
            pp::foundation::Status::out_of_range("readback buffer size exceeds the configured limit"));
    }

    return make_recording_resource<RecordingReadbackBuffer, IReadbackBuffer>(size_bytes);
}

ICommandContext& RecordingRenderDevice::immediate_context() noexcept
{
    return context_;
}

IRenderTrace* RecordingRenderDevice::trace() noexcept
{
    return &trace_;
}

std::span<const RecordedRenderCommand> RecordingRenderDevice::commands() const noexcept
{
    return commands_;
}

void RecordingRenderDevice::clear() noexcept
{
    commands_.clear();
    context_.reset();
    trace_.reset();
}

const char* recorded_render_command_kind_name(RecordedRenderCommandKind kind) noexcept
{
    switch (kind) {
    case RecordedRenderCommandKind::begin_render_pass:
        return "begin_render_pass";
    case RecordedRenderCommandKind::set_viewport:
        return "set_viewport";
    case RecordedRenderCommandKind::set_scissor:
        return "set_scissor";
    case RecordedRenderCommandKind::set_blend_state:
        return "set_blend_state";
    case RecordedRenderCommandKind::set_depth_state:
        return "set_depth_state";
    case RecordedRenderCommandKind::bind_shader:
        return "bind_shader";
    case RecordedRenderCommandKind::set_shader_uniform:
        return "set_shader_uniform";
    case RecordedRenderCommandKind::bind_texture:
        return "bind_texture";
    case RecordedRenderCommandKind::bind_sampler:
        return "bind_sampler";
    case RecordedRenderCommandKind::bind_mesh:
        return "bind_mesh";
    case RecordedRenderCommandKind::draw:
        return "draw";
    case RecordedRenderCommandKind::upload_texture:
        return "upload_texture";
    case RecordedRenderCommandKind::generate_mipmaps:
        return "generate_mipmaps";
    case RecordedRenderCommandKind::transition_texture:
        return "transition_texture";
    case RecordedRenderCommandKind::copy_texture:
        return "copy_texture";
    case RecordedRenderCommandKind::read_texture:
        return "read_texture";
    case RecordedRenderCommandKind::capture_frame:
        return "capture_frame";
    case RecordedRenderCommandKind::blit_render_target:
        return "blit_render_target";
    case RecordedRenderCommandKind::end_render_pass:
        return "end_render_pass";
    case RecordedRenderCommandKind::trace_marker:
        return "trace_marker";
    case RecordedRenderCommandKind::trace_begin_scope:
        return "trace_begin_scope";
    case RecordedRenderCommandKind::trace_end_scope:
        return "trace_end_scope";
    }

    return "unknown";
}

}