/* Determine if file was included from C or from HLSL */
#if !LanguageIsGpu
# define K_IS_CPU 1
#else
# define K_IS_CPU 0
#endif

#if K_IS_CPU

#define K_STRUCT(s) Packed(struct s)
#define K_DECL(t, n) struct Cat(K_, t) n
#define K_DECLS(t, n) K_DECL(t, n)
#define K_STATIC_ASSERT(c) StaticAssert(c)

typedef struct K_uint K_uint;
struct K_uint { u32 v; };
Inline struct K_uint K_UintFromU32(u32 v)
{
    return (struct K_uint) { .v = v };
}

typedef struct K_int K_int;
struct K_int { i32 v; };
Inline struct K_int K_IntFromI32(i32 v)
{
    return (struct K_int) { .v = v };
}

typedef struct K_uint2 K_uint2;
struct K_uint2 { u32 v[2]; };
Inline struct K_uint2 K_Uint2FromU32(u32 x, u32 y)
{
    return (struct K_uint2) { .v[0] = x, .v[1] = y };
}

typedef struct K_uint3 K_uint3;
struct K_uint3 { u32 v[3]; };
Inline struct K_uint3 K_Uint3FromU32(u32 x, u32 y, u32 z)
{
    return (struct K_uint3) { .v[0] = x, .v[1] = y, .v[2] = z };
}

typedef struct K_uint4 K_uint4;
struct K_uint4 { u32 v[4]; };
Inline struct K_uint4 K_Uint4FromU32(u32 x, u32 y, u32 z, u32 w)
{
    return (struct K_uint4) { .v[0] = x, .v[1] = y, .v[2] = z, .v[3] = w };
}

typedef struct K_float K_float;
struct K_float { f32 v; };
Inline struct K_float K_FloatFromF32(f32 v)
{
    return (struct K_float) { .v = v };
}

typedef struct K_float2 K_float2;
struct K_float2 { f32 v[2]; };
Inline struct K_float2 K_Float2FromV2(Vec2 v)
{
    return (struct K_float2) { .v[0] = v.x, .v[1] = v.y };
}

typedef struct K_float3 K_float3;
struct K_float3 { f32 v[3]; };
Inline struct K_float3 K_Float3FromV3(Vec3 v)
{
    return (struct K_float3) { .v[0] = v.x, .v[1] = v.y, .v[2] = v.z };
}

typedef struct K_float4x4 K_float4x4;
struct K_float4x4 { f32 v[4][4]; };
Inline struct K_float4x4 K_Float4x4FromMat4x4(Mat4x4 v)
{
    struct K_float4x4 result;
    StaticAssert(sizeof(result) == sizeof(v));
    CopyBytes(&result, v.e, sizeof(result));
    return result;
}

struct K_float2x3 { f32 v[2][3]; };
Inline struct K_float2x3 K_Float2x3FromXform(Xform v)
{
    struct K_float2x3 result;
    StaticAssert(sizeof(result) == sizeof(v));
    CopyBytes(&result, &v, sizeof(result));
    return result;
}

#else

#define K_STRUCT(s) struct s
#define K_DECL(t, n) t n
#define K_DECLS(t, n) t n : n
#define K_STATIC_ASSERT(c) _Static_assert(c, "")

#define DECLS(t, n) t n : n

#define resource_from_urid(urid) ResourceDescriptorHeap[urid]
#define resource_from_nurid(nurid) ResourceDescriptorHeap[NonUniformResourceIndex(nurid)]

float4 float4_from_uint_norm(uint v)
{
    float4 result;
    result.r = ((v >> 0) & 0xFF) / 255.0;
    result.g = ((v >> 8) & 0xFF) / 255.0;
    result.b = ((v >> 16) & 0xFF) / 255.0;
    result.a = ((v >> 24) & 0xFF) / 255.0;
    return result;
}

/* Linear color from normalized sRGB */
float4 linear_from_srgb(float4 srgb)
{
    return float4(pow(abs(srgb.rgb), 2.2), srgb.a);
}

/* Linear color from R8G8B8A8 sRGB  */
float4 linear_from_srgb32(uint srgb32)
{
    return linear_from_srgb(float4_from_uint_norm(srgb32));
}

/* ========================== *
 * Root signature
 * ========================== */

#define K_ROOTSIG \
    "RootFlags(CBV_SRV_UAV_HEAP_DIRECTLY_INDEXED | SAMPLER_HEAP_DIRECTLY_INDEXED | ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT), " \
    "RootConstants(b0, num32BitConstants = 64), " \
    \
    "StaticSampler(s0, " \
                  "filter = FILTER_MIN_MAG_MIP_POINT, " \
                  "addressU = TEXTURE_ADDRESS_CLAMP, " \
                  "addressV = TEXTURE_ADDRESS_CLAMP, " \
                  "addressW = TEXTURE_ADDRESS_CLAMP, " \
                  "maxAnisotropy = 1)"

SamplerState s_point_clamp : register(s0);

#define K_ENTRY [RootSignature(K_ROOTSIG)]


#endif

#define K_ASSERT_ROOT_CONST(s, n)  K_STATIC_ASSERT((sizeof(s) % 16 == 0) &&  /* Root constant struct should pad to 16 byte alignment */ \
                                                   ((sizeof(s) / 4) == n) &&  /* Root constant struct size should match the specified 32-bit-constant count */ \
                                                   (sizeof(s) <= 256))        /* Root constant struct can only fit 64 DWORDS */

/* ========================== *
 * Global textures
 * ========================== */

/* Blue noise */
#define K_BLUE_NOISE_TEX_ID        0
#define K_BLUE_NOISE_TEX_WIDTH     128
#define K_BLUE_NOISE_TEX_HEIGHT    128
#define K_BLUE_NOISE_TEX_DEPTH     64

/* ========================== *
 * Material shader structs
 * ========================== */

#if 0
K_STRUCT(k_material_sig {
    /* ----------------------------------------------------- */
    K_DECL(float4x4, projection);              /* 16 consts */
    /* ----------------------------------------------------- */
    K_DECL(uint, instances_urid);              /* 01 consts */
    K_DECL(uint, grids_urid);                  /* 01 consts */
    K_DECL(uint, _pad0);                       /* 01 consts (padding) */
    K_DECL(uint, _pad1);                       /* 01 consts (padding) */
    /* ----------------------------------------------------- */
});
K_ASSERT_ROOT_CONST(struct k_material_sig, 20);
#else
Struct(k_material_sig)
{
    /* ----------------------------------------------------- */
    Mat4x4 projection;                          /* 16 consts */
    /* ----------------------------------------------------- */
    u32 instances_urid;                         /* 01 consts */
    u32 grids_urid;                             /* 01 consts */
    u32 _pad0;                                  /* 01 consts (padding) */
    u32 _pad1;                                  /* 01 consts (padding) */
    /* ----------------------------------------------------- */
};
K_ASSERT_ROOT_CONST(struct k_material_sig, 20);
#endif

K_STRUCT(k_material_instance {
    K_DECL(uint, tex_nurid);
    K_DECL(uint, grid_id);
    K_DECL(float2x3, xf);
    K_DECL(float2, uv0);
    K_DECL(float2, uv1);
    K_DECL(uint, tint_srgb);
    K_DECL(uint, is_light);
    K_DECL(float3, light_emittance_srgb);
});

K_STRUCT(k_material_grid {
    K_DECL(float, line_thickness);
    K_DECL(float, line_spacing);
    K_DECL(float2, offset);
    K_DECL(uint, bg0_srgb);
    K_DECL(uint, bg1_srgb);
    K_DECL(uint, line_srgb);
    K_DECL(uint, x_srgb);
    K_DECL(uint, y_srgb);
});

/* ========================== *
 * Flood shader structs
 * ========================== */

K_STRUCT(k_flood_sig {
    /* ----------------------------------------------------- */
    K_DECL(int, step_len);                     /* 01 consts */
    K_DECL(uint, emittance_tex_urid);          /* 01 consts */
    K_DECL(uint, read_flood_tex_urid);         /* 01 consts */
    K_DECL(uint, target_flood_tex_urid);       /* 01 consts */
    /* ----------------------------------------------------- */
    K_DECL(uint, tex_width);                   /* 01 consts */
    K_DECL(uint, tex_height);                  /* 01 consts */
    K_DECL(uint, _pad0);                       /* 01 consts (padding) */
    K_DECL(uint, _pad1);                       /* 01 consts (padding) */
    /* ----------------------------------------------------- */
});
K_ASSERT_ROOT_CONST(struct k_flood_sig, 8);

/* ========================== *
 * Shade shader structs
 * ========================== */

#define K_SHADE_FLAG_NONE              (0 << 0)
#define K_SHADE_FLAG_DISABLE_EFFECTS   (1 << 0)

K_STRUCT(k_shade_sig {
    /* ----------------------------------------------------- */
    K_DECL(uint4, frame_seed);                 /* 04 consts */
    /* ----------------------------------------------------- */
    K_DECL(uint, flags);                       /* 01 consts */
    K_DECL(uint, _pad0);                       /* 01 consts (padding) */
    K_DECL(uint, tex_width);                   /* 01 consts */
    K_DECL(uint, tex_height);                  /* 01 consts */
    /* ----------------------------------------------------- */
    K_DECL(float2, camera_offset);             /* 02 consts */
    K_DECL(uint, frame_index);                 /* 01 consts */
    K_DECL(uint, albedo_tex_urid);             /* 01 consts */
    /* ----------------------------------------------------- */
    K_DECL(uint, emittance_tex_urid);          /* 01 consts */
    K_DECL(uint, emittance_flood_tex_urid);    /* 01 consts */
    K_DECL(uint, read_tex_urid);               /* 01 consts */
    K_DECL(uint, target_tex_urid);             /* 01 consts */
    /* ----------------------------------------------------- */
});
K_ASSERT_ROOT_CONST(struct k_shade_sig, 16);

/* ========================== *
 * Shape shader structs
 * ========================== */

K_STRUCT(k_shape_sig {
    /* ----------------------------------------------------- */
    K_DECL(float4x4, projection);              /* 16 consts */
    /* ----------------------------------------------------- */
    K_DECL(uint, verts_urid);                  /* 01 consts */
    K_DECL(uint, _pad0);                       /* 01 consts (padding) */
    K_DECL(uint, _pad1);                       /* 01 consts (padding) */
    K_DECL(uint, _pad2);                       /* 01 consts (padding) */
    /* ----------------------------------------------------- */
});
K_ASSERT_ROOT_CONST(struct k_shape_sig, 20);

K_STRUCT(k_shape_vert {
    K_DECL(float2, pos);
    K_DECL(uint, color_srgb);
});

/* ========================== *
 * UI shader structs
 * ========================== */

K_STRUCT(k_ui_sig {
    /* ----------------------------------------------------- */
    K_DECL(float4x4, projection);              /* 16 consts */
    /* ----------------------------------------------------- */
    K_DECL(uint, instances_urid);              /* 01 consts */
    K_DECL(uint, _pad0);                       /* 01 consts (padding) */
    K_DECL(uint, _pad1);                       /* 01 consts (padding) */
    K_DECL(uint, _pad2);                       /* 01 consts (padding) */
    /* ----------------------------------------------------- */
});
K_ASSERT_ROOT_CONST(struct k_ui_sig, 20);

K_STRUCT(k_ui_instance {
    K_DECL(uint, tex_nurid);
    K_DECL(uint, grid_id);
    K_DECL(float2x3, xf);
    K_DECL(float2, uv0);
    K_DECL(float2, uv1);
    K_DECL(uint, tint_srgb);
});

/* ========================== *
 * Blit shader structs
 * ========================== */

 #define K_BLIT_FLAG_NONE              (0 << 0)
 #define K_BLIT_FLAG_TONE_MAP          (1 << 0)
 #define K_BLIT_FLAG_GAMMA_CORRECT     (1 << 1)

K_STRUCT(k_blit_sig {
    /* ----------------------------------------------------- */
    K_DECL(float4x4, projection);              /* 16 consts */
    /* ----------------------------------------------------- */
    K_DECL(uint, flags);                       /* 01 consts */
    K_DECL(uint, tex_urid);                    /* 01 consts */
    K_DECL(float, exposure);                   /* 01 consts */
    K_DECL(float, gamma);                      /* 01 consts */
    /* ----------------------------------------------------- */
});
K_ASSERT_ROOT_CONST(struct k_blit_sig, 20);