#include "mirage_util.slang"

cbuffer ParamBuffer
{
    matrix transform;
};

struct PSInput {
    float4 position : SV_POSITION; 
    float2 uv : TEXCOORD0;
    float4 color : COLOR;
    float2 size : TEXCOORD1;     // 矩形大小(像素单位)
    float4 radius : TEXCOORD3;   // 四角圆角(像素单位): [左上,右上,左下,右下]
};

[shader("vertex")]
PSInput vertex_main(VSInput input)
{
    PSInput output;

    // 并通过常量缓冲区或结构化缓冲区传递
    output.position = mul(float4(input.position, 0.0, 1.0), transform);
    output.uv = input.uv;
    output.color = input.color;

    // 传递尺寸和半径数据
    output.size = input.param_a.xy;
    output.radius = input.param_b;

    return output;
}

// 改进的距离函数计算
float distance_from_rect_uv(float2 p, float2 size, float corner_radius) {
    // 保护除零
    corner_radius = max(corner_radius, 0.001);

    // 转换圆角半径为UV空间
    float2 corner_radius_uv = corner_radius / size;

    // 计算内部矩形边界（圆角内切矩形）
    float2 inner_rect = float2(1.0, 1.0) - corner_radius_uv;

    // 计算到内部矩形的距离
    float2 q = abs(p) - inner_rect;

    // 转换回像素空间进行距离计算
    q *= size;

    // 计算SDF（有符号距离场）
    return length(max(q, 0.0)) + min(max(q.x, q.y), 0.0) - corner_radius;
}

[shader("pixel")]
float4 pixel_main(PSInput input) : SV_Target
{
    // 将UV坐标转换为从中心点(-1,1)的归一化设备坐标
    float2 p = uv_to_ndc(input.uv);

    // 确定当前像素位于哪个象限，以选择正确的圆角半径
    // 0=左上(x<0,y<0), 1=右上(x>0,y<0), 2=左下(x<0,y>0), 3=右下(x>0,y>0)
    int idx = (p.x < 0 ? 0 : 1) + (p.y < 0 ? 0 : 2);

    // 获取当前象限的圆角半径 * 2
    float r = input.radius[idx];

    // 计算到圆角矩形边界的有符号距离
    float d = distance_from_rect_uv(p, input.size, r);

    // 改进的抗锯齿算法 - 使用特定的像素宽度计算
    float pixelWidth = 1.0;
    float2 ddx_p = ddx(p * input.size);
    float2 ddy_p = ddy(p * input.size);

    // 计算基于像素梯度的边缘宽度
    float edge_width = 0.5 * length(float2(length(ddx_p), length(ddy_p)));

    // 应用平滑过渡 - 提高抗锯齿质量
    float alpha_factor = 1.0 - smoothstep(-edge_width, edge_width, d);

    // 输出最终颜色
    float4 final_color = input.color;
    final_color.a *= alpha_factor;
    return final_color;
}