import express from "express";
// import * as tf from "@tensorflow/tfjs";
import { fileURLToPath } from "url";
import { dirname, join } from "path";
import multer from "multer";
import sharp from "sharp";

let tf;
async function getTF() {
  if (!tf) {
    tf = await import("@tensorflow/tfjs");
  }
  return tf;
}

// Suppress TensorFlow and Node.js warnings
process.env.TF_CPP_MIN_LOG_LEVEL = "2"; // Suppress TensorFlow info messages
process.removeAllListeners("warning"); // Suppress Node.js deprecation warnings

// Initialize paths
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
const MODEL_PATH = join(__dirname, "../models/tfjs_model_dir/model.json");

const router = express.Router();
const upload = multer();

// 1. Model Loading
let model;
async function loadModel() {
  const tf = await getTF(); // Dynamic import
  try {
    console.log("⌛ Loading model...");
    model = await tf.loadGraphModel(
      "https://y23hqgxjqrmemzys.public.blob.vercel-storage.com/model.json",
    );

    // Warmup
    const warmupTensor = tf.zeros([1, 256, 256, 3]);
    model.execute(warmupTensor);
    tf.dispose(warmupTensor);

    console.log("✅ Model loaded");
  } catch (error) {
    console.error("❌ Model loading failed:", error);
    throw error;
  }
}
// await loadModel();

// 2. Image Preprocessing
const preprocessImage = async (buffer) => {
  try {
    const processed = await sharp(buffer)
      .resize(256, 256)
      .removeAlpha()
      .raw() // Get raw pixel data
      .toBuffer({ resolveWithObject: true });

    return tf.tidy(() => {
      const tensor = tf.tensor3d(processed.data, [
        processed.info.height,
        processed.info.width,
        3,
      ]);
      return tf.div(
        tf.sub(tf.expandDims(tensor, 0), [127.5, 127.5, 127.5]),
        [127.5, 127.5, 127.5],
      );
    });
  } catch (error) {
    console.error("Image preprocessing failed:", error);
    throw new Error("Failed to process image");
  }
};

// 3. Prediction Endpoint
router.post("/predict", upload.single("image"), async (req, res) => {
  if (!req.file) {
    return res.status(400).json({ error: "No image provided" });
  }
  if (!model) {
    try {
      await loadModel();
    } catch (error) {
      return res.status(500).json({ error: "Model failed to load" });
    }
  }

  const tf = await getTF(); // Get TF instance
  let inputTensor, outputTensor;
  try {
    // 1. Preprocess with verification
    inputTensor = await preprocessImage(req.file.buffer);
    console.log("✅ Input tensor shape:", inputTensor.shape);
    console.log(
      "✅ Input tensor values sample:",
      Array.from(inputTensor.dataSync()).slice(0, 5),
    );

    // 2. Get predictions with verification
    outputTensor = model.execute(inputTensor);
    const predictions = await outputTensor.data();
    console.log("🔥 Raw predictions:", Array.from(predictions));

    // 3. Verify model output
    if (!predictions || predictions.length === 0) {
      throw new Error("Model returned empty predictions");
    }

    // 4. Format results with safeguards
    const classNames = ["Karsinoma", "Kulit_sehat", "Melanoma"];

    if (predictions.length !== classNames.length) {
      console.warn(
        `⚠️ Warning: Expected ${classNames.length} outputs, got ${predictions.length}`,
      );
    }

    const probabilities = {};
    const predictionArray = Array.from(predictions);

    classNames.forEach((className, index) => {
      probabilities[className] =
        index < predictionArray.length
          ? parseFloat(predictionArray[index].toFixed(6))
          : 0; // Default to 0 if missing
    });

    const maxIndex = predictionArray.indexOf(Math.max(...predictionArray));
    const predictedClass =
      maxIndex !== -1
        ? classNames[Math.min(maxIndex, classNames.length - 1)]
        : "unknown";

    res.json({
      status: "success",
      prediction: {
        predicted_class: predictedClass,
        probabilities: probabilities,
        confidence:
          predictionArray.length > 0
            ? Math.max(...predictionArray).toFixed(6)
            : 0,
      },
    });
  } catch (error) {
    console.error("❌ Prediction error:", error);
    res.status(500).json({
      status: "error",
      error: error.message,
      stack: process.env.NODE_ENV === "development" ? error.stack : undefined,
    });
  } finally {
    if (tf && inputTensor) tf.dispose([inputTensor, outputTensor]);
  }
});

export default router;
