//====================================================================
// fragment_shader.v - 片段着色器
// 第05课：颜色插值、Lambert光照、雾效计算
//====================================================================
module fragment_shader #(
    parameter COORD_WIDTH = 16,
    parameter COLOR_WIDTH = 24,
    parameter FRAC_BITS  = 12
)(
    input  wire                          clk, rst_n,
    input  wire                          frag_in_valid,
    input  wire [COORD_WIDTH-1:0]        frag_in_x, frag_in_y,
    input  wire [COLOR_WIDTH-1:0]        frag_in_color,
    input  wire signed [COORD_WIDTH-1:0] frag_in_depth,
    input  wire signed [COORD_WIDTH-1:0] frag_in_nx, frag_in_ny, frag_in_nz,
    output reg                           frag_in_ready,
    input  wire signed [COORD_WIDTH-1:0] light_dir_x, light_dir_y, light_dir_z,
    input  wire [COLOR_WIDTH-1:0]        light_color, ambient_color,
    input  wire                          fog_enable,
    input  wire [COLOR_WIDTH-1:0]        fog_color,
    input  wire signed [COORD_WIDTH-1:0] fog_start, fog_end,
    output reg                           frag_out_valid,
    output reg  [COORD_WIDTH-1:0]        frag_out_x, frag_out_y,
    output reg  [COLOR_WIDTH-1:0]        frag_out_color,
    output reg  signed [COORD_WIDTH-1:0] frag_out_depth,
    input  wire                          frag_out_ready
);
/* verilator lint_off WIDTHEXPAND */
/* verilator lint_off WIDTHTRUNC */
/* verilator lint_off CASEOVERLAP */
/* verilator lint_off CMPCONST */
/* verilator lint_off UNSIGNED */
/* verilator lint_off WIDTHCONCAT */

    reg [7:0] color_r, color_g, color_b, lit_r, lit_g, lit_b, final_r, final_g, final_b;
    reg signed [2*COORD_WIDTH-1:0] dot_nl;
    reg [7:0] diffuse_factor, fog_factor;
    reg [COORD_WIDTH-1:0] pass_x, pass_y;
    reg signed [COORD_WIDTH-1:0] pass_depth;
    localparam S_IDLE=3'd0, S_LIGHT=3'd1, S_FOG=3'd2, S_OUTPUT=3'd3;
    reg [2:0] state;
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin state<=S_IDLE; frag_in_ready<=1'b1; frag_out_valid<=1'b0; end
        else begin
            frag_out_valid <= 1'b0;
            case (state)
                S_IDLE: begin
                    frag_in_ready <= 1'b1;
                    if (frag_in_valid) begin
                        frag_in_ready <= 1'b0;
                        color_r <= frag_in_color[23:16]; color_g <= frag_in_color[15:8]; color_b <= frag_in_color[7:0];
                        pass_x <= frag_in_x; pass_y <= frag_in_y; pass_depth <= frag_in_depth;
                        state <= S_LIGHT;
                    end
                end
                S_LIGHT: begin
                    dot_nl = frag_in_nx * light_dir_x + frag_in_ny * light_dir_y + frag_in_nz * light_dir_z;
                    diffuse_factor = (dot_nl > 0) ? ((dot_nl[19:12] > 8'd255) ? 8'd255 : dot_nl[19:12]) : 8'd0;
                    lit_r = (ambient_color[23:16] >> 2) + ((color_r * diffuse_factor) >> 8);
                    lit_g = (ambient_color[15:8] >> 2) + ((color_g * diffuse_factor) >> 8);
                    lit_b = (ambient_color[7:0] >> 2) + ((color_b * diffuse_factor) >> 8);
                    if (lit_r > 8'd255) lit_r = 8'd255;
                    if (lit_g > 8'd255) lit_g = 8'd255;
                    if (lit_b > 8'd255) lit_b = 8'd255;
                    state <= S_FOG;
                end
                S_FOG: begin
                    if (fog_enable) begin
                        fog_factor = (pass_depth < fog_start) ? 8'd255 : ((pass_depth > fog_end) ? 8'd0 : 8'd128);
                        final_r = (lit_r * fog_factor + fog_color[23:16] * (8'd255 - fog_factor)) >> 8;
                        final_g = (lit_g * fog_factor + fog_color[15:8] * (8'd255 - fog_factor)) >> 8;
                        final_b = (lit_b * fog_factor + fog_color[7:0] * (8'd255 - fog_factor)) >> 8;
                    end else begin final_r=lit_r; final_g=lit_g; final_b=lit_b; end
                    state <= S_OUTPUT;
                end
                S_OUTPUT: begin
                    frag_out_valid <= 1'b1; frag_out_x <= pass_x; frag_out_y <= pass_y;
                    frag_out_color <= {final_r, final_g, final_b}; frag_out_depth <= pass_depth;
                    state <= S_IDLE;
                end
                default: state <= S_IDLE;
            endcase
        end
    end
endmodule