import fs from "fs/promises";
import path from "path";

const API_ENDPOINT = "http://localhost:5229/api/image/generate";
const NUM_REQUESTS_PER_SCENARIO = 10;
const OUTPUT_DIR = "benchmark_output";

const colors = {
  reset: "\x1b[0m",
  bright: "\x1b[1m",
  dim: "\x1b[2m",
  fg: {
    cyan: "\x1b[36m",
    green: "\x1b[32m",
    yellow: "\x1b[33m",
    red: "\x1b[31m",
  },
};

const scenarios = [
  {
    name: "Small Canvas, Small Crop",
    payload: {
      canvas_rect: "A1:B2",
      crop_offset: [0.25, 0.25],
      crop_size: [0.5, 0.5],
      output_scale: 1.0,
    },
  },
  {
    name: "Medium Canvas, Full Crop (Stitching)",
    payload: {
      canvas_rect: "C3:F6",
      crop_offset: [0.0, 0.0],
      crop_size: [1.0, 1.0],
      output_scale: 0.5,
    },
  },
  {
    name: "Large Canvas, Small Corner Crop",
    payload: {
      canvas_rect: "A1:H12",
      crop_offset: [0.0, 0.0],
      crop_size: [0.1, 0.1],
      output_scale: 1.0,
    },
  },
  {
    name: "Large Canvas, Center Crop",
    payload: {
      canvas_rect: "A1:AE55",
      crop_offset: [0.45, 0.45],
      crop_size: [0.1, 0.1],
      output_scale: 1.0,
    },
  },
  {
    name: "Tall Canvas, Full Crop & Scale",
    payload: {
      canvas_rect: "A1:A20",
      crop_offset: [0.0, 0.0],
      crop_size: [1.0, 1.0],
      output_scale: 0.1,
    },
  },
  {
    name: "Wide Canvas, Full Crop & Scale",
    payload: {
      canvas_rect: "A1:T1",
      crop_offset: [0.0, 0.0],
      crop_size: [1.0, 1.0],
      output_scale: 0.1,
    },
  },
  {
    name: "Single Tile, Full Crop",
    payload: {
      canvas_rect: "K16:K16",
      crop_offset: [0.0, 0.0],
      crop_size: [1.0, 1.0],
      output_scale: 1.0,
    },
  },
  {
    name: "Single Tile, Partial Crop",
    payload: {
      canvas_rect: "K16:K16",
      crop_offset: [0.1, 0.1],
      crop_size: [0.25, 0.25],
      output_scale: 1.0,
    },
  },
  {
    name: "Large Canvas, Bottom-Right Crop",
    payload: {
      canvas_rect: "A1:AE55",
      crop_offset: [0.95, 0.95],
      crop_size: [0.05, 0.05],
      output_scale: 1.0,
    },
  },
  {
    name: "Wide Canvas, Thin Horizontal Crop",
    payload: {
      canvas_rect: "A1:AE10",
      crop_offset: [0.0, 0.5],
      crop_size: [1.0, 0.01],
      output_scale: 1.0,
    },
  },
  {
    name: "Tall Canvas, Thin Vertical Crop",
    payload: {
      canvas_rect: "A1:J55",
      crop_offset: [0.5, 0.0],
      crop_size: [0.01, 1.0],
      output_scale: 1.0,
    },
  },
  {
    name: "Medium Canvas, Heavy Scaling",
    payload: {
      canvas_rect: "D4:G8",
      crop_offset: [0.0, 0.0],
      crop_size: [1.0, 1.0],
      output_scale: 0.05,
    },
  },
  {
    name: "MAXIMUM CANVAS (Streaming Test)",
    payload: {
      canvas_rect: "A1:AE55",
      crop_offset: [0.0, 0.0],
      crop_size: [1.0, 1.0],
      output_scale: 0.01,
    },
  },
];

const calculateStats = (times) => {
  if (times.length === 0) return null;
  const sum = times.reduce((a, b) => a + b, 0);
  const avg = sum / times.length;
  const stdDev = Math.sqrt(
    times.map((x) => Math.pow(x - avg, 2)).reduce((a, b) => a + b, 0) /
      times.length
  );
  const sorted = [...times].sort((a, b) => a - b);
  return {
    avg: avg,
    stdDev: stdDev,
    median: sorted[Math.floor(sorted.length / 2)],
    min: sorted[0],
    max: sorted[sorted.length - 1],
    throughput: 1000 / avg,
  };
};

async function runBenchmark() {
  console.log(
    `${colors.bright}${colors.fg.cyan}--- Starting Image Stitcher API Benchmark ---${colors.reset}`
  );

  try {
    await fs.mkdir(OUTPUT_DIR, { recursive: true });
  } catch (e) {
    console.error(
      `${colors.fg.red}Error creating output directory: ${e.message}${colors.reset}`
    );
    return;
  }

  const allResults = [];

  for (const scenario of scenarios) {
    console.log(
      `\n${colors.fg.yellow}Running Scenario: '${scenario.name}' (${NUM_REQUESTS_PER_SCENARIO} requests)...${colors.reset}`
    );
    const responseTimes = [];

    for (let i = 0; i < NUM_REQUESTS_PER_SCENARIO; i++) {
      try {
        const startTime = performance.now();
        const response = await fetch(API_ENDPOINT, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(scenario.payload),
        });
        const endTime = performance.now();

        if (response.ok) {
          responseTimes.push(endTime - startTime);
          if (i === 0) {
            const arrayBuffer = await response.arrayBuffer();
            const buffer = Buffer.from(arrayBuffer);
            const fileName = `${scenario.name
              .replace(/[\s(),]/g, "_")
              .toLowerCase()}.png`;
            const filePath = path.join(OUTPUT_DIR, fileName);
            await fs.writeFile(filePath, buffer);
            console.log(
              `  ${colors.dim}Saved output image to '${filePath}'${colors.reset}`
            );
          } else {
            await response.arrayBuffer();
          }
        } else {
          console.log(
            `  ${colors.fg.red}Request ${i + 1} failed with status ${
              response.status
            }${colors.reset}`
          );
        }
      } catch (e) {
        console.error(
          `  ${colors.fg.red}Request ${i + 1} failed with an exception: ${
            e.message
          }${colors.reset}`
        );

        break;
      }
    }

    const stats = calculateStats(responseTimes);
    if (stats) {
      allResults.push({ name: scenario.name, stats });
      console.log(
        `  ${colors.fg.green}Average Latency: ${stats.avg.toFixed(2)} ms${
          colors.reset
        }`
      );
    }
  }

  console.log(
    `\n${colors.bright}${colors.fg.cyan}--- Benchmark Summary ---${colors.reset}`
  );
  for (const result of allResults) {
    console.log(`\n${colors.bright}Scenario: ${result.name}${colors.reset}`);
    const { stats } = result;
    console.log(
      `  Avg Latency:    ${stats.avg.toFixed(2)} ms (+/- ${stats.stdDev.toFixed(
        2
      )} ms)`
    );
    console.log(
      `  Details (ms):   Median: ${stats.median.toFixed(
        2
      )} | Min: ${stats.min.toFixed(2)} | Max: ${stats.max.toFixed(2)}`
    );
    console.log(`  Avg Throughput: ${stats.throughput.toFixed(2)} req/s`);
  }
  console.log(
    `\n${colors.bright}${colors.fg.cyan}--- Benchmark Complete ---${colors.reset}`
  );
}

runBenchmark();