server.ts

import express from "express";
import path from "path";
import fs from "fs";
import { createServer as createViteServer } from "vite";
import dotenv from "dotenv";

dotenv.config();

const app = express();
const PORT = 3000;

// Set high limits for image base64 uploads
app.use(express.json({ limit: "50mb" }));
app.use(express.urlencoded({ limit: "50mb", extended: true }));

// Ensure base datasets directory exists
const DATASETS_ROOT = path.join(process.cwd(), "datasets");
if (!fs.existsSync(DATASETS_ROOT)) {
  fs.mkdirSync(DATASETS_ROOT, { recursive: true });
}

// Utility to resolve relative or absolute paths safely
function resolvePath(dirPath: string): string {
  if (path.isAbsolute(dirPath)) {
    return dirPath;
  }
  return path.join(process.cwd(), dirPath);
}

// Helper to scan for images
function scanImagesInDirectory(dirPath: string): string[] {
  const resolved = resolvePath(dirPath);
  if (!fs.existsSync(resolved)) {
    return [];
  }
  const files = fs.readdirSync(resolved);
  const extensions = [".png", ".jpg", ".jpeg", ".webp", ".bmp"];
  return files
    .filter((file) => extensions.includes(path.extname(file).toLowerCase()))
    .sort();
}

// Helper to check if mask exists
function getMaskForImage(imageName: string, outputDir: string): { hasMask: boolean; maskPath: string } {
  const resolvedOut = resolvePath(outputDir);
  // Masks are saved as PNG with matching base name or exact name
  const baseName = path.parse(imageName).name;
  const maskPath = path.join(resolvedOut, `${baseName}_mask.png`);
  return {
    hasMask: fs.existsSync(maskPath),
    maskPath: maskPath,
  };
}

// Helper to check if JSON annotations exist
function getAnnotationsForImage(imageName: string, outputDir: string): { annotations: any | null } {
  const resolvedOut = resolvePath(outputDir);
  const baseName = path.parse(imageName).name;
  const jsonPath = path.join(resolvedOut, `${baseName}_annotations.json`);
  if (fs.existsSync(jsonPath)) {
    try {
      const content = fs.readFileSync(jsonPath, "utf-8");
      return { annotations: JSON.parse(content) };
    } catch (e) {
      console.error("Error parsing annotations JSON:", e);
    }
  }
  return { annotations: null };
}

// API Routes

// 1. Get standard or custom datasets
app.get("/api/datasets", (req, res) => {
  try {
    const subfolders = fs.readdirSync(DATASETS_ROOT).filter((file) => {
      return fs.statSync(path.join(DATASETS_ROOT, file)).isDirectory();
    });

    const datasets = subfolders.map((folder) => {
      const inputDir = path.join("datasets", folder, "input");
      const outputDir = path.join("datasets", folder, "masks");
      const images = scanImagesInDirectory(inputDir);
      return {
        name: folder,
        inputDir,
        outputDir,
        imageCount: images.length,
      };
    });

    res.json({ success: true, datasets });
  } catch (err: any) {
    res.status(500).json({ success: false, error: err.message });
  }
});

// 2. Scan custom directory (manual input/output folders)
app.post("/api/scan-directory", (req, res) => {
  try {
    const { inputDir, outputDir } = req.body;
    if (!inputDir || !outputDir) {
      return res.status(400).json({ success: false, error: "Input and Output directory paths are required." });
    }

    const resolvedInput = resolvePath(inputDir);
    const resolvedOutput = resolvePath(outputDir);

    // Create directories if they do not exist
    if (!fs.existsSync(resolvedInput)) {
      fs.mkdirSync(resolvedInput, { recursive: true });
    }
    if (!fs.existsSync(resolvedOutput)) {
      fs.mkdirSync(resolvedOutput, { recursive: true });
    }

    const images = scanImagesInDirectory(inputDir);
    const imageDetails = images.map((img) => {
      const { hasMask } = getMaskForImage(img, outputDir);
      const { annotations } = getAnnotationsForImage(img, outputDir);
      return {
        name: img,
        hasMask,
        annotations,
      };
    });

    res.json({
      success: true,
      inputDir,
      outputDir,
      images: imageDetails,
    });
  } catch (err: any) {
    res.status(500).json({ success: false, error: err.message });
  }
});

// 3. Upload a sample image (handles client-side offscreen generated sample images)
app.post("/api/upload-sample", (req, res) => {
  try {
    const { datasetName, imageName, imageDataUrl } = req.body;
    if (!datasetName || !imageName || !imageDataUrl) {
      return res.status(400).json({ success: false, error: "Missing required parameters." });
    }

    const targetInputFolder = path.join(DATASETS_ROOT, datasetName, "input");
    const targetOutputFolder = path.join(DATASETS_ROOT, datasetName, "masks");

    if (!fs.existsSync(targetInputFolder)) {
      fs.mkdirSync(targetInputFolder, { recursive: true });
    }
    if (!fs.existsSync(targetOutputFolder)) {
      fs.mkdirSync(targetOutputFolder, { recursive: true });
    }

    // Decode base64 PNG/JPG
    const matches = imageDataUrl.match(/^data:image\/([a-zA-Z+]+);base64,(.+)$/);
    if (!matches || matches.length !== 3) {
      return res.status(400).json({ success: false, error: "Invalid image data URL format." });
    }

    const imageBuffer = Buffer.from(matches[2], "base64");
    const filePath = path.join(targetInputFolder, imageName);
    fs.writeFileSync(filePath, imageBuffer);

    res.json({ success: true, filePath });
  } catch (err: any) {
    res.status(500).json({ success: false, error: err.message });
  }
});

// 4. Serve an image from any directory
app.get("/api/image", (req, res) => {
  try {
    const relativePath = req.query.path as string;
    if (!relativePath) {
      return res.status(400).json({ error: "Path parameter is required" });
    }

    const resolved = resolvePath(relativePath);
    if (!fs.existsSync(resolved)) {
      return res.status(404).json({ error: "Image file not found: " + relativePath });
    }

    res.sendFile(resolved);
  } catch (err: any) {
    res.status(500).json({ error: err.message });
  }
});

// 5. Serve a mask PNG from output directory
app.get("/api/mask", (req, res) => {
  try {
    const relativePath = req.query.path as string;
    if (!relativePath) {
      return res.status(400).json({ error: "Path parameter is required" });
    }

    const resolved = resolvePath(relativePath);
    if (!fs.existsSync(resolved)) {
      return res.status(404).json({ error: "Mask file not found: " + relativePath });
    }

    res.sendFile(resolved);
  } catch (err: any) {
    res.status(500).json({ error: err.message });
  }
});

// 6. Save annotation PNG mask and JSON coordinates
app.post("/api/save-mask", (req, res) => {
  try {
    const { imageName, outputDir, maskDataUrl, annotations } = req.body;
    if (!imageName || !outputDir) {
      return res.status(400).json({ success: false, error: "Missing required parameters." });
    }

    const resolvedOutput = resolvePath(outputDir);
    if (!fs.existsSync(resolvedOutput)) {
      fs.mkdirSync(resolvedOutput, { recursive: true });
    }

    const baseName = path.parse(imageName).name;

    // If annotations are cleared/empty, delete both the mask PNG and annotations JSON files
    if (annotations && annotations.length === 0) {
      const maskPath = path.join(resolvedOutput, `${baseName}_mask.png`);
      const jsonPath = path.join(resolvedOutput, `${baseName}_annotations.json`);
      if (fs.existsSync(maskPath)) fs.unlinkSync(maskPath);
      if (fs.existsSync(jsonPath)) fs.unlinkSync(jsonPath);
      return res.json({ success: true, message: `Successfully cleared mask and annotations for ${imageName}` });
    }

    if (!maskDataUrl) {
      return res.status(400).json({ success: false, error: "Missing maskDataUrl." });
    }

    // Decode and save PNG mask
    const matches = maskDataUrl.match(/^data:image\/([a-zA-Z+]+);base64,(.+)$/);
    if (matches && matches.length === 3) {
      const maskBuffer = Buffer.from(matches[2], "base64");
      const maskPath = path.join(resolvedOutput, `${baseName}_mask.png`);
      fs.writeFileSync(maskPath, maskBuffer);
    } else {
      return res.status(400).json({ success: false, error: "Invalid mask data URL." });
    }

    // Save JSON coordinates for persistence & editing
    const jsonPath = path.join(resolvedOutput, `${baseName}_annotations.json`);
    fs.writeFileSync(jsonPath, JSON.stringify(annotations, null, 2));

    res.json({ success: true, message: `Successfully saved mask for ${imageName}` });
  } catch (err: any) {
    res.status(500).json({ success: false, error: err.message });
  }
});

// Setup Vite Dev Server / Serve static build
async function startServer() {
  if (process.env.NODE_ENV !== "production") {
    const vite = await createViteServer({
      server: { middlewareMode: true },
      appType: "spa",
    });
    app.use(vite.middlewares);
  } else {
    const distPath = path.join(process.cwd(), "dist");
    app.use(express.static(distPath));
    app.get("*", (req, res) => {
      res.sendFile(path.join(distPath, "index.html"));
    });
  }

  app.listen(PORT, "0.0.0.0", () => {
    console.log(`Express server running on http://localhost:${PORT}`);
  });
}

startServer();