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;

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 {
    // Auth is handled cleanly right here
    const client = await Client.connect("pockit-cloud/main", { hf_token: HF_TOKEN });
    
    // 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 = `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.file.originalname; // Required - use filename to avoid "blob"
  const fileData = {
    data: [new Uint8Array(req.file.buffer)],
    orig_name: req.file.originalname,
    size: req.file.size,
    mime_type: req.file.mimetype,
    is_stream: false
  };
  try {
    const client = await Client.connect("pockit-cloud/main", { hf_token: HF_TOKEN });
    const result = await client.predict("/upload_file_secure", { 
      user_id, 
      password, 
      filepath: fileData,
      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.use(express.static("."));

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