SnapOtter/apps/api/src/routes/tools/image-enhancement.ts

import { randomUUID } from "node:crypto";
import { mkdir } from "node:fs/promises";
import { join } from "node:path";
import { noiseRemoval } from "@snapotter/ai";
import { analyzeImage, applyCorrections } from "@snapotter/image-engine";
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
import sharp from "sharp";
import { z } from "zod";
import { autoOrient } from "../../lib/auto-orient.js";
import { isToolInstalled } from "../../lib/feature-status.js";
import { validateImageBuffer } from "../../lib/file-validation.js";
import { decodeToSharpCompat, needsCliDecode } from "../../lib/format-decoders.js";
import { decodeHeic } from "../../lib/heic-converter.js";
import { resolveOutputFormat } from "../../lib/output-format.js";
import { createWorkspace } from "../../lib/workspace.js";
import { createToolRoute } from "../tool-factory.js";

const settingsSchema = z.object({
  mode: z.enum(["auto", "portrait", "landscape", "low-light", "food", "document"]).default("auto"),
  intensity: z.number().min(0).max(100).default(50),
  corrections: z
    .object({
      exposure: z.boolean().default(true),
      contrast: z.boolean().default(true),
      whiteBalance: z.boolean().default(true),
      saturation: z.boolean().default(true),
      sharpness: z.boolean().default(true),
      denoise: z.boolean().default(true),
    })
    .default({}),
  deepEnhance: z.boolean().default(false),
});

type EnhancementSettings = z.infer<typeof settingsSchema>;

async function processImageEnhancement(
  rawBuffer: Buffer,
  settings: EnhancementSettings,
  filename: string,
) {
  const outputFormat = await resolveOutputFormat(rawBuffer, filename);

  // HDR/EXR decodes can produce 16-bit buffers; CLAHE requires 8-bit (VIPS_FORMAT_UCHAR)
  let inputBuffer = rawBuffer;
  const inputMeta = await sharp(inputBuffer).metadata();
  if (inputMeta.depth && inputMeta.depth !== "uchar") {
    inputBuffer = await sharp(inputBuffer).toColourspace("srgb").png().toBuffer();
  }

  const analysis = await analyzeImage(inputBuffer);
  const meta = await sharp(inputBuffer).metadata();
  const hasAlpha = meta.hasAlpha === true;

  let alphaBuffer: Buffer | undefined;
  if (hasAlpha) {
    alphaBuffer = await sharp(inputBuffer).extractChannel(3).toBuffer();
  }

  let image = sharp(inputBuffer);
  if (hasAlpha) {
    image = image.removeAlpha();
  }

  image = applyCorrections(
    image,
    analysis.corrections,
    settings.mode,
    settings.intensity,
    settings.corrections,
    { width: meta.width ?? 1, height: meta.height ?? 1 },
  );

  let buffer = await image
    .toFormat(outputFormat.format, { quality: outputFormat.quality })
    .toBuffer();

  if (alphaBuffer) {
    buffer = await sharp(buffer)
      .joinChannel(alphaBuffer)
      .toFormat(outputFormat.format, { quality: outputFormat.quality })
      .toBuffer();
  }

  if (settings.deepEnhance && isToolInstalled("noise-removal")) {
    try {
      const jobId = randomUUID();
      const workspacePath = await createWorkspace(jobId);
      const outputDir = join(workspacePath, "output");
      await mkdir(outputDir, { recursive: true });
      const result = await noiseRemoval(buffer, outputDir, {
        tier: "quality",
        strength: 35,
        detailPreservation: 70,
        colorNoise: 20,
      });
      buffer = result.buffer;
    } catch {
      // SCUNet unavailable -- fall back to Sharp-only result
    }
  }

  return { buffer, filename, contentType: outputFormat.contentType };
}

export function registerImageEnhancement(app: FastifyInstance) {
  createToolRoute(app, {
    toolId: "image-enhancement",
    settingsSchema,
    process: processImageEnhancement,
  });

  app.post(
    "/api/v1/tools/image-enhancement/analyze",
    async (request: FastifyRequest, reply: FastifyReply) => {
      let fileBuffer: Buffer | null = null;
      let filename = "image";

      try {
        const parts = request.parts();
        for await (const part of parts) {
          if (part.type === "file") {
            const chunks: Buffer[] = [];
            for await (const chunk of part.file) {
              chunks.push(chunk);
            }
            fileBuffer = Buffer.concat(chunks);
            filename = part.filename ?? "image";
            break;
          }
        }
      } catch (err) {
        return reply.status(400).send({
          error: "Failed to parse request",
          details: err instanceof Error ? err.message : String(err),
        });
      }

      if (!fileBuffer || fileBuffer.length === 0) {
        return reply.status(400).send({ error: "No image file provided" });
      }

      const validation = await validateImageBuffer(fileBuffer, filename);
      if (!validation.valid) {
        return reply.status(400).send({ error: `Invalid image: ${validation.reason}` });
      }

      if (validation.format === "heif") {
        try {
          fileBuffer = await decodeHeic(fileBuffer);
        } catch (err) {
          return reply.status(422).send({
            error: "Failed to decode HEIC file",
            details: err instanceof Error ? err.message : String(err),
          });
        }
      }

      // Decode CLI-decoded formats (RAW, TGA, PSD, EXR, HDR)
      if (needsCliDecode(validation.format)) {
        try {
          fileBuffer = await decodeToSharpCompat(fileBuffer, validation.format);
        } catch (err) {
          return reply.status(422).send({
            error: `Failed to decode ${validation.format} file`,
            details: err instanceof Error ? err.message : String(err),
          });
        }
      }

      try {
        fileBuffer = await autoOrient(fileBuffer);
        const analysis = await analyzeImage(fileBuffer);
        return reply.send(analysis);
      } catch (err) {
        return reply.status(422).send({
          error: "Analysis failed",
          details: err instanceof Error ? err.message : String(err),
        });
      }
    },
  );
}