server.js

const express = require("express");
const cors = require("cors");
const dotenv = require("dotenv");
const { GoogleGenerativeAI } = require("@google/generative-ai");
const {
  GoogleAIFileManager,
  GoogleAICacheManager,
} = require("@google/generative-ai/server");

const multer = require("multer");
const fs = require("fs");
const path = require("path");
const os = require("os");

const { Storage } = require("@google-cloud/storage"); // Import Cloud Storage
const { v4: uuidv4 } = require("uuid");

const { spawn } = require("node:child_process");
const { exec } = require("node:child_process");

dotenv.config();
const PORT = process.env.PORT || 10000;
const app = express();

// Configure Multer (Save files to disk temporarily, not memory)
const uploadDir = path.join(os.tmpdir(), "uploads"); // Ensure this directory exists
fs.existsSync(uploadDir) || fs.mkdirSync(uploadDir); // Create if it doesn't exist

const storage = multer.diskStorage({
  destination: uploadDir,
  filename: (req, file, cb) => {
    const ext = path.extname(file.originalname);
    cb(null, `${Date.now()}-${uuidv4()}${ext}`); // Unique filenames
  },
});
const upload = multer({
  storage,
  limits: { fileSize: 10 * 1024 * 1024 },
}).single("file");

app.use(
  cors({
    origin: ["http://localhost:5173", "https://prototype.datarai.com"],
    methods: ["GET", "POST"],
    credentials: true,
  })
);
app.use(express.json());

const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY);
const fileManager = new GoogleAIFileManager(process.env.GEMINI_API_KEY);
const cacheManager = new GoogleAICacheManager(process.env.GEMINI_API_KEY);

// Cloud Storage setup
const gcs = new Storage({ keyFilename: process.env.GCS_KEY_PATH }); // Make sure to set the path in .env
const bucketName = process.env.GCS_BUCKET_NAME; // Make sure to set the bucket name in .env

app.post("/upload", upload, async (req, res) => {
  try {
    const { projectId } = req.body;
    const file = req.file;

    if (!projectId || !file) {
      return res
        .status(400)
        .json({ error: "Project ID and file are required" });
    }

    //1. Cloud Storage Upload
    const objectName = file.filename; // Use the Multer-generated filename

    let gcsUri;
    try {
      gcsUri = await uploadToGCS(file.path, bucketName, objectName); // file.path is the temp file path
    } catch (gcsError) {
      console.error("GCS Upload Failed:", gcsError);
      fs.unlinkSync(file.path); // Clean up the temp file
      return res
        .status(500)
        .json({ error: "Failed to upload to Cloud Storage" });
    }

    // 2. GoogleAIFileManager upload (optional, for initial Gemini analysis, but not for download URL)
    // Remove this block if the uri from it is not used
    let uploadedFile;
    try {
      uploadedFile = await fileManager.uploadFile(file.path, {
        // Pass the actual existing path
        displayName: file.originalname,
        mimeType: file.mimetype,
      });
    } catch (fileManagerError) {
      console.error("File Manager Upload Failed:", fileManagerError);
      //Attempt to delete GCS file, but don't halt if it fails.
      try {
        await gcs.bucket(bucketName).file(objectName).delete();
        console.log(
          `Successfully deleted ${objectName} from GCS due to FileManager failure.`
        );
      } catch (deleteError) {
        console.warn(
          `Failed to delete ${objectName} from GCS after FileManager failure:`,
          deleteError
        );
      }

      fs.unlinkSync(file.path); // Clean up local temp file
      return res
        .status(500)
        .json({ error: "File upload to Google AI File Manager failed." });
    } finally {
      fs.unlinkSync(file.path); // Clean up local temp file always, regardless of filemanager success
    }

    console.log("Uploaded File Manager:", uploadedFile);

    return res.json({
      success: true,
      fileUri: uploadedFile.file.uri, // Keep if file manager uri used
      downloadUri: gcsUri, //Return gcs uri or file manager uri if you are using file manager uri for initial processing
      // downloadUri: uploadedFile.file.downloadUri,  // Don't return this anymore.
    });
  } catch (error) {
    console.error("Upload Error:", error);
    return res.status(500).json({ error: "File upload failed" });
  }
});

//Helper Functions
const uploadToGCS = async (filePath, bucketName, objectName) => {
  try {
    await gcs.bucket(bucketName).upload(filePath, {
      destination: objectName,
    });
    console.log(`${filePath} uploaded to ${bucketName}/${objectName}`);
    return `${bucketName}/${objectName}`; // Return GCS URI
  } catch (error) {
    console.error("GCS Upload Error:", error);
    throw error; // Propagate the error
  }
};

const generateSignedUrl = async (bucketName, objectName) => {
  try {
    const bucket = gcs.bucket(bucketName);
    const file = bucket.file(objectName);

    const config = {
      action: "read",
      expires: Date.now() + 60 * 60 * 1000, // 1 hour
    };

    const signedUrls = await file.getSignedUrl(config);
    const signedUrl = signedUrls[0];
    // console.log(`Generated signed URL for ${objectName}: ${signedUrl}`);
    return signedUrl;
  } catch (err) {
    console.error("ERROR: getting signed url", err);
    throw err;
  }
};

// Preprocesses Python code to include the signed URL
const preprocessPythonCode = (code, downloadUri) => {
  try {
    // Add any preprocessing steps here
    // replace anything inside read_csv() with the download uri
    code = code.replace(/read_csv\((.*?)\)/, `read_csv('${downloadUri}')`);

    code = "import io\nimport base64\n" + code;

    // return base64 string
    code = code.replace(
      /plt\.savefig\((.*?)\)/,
      `ioBytes = io.BytesIO()\nplt.savefig(ioBytes, format='png')\nioBytes.seek(0)\nbase64Image = base64.b64encode(ioBytes.read()).decode('utf-8')`
    );
    return code;
  } catch (error) {
    console.error("Error in preprocessPythonCode:", error);
    throw error;
  }
};

// // Executes code provided by gemini
// const executePythonCode = (code) => {
//   return new Promise((resolve, reject) => {
//       console.log("Executing Python code...");
//       const pythonProcess = spawn("python3", ["-c", code], {
//           maxBuffer: 1024 * 1024 * 1, // 1MB (Adjust as needed)
//       });

//       const outputChunks = []; // Array to hold Buffer chunks
//       let errorOutput = "";

//       pythonProcess.stdout.on("data", (data) => {
//           console.log("Data received from stdout:", data.toString());
//           outputChunks.push(data); // Accumulate Buffer chunks
//       });

//       pythonProcess.stderr.on("data", (data) => {
//           errorOutput += data.toString();
//       });

//       pythonProcess.on("close", (code) => {
//           if (code === 0) {
//               console.log("Python code executed successfully");
//               const output = Buffer.concat(outputChunks).toString();
//               console.log("Full base64 output length: ", output.length);
//               console.log("Full base64 output (first 200 chars): ", output.substring(0, 200));

//               resolve(output);
//           } else {
//               console.error("Python code execution failed:", errorOutput);
//               reject(new Error(`Python code execution failed with code ${code}: ${errorOutput}`));
//           }
//       });

//       pythonProcess.on("error", (err) => { // Handle process spawning errors
//           console.error("Failed to spawn Python process:", err);
//           reject(new Error(`Failed to spawn Python process: ${err.message}`));
//       });
//   });
// };

const executePythonCode = (code) => {
  return new Promise((resolve, reject) => {
    console.log("Executing Python code...");
    exec(
      `python3 -c "${code.replace(/"/g, '\\"')}"`,
      { maxBuffer: 1024 * 1024 * 1 },
      (error, stdout, stderr) => {
        if (error) {
          console.error("Python code execution failed:", stderr);
          reject(new Error(`Python code execution failed: ${stderr}`));
        } else {
          resolve(stdout.trim());
        }
      }
    );
  });
};

app.post("/gemini", async (req, res) => {
  try {
    const {
      messages,
      projectId,
      fileUri,
      downloadUri,
      visualizationMode,
      codeExecutionMode,
    } = req.body;

    if (!messages?.length || !projectId || !fileUri) {
      return res.status(400).json({ error: "Invalid request data" });
    }

    const session = {
      messages: [],
      fileUri: fileUri, //Store GCS URI, or file manager URI if you choose to use file manager uri
      downloadUri: downloadUri,
      projectId,
    };

    const messageHistory = messages.map((m) => ({
      role: m.sender === "user" ? "user" : "model",
      parts: [{ text: m.value }],
    }));

    let signedUrl = null;

    if (visualizationMode || codeExecutionMode) {
      // Only generate signed URL if needed
      // Extract object name from GCS URI

      try {
        console.log(downloadUri);
        const urlParts = downloadUri.split("/");
        const objectName = urlParts[3];
        signedUrl = await generateSignedUrl(bucketName, objectName); // generate for user selected file
      } catch (signedUrlError) {
        console.error("Failed to generate signed URL:", signedUrlError);
        return res
          .status(500)
          .json({ error: "Failed to generate signed URL." });
      }
    }

    const displayName = "data.csv";
    const useModel = "models/gemini-1.5-flash-001";
    const systemInstruction =
      `You are a data analysis expert. Use the provided data file ${fileUri} to assist users.\n` +
      "You offer short, concise answers to simple questions, unless the user asks for more depth.\n" +
      (!visualizationMode && !codeExecutionMode ? "If they ask for a statistic or something simple about the data, give the answer to them, then ask if they want code.\n" : "") +
      "Structure responses with good markdown formatting for ease of reading.\n" +
      (visualizationMode
        ? "Your response MUST include a Python script in triple backticks (```python code```) that: \n" +
          `- Loads data using pandas from ${signedUrl} \n` +
          "- Follows exact instructions to create good looking graph/chart \n" +
          `- Saves visualisations as a base64 string and prints the base64 representation with utf-8 encoding then add plt.close() \n` +
          "Efficient explanations of the code only \n"
        : "") +
      (codeExecutionMode
        ? "Your response MUST include a Python script in triple backticks (```python code```) that: \n" +
          `- Loads data using pandas from ${signedUrl} \n` +
          "- Follows exact instructions to answer the user's question \n" +
          `- Efficiently explains the answer and any interesting facts or observations \n` +
          `- Add a newline to every print statement for readability! \n` +
          "- Only accept the code output, not your own answer \n"
        : "") +
        "IF THE REQUEST DOES NOT PERTAIN TO THE DATA FILE, RESPOND WITH 'I'm sorry, I can't help with that.'";
    let ttlSeconds = 1 * 60 * 30; // 30 min

    let cache;

    // try {
    //   cache = await cacheManager.create({
    //     model: useModel,
    //     displayName,
    //     systemInstruction,
    //     contents: [
    //       {
    //         role: "user",
    //         parts: [
    //           {
    //             fileData: {
    //               mimeType: "text/csv",
    //               fileUri: fileUri,
    //             },
    //           }
    //         ],
    //       },
    //     ],
    //     ttlSeconds,
    //   });
    // } catch (cacheError) {
    //   console.error("Failed to create cache:", cacheError);
    //   return res.status(500).json({ error: "Failed to create cache entry." });
    // }

    // Add user messages
    session.messages.push({
      role: "user",
      parts: [
        {
          fileData: {
            fileUri: fileUri,
            mimeType: "text/csv",
          },
        },
      ],
    });
    session.messages.push(...messageHistory);
    session.messages.push({
      role: "user",
      parts: [{ text: systemInstruction }],
    });

    // const model = genAI.getGenerativeModelFromCachedContent(cache);
    const model = genAI.getGenerativeModel({ model: useModel });

    const chat = model.startChat({ history: session.messages });

    let result;
    try {
      result = await chat.sendMessage(messages[messages.length - 1].value, {
        maxTokens: 10000,
      });
    } catch (geminiError) {
      console.error("Gemini API call failed:", geminiError);
      return res.status(500).json({ error: "Gemini API call failed." });
    }

    let aiMessage = result.response.text();
    console.log(result.response.usageMetadata);
    // console.log("AI Response:", aiMessage);

    // Extract Python code from Gemini's response (assuming it's in triple backticks)
    const codeRegex = /```python\n([\s\S]*?)\n```/; //Matches a triple backtick block that specifies python
    const match = aiMessage.match(codeRegex);

    let pythonCode = null;

    if (match) {
      console.log("Python code found in Gemini response");
      pythonCode = match[1]; //Extracted code
      // console.log("Python code:", pythonCode);
    }

    if (visualizationMode || codeExecutionMode) {
      if (pythonCode) {
        //Execute code
        try {
          if (visualizationMode) {
            // let preprocessedCode;
            // try {
            //   preprocessedCode = preprocessPythonCode(pythonCode, signedUrl);
            // } catch (preprocessError) {
            //   console.error(
            //     "Error during preprocessPythonCode",
            //     preprocessError
            //   );
            //   return res
            //     .status(500)
            //     .json({ error: "Failed to preprocess Python code." });
            // }

            let codeOutput = "";
            let wasError = false;
            await executePythonCode(pythonCode)
              .then((base64) => {
                codeOutput = base64;
                console.log(codeOutput);
              })
              .catch((err) => {
                codeOutput = err;
                wasError = true;
                console.error("Python Execution Error:", err);
              });

            if (!wasError) {
              console.log(codeOutput);
              return res.json({
                type: "image",
                message: aiMessage,
                image: codeOutput,
              });
            } else {
              return res.json({
                type: "code",
                message: aiMessage + `\n\n**Error Message:**\n${codeOutput}`,
              });
            }
          }

          if (codeExecutionMode) {
            let codeOutput;
            await executePythonCode(pythonCode)
            .then((base64) => {
              codeOutput = base64;
              console.log(codeOutput);
            })
            .catch((err) => {
              codeOutput = err;
              wasError = true;
              console.error("Python Execution Error:", err);
            });

            return res.json({
              type: "code",
              message: aiMessage + `\n\n**Code Output:**\n${codeOutput}`,
            });
          }
        } catch (generalError) {
          // Catch-all for visualization/execution mode
          console.error(
            "Error during visualization/code execution mode:",
            generalError
          );
          return res
            .status(500)
            .json({ error: "Error during visualization/code execution." });
        }
      }
    }

    session.messages.push({ role: "model", parts: [{ text: aiMessage }] });

    return res.json({ type: "normal", message: aiMessage });
  } catch (error) {
    console.error("Gemini endpoint error:", error);
    res.status(500).json({ error: "Internal server error" });
  }
});

app.listen(PORT, () => {
  console.log(`Server running on port ${PORT}`);
});