import { afterAll, beforeAll, expect, test } from "vitest";
import {
  checkerboardTexture,
  createSampler,
  gradientTexture,
  solidTexture,
} from "wesl-gpu";
import { testFragment } from "../TestFragmentShader.ts";
import { destroySharedDevice, getGPUDevice } from "../WebGPUTestSetup.ts";

let device: GPUDevice;

beforeAll(async () => {
  device = await getGPUDevice();
});

afterAll(() => {
  destroySharedDevice();
});

test("renders simple constant color", async () => {
  const src = `
    @fragment
    fn fs_main() -> @location(0) vec4f {
      return vec4f(0.5, 0.25, 0.75, 1.0);
    }
  `;
  const projectDir = import.meta.url;
  const textureFormat: GPUTextureFormat = "rgba32float";
  const params = { projectDir, device, src, textureFormat };
  const result = await testFragment(params);

  expect(result).toHaveLength(4);
  expect(result[0]).toBeCloseTo(0.5);
  expect(result[1]).toBeCloseTo(0.25);
  expect(result[2]).toBeCloseTo(0.75);
  expect(result[3]).toBeCloseTo(1.0);
});

test("derivative of x coordinate", async () => {
  const src = `
    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      let dx = dpdx(pos.x);
      return vec4f(pos.x, dx, 0.0, 1.0);
    }
  `;
  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    textureFormat: "rg32float",
    size: [2, 2],
  });
  // result at pixel (0, 0)
  const [x, dx] = result;

  expect(x).toBeCloseTo(0.5);
  expect(dx).toBeCloseTo(1);
});

test("samples solid color texture", async () => {
  const inputTex = solidTexture(device, [0.5, 0.5, 0.5, 1.0], 256, 256);
  const sampler = createSampler(device);

  const src = `
    @group(0) @binding(1) var input_tex: texture_2d<f32>;
    @group(0) @binding(2) var input_samp: sampler;

    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      let uv = pos.xy / 256.0;
      return textureSample(input_tex, input_samp, uv);
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    textures: [inputTex],
    samplers: [sampler],
  });

  expect(result[0]).toBeCloseTo(0.5);
  expect(result[1]).toBeCloseTo(0.5);
  expect(result[2]).toBeCloseTo(0.5);
  expect(result[3]).toBeCloseTo(1.0);
});

test("samples gradient texture at center", async () => {
  const inputTex = gradientTexture(device, 256, 256, "horizontal");
  const sampler = createSampler(device);

  const src = `
    @group(0) @binding(1) var input_tex: texture_2d<f32>;
    @group(0) @binding(2) var input_samp: sampler;

    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      return textureSample(input_tex, input_samp, vec2f(0.5, 0.5));
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    textures: [inputTex],
    samplers: [sampler],
  });

  expect(result[0]).toBeCloseTo(0.5, 1);
  expect(result[1]).toBeCloseTo(0.5, 1);
  expect(result[2]).toBeCloseTo(0.5, 1);
});

test("samples checkerboard texture", async () => {
  const inputTex = checkerboardTexture(device, 256, 256, 128);
  const sampler = createSampler(device);

  const src = `
    @group(0) @binding(1) var input_tex: texture_2d<f32>;
    @group(0) @binding(2) var input_samp: sampler;

    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      // Sample at (0.25, 0.25) - should be black (0.0)
      return textureSample(input_tex, input_samp, vec2f(0.25, 0.25));
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    textures: [inputTex],
    samplers: [sampler],
  });

  expect(result[0]).toBeCloseTo(0.0);
  expect(result[1]).toBeCloseTo(0.0);
  expect(result[2]).toBeCloseTo(0.0);
});

test("samples multiple textures", async () => {
  const tex1 = solidTexture(device, [1.0, 0.0, 0.0, 1.0], 64, 64);
  const tex2 = solidTexture(device, [0.0, 1.0, 0.0, 1.0], 64, 64);
  const sampler = createSampler(device);

  const src = `
    @group(0) @binding(1) var tex1: texture_2d<f32>;
    @group(0) @binding(2) var tex2: texture_2d<f32>;
    @group(0) @binding(3) var samp: sampler;

    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      let uv = vec2f(0.5, 0.5);
      let c1 = textureSample(tex1, samp, uv);
      let c2 = textureSample(tex2, samp, uv);
      return c1 * 0.5 + c2 * 0.5;
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    textures: [tex1, tex2],
    samplers: [sampler],
  });

  expect(result[0]).toBeCloseTo(0.5); // (1.0 + 0.0) / 2
  expect(result[1]).toBeCloseTo(0.5); // (0.0 + 1.0) / 2
  expect(result[2]).toBeCloseTo(0.0); // (0.0 + 0.0) / 2
});

test("uses scalar constant from constants namespace", async () => {
  const src = `
    import constants::BRIGHTNESS;

    @fragment
    fn fs_main() -> @location(0) vec4f {
      return vec4f(BRIGHTNESS, 0.0, 0.0, 1.0);
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    constants: { BRIGHTNESS: 0.75 },
  });

  expect(result[0]).toBeCloseTo(0.75);
});

test("uses vector constant from constants namespace", async () => {
  const src = `
    import constants::COLOR;

    @fragment
    fn fs_main() -> @location(0) vec4f {
      return vec4f(COLOR, 0.0, 1.0);
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    constants: { COLOR: "vec2f(0.25, 0.5)" },
  });

  expect(result[0]).toBeCloseTo(0.25);
  expect(result[1]).toBeCloseTo(0.5);
  expect(result[2]).toBeCloseTo(0.0);
  expect(result[3]).toBeCloseTo(1.0);
});

test("uses conditions for conditional compilation", async () => {
  const src = `
    @fragment
    fn fs_main() -> @location(0) vec4f {
      @if(USE_RED)
      return vec4f(1.0, 0.0, 0.0, 1.0);
      @else
      return vec4f(0.0, 1.0, 0.0, 1.0);
    }
  `;

  const resultRed = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    conditions: { USE_RED: true },
  });

  expect(resultRed[0]).toBeCloseTo(1.0);
  expect(resultRed[1]).toBeCloseTo(0.0);

  const resultGreen = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    conditions: { USE_RED: false },
  });

  expect(resultGreen[0]).toBeCloseTo(0.0);
  expect(resultGreen[1]).toBeCloseTo(1.0);
});

test("uses both conditions and constants together", async () => {
  const src = `
    @if(USE_CUSTOM_COLOR)
    import constants::CUSTOM_COLOR;

    @fragment
    fn fs_main() -> @location(0) vec4f {
      @if(USE_CUSTOM_COLOR)
      return vec4f(CUSTOM_COLOR, 0.0, 1.0);
      @else
      return vec4f(0.0, 0.0, 0.0, 1.0);
    }
  `;

  const resultWithColor = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    conditions: { USE_CUSTOM_COLOR: true },
    constants: { CUSTOM_COLOR: "vec2f(0.8, 0.6)" },
  });

  expect(resultWithColor[0]).toBeCloseTo(0.8);
  expect(resultWithColor[1]).toBeCloseTo(0.6);
  expect(resultWithColor[2]).toBeCloseTo(0.0);

  const resultWithoutColor = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    conditions: { USE_CUSTOM_COLOR: false },
  });

  expect(resultWithoutColor[0]).toBeCloseTo(0.0);
  expect(resultWithoutColor[1]).toBeCloseTo(0.0);
  expect(resultWithoutColor[2]).toBeCloseTo(0.0);
});

test("shader with resolution uniform (auto-populated)", async () => {
  const src = `
    import env::u;

    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      let st = pos.xy / u.resolution;
      return vec4f(st.x, st.y, 0.0, 1.0);
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    size: [256, 256],
    // resolution auto-populated as [256, 256]
  });

  // Pixel (0,0) is at fragment coordinate (0.5, 0.5)
  // Normalized: (0.5/256, 0.5/256, 0, 1)
  expect(result[0]).toBeCloseTo(0.5 / 256, 4);
  expect(result[1]).toBeCloseTo(0.5 / 256, 4);
  expect(result[2]).toBe(0.0);
  expect(result[3]).toBe(1.0);
});

test("shader with time uniform", async () => {
  const src = `
    import env::u;

    @fragment
    fn fs_main() -> @location(0) vec4f {
      return vec4f(u.time, u.time * 2.0, u.time * 3.0, 1.0);
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    uniforms: { time: 5.0 },
  });

  expect(result[0]).toBeCloseTo(5.0);
  expect(result[1]).toBeCloseTo(10.0);
  expect(result[2]).toBeCloseTo(15.0);
  expect(result[3]).toBeCloseTo(1.0);
});

test("shader with resolution and time uniforms", async () => {
  const src = `
    import env::u;

    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      let st = pos.xy / u.resolution;
      return vec4f(st.x, st.y, u.time / 10.0, 1.0);
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    size: [256, 256],
    uniforms: { time: 5.0 },
  });

  // pixel (0,0) is at fragment (0.5, 0.5)
  expect(result[0]).toBeCloseTo(0.5 / 256, 4); // st.x
  expect(result[1]).toBeCloseTo(0.5 / 256, 4); // st.y
  expect(result[2]).toBeCloseTo(0.5, 4); // time/10
});

test("shader with uniforms and texture", async () => {
  const inputTex = solidTexture(device, [0.5, 0.5, 0.5, 1.0], 64, 64);
  const sampler = createSampler(device);

  const src = `
    import env::u;
    @group(0) @binding(1) var input_tex: texture_2d<f32>;
    @group(0) @binding(2) var input_samp: sampler;

    @fragment
    fn fs_main(@builtin(position) pos: vec4f) -> @location(0) vec4f {
      let uv = pos.xy / u.resolution;
      let tex_color = textureSample(input_tex, input_samp, uv);
      let time_mod = vec4f(u.time * 0.1);
      return tex_color + time_mod;
    }
  `;

  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
    size: [64, 64],
    uniforms: { time: 10.0 },
    textures: [inputTex],
    samplers: [sampler],
  });

  // 0.5 (texture) + 1.0 (time * 0.1 where time=10) = 1.5
  expect(result[0]).toBeCloseTo(1.5, 2);
});

test("fragment with @test_texture and @sampler annotations", async () => {
  const src = `
    @test_texture(solid, 1, 0, 0, 1) var tex: texture_2d<f32>;
    @sampler(linear) var samp: sampler;

    @fragment
    fn fs_main() -> @location(0) vec4f {
      return textureSampleLevel(tex, samp, vec2f(0.5), 0.0);
    }
  `;
  const result = await testFragment({
    projectDir: import.meta.url,
    device,
    src,
  });
  expect(result[0]).toBeCloseTo(1.0);
  expect(result[1]).toBeCloseTo(0.0);
  expect(result[2]).toBeCloseTo(0.0);
  expect(result[3]).toBeCloseTo(1.0);
});