server.js

// server.js
import express from "express";
import { Client } from "@gradio/client";
import dotenv from "dotenv";
import fetch from "node-fetch";
import multer from "multer";
const upload = multer();

dotenv.config();

const app = express();
const PORT = process.env.PORT || 7860;
const HF_TOKEN = process.env.HF_TOKEN;

// Reuse a single Gradio client so state (like dropdown choices)
// is preserved across related function calls.
const clientPromise = Client.connect("pockit-cloud/main", { hf_token: HF_TOKEN });

if (!HF_TOKEN) {
  console.error("HF_TOKEN is not set in environment variables!");
  process.exit(1);
}

app.use(express.json({ limit: "10mb" }));

app.post("/api/proxy", async (req, res) => {
  let { fn, args } = req.body;
  
  if (!fn) return res.status(400).json({ error: "Missing 'fn' in request body" });

  // THE FIX: Scrub the 'args' object clean before passing it to Gradio.
  // If your frontend accidentally sent 'read_token' in the payload, this deletes it.
  if (args && typeof args === 'object' && !Array.isArray(args)) {
    delete args.read_token;
    delete args.hf_token;
    delete args.write_token;
  }

  try {
    const client = await clientPromise;
    
    // Now 'args' only contains the actual model inputs
    const result = await client.predict(fn, args);
    res.json({ data: result.data });
  } catch (err) {
    console.error("Proxy error:", err);
    res.status(500).json({ error: "Proxy call failed", details: err.message });
  }
});

// Download endpoint
app.get("/api/download", async (req, res) => {
  const { file } = req.query;
  if (!file) return res.status(400).json({ error: "Missing 'file' query param" });

  const url =
    typeof file === "string" && file.startsWith("http")
      ? file
      : `https://huggingface.co/datasets/${file}`;
  try {
    const response = await fetch(url, {
      headers: { Authorization: `Bearer ${HF_TOKEN}` },
    });
    
    if (!response.ok) {
      return res.status(response.status).json({ error: "File fetch failed" });
    }
    
    res.setHeader("Content-Disposition", `attachment; filename=\"${file.split("/").pop()}\"`);
    res.setHeader("Content-Type", response.headers.get("content-type") || "application/octet-stream");
    
    response.body.pipe(res);
  } catch (err) {
    console.error("Download error:", err);
    res.status(500).json({ error: "Download failed", details: err.message });
  }
});

app.post("/api/upload", upload.single("file"), async (req, res) => {
  const { user_id, password } = req.body;
  if (!req.file || !user_id || !password) {
    return res.status(400).json({ error: "Missing file, user_id, or password" });
  }
  const custom_name = req.body.custom_name || req.file.originalname;
  try {
    const client = await clientPromise;
    const result = await client.predict("/upload_file_secure", {
      user_id,
      password,
      filepath: req.file.buffer,
      custom_name
    });
    res.json({ data: result.data });
  } catch (err) {
    console.error("Upload error:", err);
    res.status(500).json({ error: "Upload failed", details: err.message });
  }
});

app.post("/api/download-link", async (req, res) => {
  const { user_id, password, filename } = req.body || {};

  if (!user_id || !password || !filename) {
    return res.status(400).json({ error: "Missing user_id, password, or filename" });
  }

  try {
    const client = await clientPromise;
    const result = await client.predict("/get_download_link_action", {
      user_id,
      password,
      filename,
    });
    res.json({ data: result.data });
  } catch (err) {
    console.error("Download-link error:", err);
    res.status(500).json({ error: "Download link failed", details: err.message });
  }
});

app.use(express.static("."));

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