import fs from "node:fs/promises";
import path from "node:path";
import { PNG } from "pngjs";
import pixelmatch from "pixelmatch";
import { chromium } from "playwright";
import { config } from "./config.js";
import { buildHtmlDocument, type CodeBundle } from "./codeDocument.js";
import type { CheckDefinition, CheckResult, JudgeResult } from "./types.js";

async function renderScreenshot(bundle: CodeBundle, outputPath: string) {
  await fs.mkdir(path.dirname(outputPath), { recursive: true });

  const browser = await chromium.launch({ headless: true });
  const context = await browser.newContext({
    viewport: config.viewport,
    javaScriptEnabled: true
  });

  await context.route("**/*", (route) => {
    const url = route.request().url();
    if (url.startsWith("data:") || url.startsWith("about:")) {
      route.continue();
      return;
    }

    route.abort();
  });

  const page = await context.newPage();
  page.setDefaultTimeout(1500);
  page.on("dialog", (dialog) => dialog.dismiss().catch(() => undefined));
  await page.setContent(buildHtmlDocument(bundle), {
    waitUntil: "domcontentloaded",
    timeout: 2000
  });
  await page.waitForTimeout(150);
  await page.screenshot({ path: outputPath, fullPage: false });
  return { browser, page };
}

function parseChecks(checksJson: string): CheckDefinition[] {
  try {
    const checks = JSON.parse(checksJson);
    return Array.isArray(checks) ? checks : [];
  } catch {
    return [];
  }
}

function formatPercent(value: number) {
  return `${(value * 100).toFixed(1)}%`;
}

function checkWeight(check: CheckDefinition) {
  return typeof check.weight === "number" && Number.isFinite(check.weight) && check.weight > 0 ? check.weight : 1;
}

export function summarizeWeightedResults(results: CheckResult[], passingScore = 0.8) {
  const totalWeight = results.reduce((sum, result) => sum + (result.weight ?? 1), 0);
  const passedWeight = results.reduce((sum, result) => sum + (result.passed ? result.weight ?? 1 : 0), 0);
  const score = totalWeight === 0 ? 1 : Number((passedWeight / totalWeight).toFixed(4));

  return {
    score,
    passed: score >= passingScore
  };
}

type Pixel = {
  r: number;
  g: number;
  b: number;
  a: number;
};

function pixelAt(image: PNG, index: number): Pixel {
  const offset = index * 4;
  return {
    r: image.data[offset],
    g: image.data[offset + 1],
    b: image.data[offset + 2],
    a: image.data[offset + 3]
  };
}

function colorDistance(a: Pixel, b: Pixel) {
  const dr = a.r - b.r;
  const dg = a.g - b.g;
  const db = a.b - b.b;
  const da = (a.a - b.a) / 2;
  return Math.sqrt(dr * dr + dg * dg + db * db + da * da);
}

function cornerBackground(image: PNG) {
  return pixelAt(image, 0);
}

function contentComparison(actual: PNG, expected: PNG) {
  const actualBackground = cornerBackground(actual);
  const expectedBackground = cornerBackground(expected);
  const contentTolerance = 18;
  const mismatchTolerance = 30;
  let contentPixels = 0;
  let contentMismatches = 0;

  for (let index = 0; index < actual.width * actual.height; index += 1) {
    const actualPixel = pixelAt(actual, index);
    const expectedPixel = pixelAt(expected, index);
    const isExpectedContent = colorDistance(expectedPixel, expectedBackground) > contentTolerance;
    const isActualContent = colorDistance(actualPixel, actualBackground) > contentTolerance;
    const isMismatch = colorDistance(actualPixel, expectedPixel) > mismatchTolerance;

    if (isExpectedContent || isActualContent || isMismatch) {
      contentPixels += 1;
      if (isMismatch) {
        contentMismatches += 1;
      }
    }
  }

  if (contentPixels === 0) {
    return {
      similarity: 1,
      contentPixels,
      contentMismatches
    };
  }

  return {
    similarity: Number((1 - contentMismatches / contentPixels).toFixed(4)),
    contentPixels,
    contentMismatches
  };
}

async function compareScreenshots(actualPath: string, expectedPath: string, diffPath: string) {
  const [actualBuffer, expectedBuffer] = await Promise.all([
    fs.readFile(actualPath),
    fs.readFile(expectedPath)
  ]);
  const actual = PNG.sync.read(actualBuffer);
  const expected = PNG.sync.read(expectedBuffer);

  if (actual.width !== expected.width || actual.height !== expected.height) {
    return {
      similarity: 0,
      diffPath: undefined,
      details: `Screenshot size ${actual.width}x${actual.height} did not match expected ${expected.width}x${expected.height}.`
    };
  }

  const diff = new PNG({ width: actual.width, height: actual.height });
  const mismatchedPixels = pixelmatch(actual.data, expected.data, diff.data, actual.width, actual.height, {
    threshold: 0.12
  });
  await fs.mkdir(path.dirname(diffPath), { recursive: true });
  await fs.writeFile(diffPath, PNG.sync.write(diff));

  const totalPixels = actual.width * actual.height;
  const content = contentComparison(actual, expected);
  return {
    similarity: content.similarity,
    diffPath,
    details: `${content.contentMismatches} of ${content.contentPixels} content pixels differed. Full screenshot: ${mismatchedPixels} of ${totalPixels} pixels differed.`
  };
}

export async function generateExpectedScreenshot(taskId: string, bundle: CodeBundle) {
  const outputPath = path.join(config.generatedDir, "expected", `${taskId}.png`);
  const { browser } = await renderScreenshot(bundle, outputPath);
  await browser.close();
  return path.relative(config.dataDir, outputPath).replaceAll(path.sep, "/");
}

export async function judgeSubmission(options: {
  taskId: string;
  checksJson: string;
  expectedScreenshotPath?: string | null;
  visualThreshold: number;
  passingScore?: number;
  bundle: CodeBundle;
}): Promise<JudgeResult> {
  const submissionScreenshot = path.join(
    config.generatedDir,
    "submissions",
    `${options.taskId}-${Date.now()}.png`
  );
  const { browser, page } = await renderScreenshot(options.bundle, submissionScreenshot);
  const checks = parseChecks(options.checksJson);
  const results: CheckResult[] = [];
  let diffScreenshotPath: string | undefined;

  try {
    for (const check of checks) {
      if (check.type === "visual-match") {
        const threshold = check.threshold ?? options.visualThreshold;
        if (!options.expectedScreenshotPath) {
          results.push({
            type: check.type,
            message: check.message,
            passed: false,
            weight: checkWeight(check),
            details: "No expected screenshot has been generated for this task."
          });
          continue;
        }

        const expectedPath = path.join(config.dataDir, options.expectedScreenshotPath);
        const diffPath = path.join(config.generatedDir, "diffs", `${options.taskId}-${Date.now()}.png`);
        const comparison = await compareScreenshots(submissionScreenshot, expectedPath, diffPath);
        if (comparison.diffPath) {
          diffScreenshotPath = path.relative(config.dataDir, comparison.diffPath).replaceAll(path.sep, "/");
        }
        results.push({
          type: check.type,
          message: check.message,
          passed: comparison.similarity >= threshold,
          score: comparison.similarity,
          weight: checkWeight(check),
          details: `${formatPercent(comparison.similarity)} match required ${formatPercent(threshold)}. ${comparison.details}`
        });
        continue;
      }

      if (!check.selector) {
        results.push({
          type: check.type,
          message: check.message,
          passed: false,
          weight: checkWeight(check),
          details: "This check is missing a CSS selector."
        });
        continue;
      }

      const locator = page.locator(check.selector);
      const count = await locator.count();

      if (check.type === "selector-exists") {
        results.push({
          type: check.type,
          message: check.message,
          passed: count > 0,
          weight: checkWeight(check),
          details: count > 0 ? `Found ${count} matching element(s).` : "No matching element was found."
        });
      }

      if (check.type === "element-count") {
        const expectedCount = check.count ?? Number(check.value ?? 0);
        results.push({
          type: check.type,
          message: check.message,
          passed: count === expectedCount,
          weight: checkWeight(check),
          details: `Found ${count}, expected ${expectedCount}.`
        });
      }

      if (check.type === "text-contains") {
        const text = count > 0 ? await locator.first().textContent() : "";
        const expected = check.value ?? "";
        results.push({
          type: check.type,
          message: check.message,
          passed: Boolean(text?.toLocaleLowerCase("en-US").includes(expected.toLocaleLowerCase("en-US"))),
          weight: checkWeight(check),
          details: `Expected text containing "${expected}".`
        });
      }

      if (check.type === "attribute-equals") {
        const actual = count > 0 ? await locator.first().getAttribute(check.attribute ?? "") : null;
        const expected = check.value ?? "";
        results.push({
          type: check.type,
          message: check.message,
          passed: actual === expected,
          weight: checkWeight(check),
          details: `Expected ${check.attribute}="${expected}", found ${actual ?? "nothing"}.`
        });
      }

      if (check.type === "css-property") {
        const actual = count > 0
          ? await locator.first().evaluate((el, property) => getComputedStyle(el).getPropertyValue(property), check.property ?? "")
          : "";
        const expected = check.value ?? "";
        results.push({
          type: check.type,
          message: check.message,
          passed: actual.trim() === expected.trim(),
          weight: checkWeight(check),
          details: `Expected ${check.property}: ${expected}, found ${actual || "nothing"}.`
        });
      }
    }
  } finally {
    await browser.close();
  }

  const passingScore = options.passingScore ?? 0.8;
  const summary = summarizeWeightedResults(results, passingScore);
  return {
    passed: summary.passed,
    score: summary.score,
    passingScore,
    results,
    screenshotPath: path.relative(config.dataDir, submissionScreenshot).replaceAll(path.sep, "/"),
    diffScreenshotPath
  };
}