src/routes/api/upload-to-dataset.ts

import { createFileRoute } from "@tanstack/react-router";
import { zipSync, strToU8 } from "fflate";
import { z } from "zod";
import { writeFileSync, unlinkSync } from "fs";
import { execSync } from "child_process";

const PostSchema = z.object({
  id: z.string(),
  fileUrl: z.string().url(),
  ext: z.string().default("jpg"),
  previewUrl: z.string().optional(),
  sampleUrl: z.string().optional(),
  tags: z.array(z.string()).default([]),
  source: z.string().optional(),
  rating: z.string().optional(),
  score: z.number().optional(),
  fileSize: z.number().optional(),
  width: z.number().optional(),
  height: z.number().optional(),
  md5: z.string().optional(),
  uploader: z.union([z.string(), z.number()]).optional(),
  createdAt: z.string().optional(),
  rawJson: z.string().optional(),
});

const BodySchema = z.object({
  siteName: z.string().default("grabber"),
  posts: z.array(PostSchema).min(1).max(200),
  excludeTags: z.array(z.string()).default([]),
  hfToken: z.string().min(1, "HF token is required"),
  datasetName: z.string().min(1, "Dataset name is required"),
  zipName: z.string().optional(),
  subfolder: z.string().optional(),
});

function sanitize(s: string) {
  return s.replace(/[^a-zA-Z0-9._-]+/g, "_").slice(0, 60);
}

function humanSize(n?: number): string {
  if (n == null || !Number.isFinite(n)) return "";
  if (n < 1024) return `${n} B`;
  if (n < 1024 * 1024) return `${(n / 1024).toFixed(1)} KB`;
  if (n < 1024 * 1024 * 1024) return `${(n / 1024 / 1024).toFixed(2)} MB`;
  return `${(n / 1024 / 1024 / 1024).toFixed(2)} GB`;
}

function buildMarkdown(post: z.infer<typeof PostSchema>, excludeSet: Set<string>): string {
  const dims = post.width && post.height ? `${post.width} x ${post.height}` : "";
  const rows: [string, string][] = [
    ["Source", post.source ?? ""],
    ["Rating", post.rating ?? ""],
    ["Score", post.score != null ? String(post.score) : ""],
    ["Dimensions", dims],
    ["File size", humanSize(post.fileSize)],
    ["MD5", post.md5 ?? ""],
    ["Uploader", post.uploader != null ? String(post.uploader) : ""],
    ["Created at", post.createdAt ?? ""],
    ["Original file", post.fileUrl],
    ["Sample", post.sampleUrl ?? ""],
    ["Preview", post.previewUrl ?? ""],
    ["Ext", post.ext],
  ];

  const meta = rows.map(([k, v]) => `- **${k}:** ${v}`).join("\n");
  const filteredTags = post.tags.filter((t) => !excludeSet.has(t.toLowerCase()));
  const tags = filteredTags.length ? filteredTags.join(" ") : "";
  let raw = "{}";
  if (post.rawJson) {
    try {
      raw = JSON.stringify(JSON.parse(post.rawJson), null, 2);
    } catch {
      raw = post.rawJson;
    }
  }

  return `# Post ${post.id}

${meta}

## Tags

${tags}

## Raw metadata

\`\`\`json
${raw}
\`\`\`
`;
}

export const Route = createFileRoute("/api/upload-to-dataset")({
  server: {
    handlers: {
      POST: async ({ request }) => {
        let body: z.infer<typeof BodySchema>;
        try {
          body = BodySchema.parse(await request.json());
        } catch (err) {
          return new Response(
            JSON.stringify({ error: err instanceof Error ? err.message : "Invalid body" }),
            { status: 400, headers: { "Content-Type": "application/json" } },
          );
        }

        const files: Record<string, Uint8Array> = {};
        const errors: string[] = [];
        const siteName = sanitize(body.siteName);
        const excludeSet = new Set((body.excludeTags || []).map((t) => t.toLowerCase()));

        // Generate files
        await Promise.all(
          body.posts.map(async (post) => {
            const idSafe = sanitize(post.id);
            const extSafe = sanitize(post.ext || "jpg");
            const base = `${siteName}_${idSafe}`;

            const filteredTags = post.tags.filter((t) => !excludeSet.has(t.toLowerCase()));

            files[`${base}.txt`] = strToU8(filteredTags.join(" "));
            files[`${idSafe}.md`] = strToU8(buildMarkdown(post, excludeSet));

            try {
              const res = await fetch(post.fileUrl, {
                headers: { "User-Agent": "LovableGrabber/1.0" },
              });
              if (!res.ok) {
                errors.push(`${post.id}: HTTP ${res.status}`);
                return;
              }
              const buf = new Uint8Array(await res.arrayBuffer());
              files[`${base}.${extSafe}`] = buf;
            } catch (err) {
              errors.push(`${post.id}: ${err instanceof Error ? err.message : "fetch failed"}`);
            }
          }),
        );

        if (errors.length) {
          files["_errors.txt"] = strToU8(errors.join("\n"));
        }

        if (Object.keys(files).length === 0) {
          return new Response(JSON.stringify({ error: "No files successfully compiled" }), {
            status: 502,
            headers: { "Content-Type": "application/json" },
          });
        }

        // Build ZIP Sync
        const zipped = zipSync(files, { level: 0 });
        const ts = new Date().toISOString().replace(/[:.]/g, "-");
        
        let filename = "";
        if (body.zipName && body.zipName.trim()) {
          filename = sanitize(body.zipName.trim());
          if (!filename.endsWith(".zip")) {
            filename += ".zip";
          }
        } else {
          filename = `grabber-${ts}.zip`;
        }

        let pathInRepo = filename;
        if (body.subfolder && body.subfolder.trim()) {
          const cleanSubfolder = body.subfolder
            .trim()
            .split("/")
            .map((p) => p.trim())
            .filter(Boolean)
            .join("/");
          if (cleanSubfolder) {
            pathInRepo = `${cleanSubfolder}/${filename}`;
          }
        }

        // Temporary paths - use safe name for local file
        const safeLocalFilename = `${ts}_${sanitize(filename)}`;
        const zipPath = `/tmp/${safeLocalFilename}`;
        const pyScriptPath = `/tmp/hf_dataset_upload_${ts}.py`;

        try {
          // Write compiled zip to local temp folder
          writeFileSync(zipPath, zipped);

          // Write python uploader helper script
          const pyCode = `
import sys
import os
from huggingface_hub import HfApi

token = sys.argv[1]
repo_id = sys.argv[2]
file_path = sys.argv[3]
path_in_repo = sys.argv[4]

api = HfApi()
try:
    # Attempt to create the dataset repository if not exists
    api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, token=token)
    
    # Upload the compiled zip to the dataset repository
    api.upload_file(
        path_or_fileobj=file_path,
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type="dataset",
        token=token
    )
    print("SUCCESS")
    sys.exit(0)
except Exception as e:
    print(f"ERROR: {str(e)}", file=sys.stderr)
    sys.exit(1)
`;
          writeFileSync(pyScriptPath, pyCode);

          // Execute python helper to create/upload the dataset
          execSync(
            `python3 "${pyScriptPath}" "${body.hfToken}" "${body.datasetName}" "${zipPath}" "${pathInRepo}"`,
          );

          return new Response(JSON.stringify({ success: true, filename: pathInRepo }), {
            status: 200,
            headers: { "Content-Type": "application/json" },
          });
        } catch (execErr: unknown) {
          let errMsg = "Execution error";
          if (execErr instanceof Error) {
            const stderr = (execErr as { stderr?: Buffer }).stderr;
            errMsg = stderr ? stderr.toString() : execErr.message;
          }
          return new Response(
            JSON.stringify({ error: `Hugging Face dataset upload failed: ${errMsg}` }),
            {
              status: 500,
              headers: { "Content-Type": "application/json" },
            },
          );
        } finally {
          // Clean up temporary files
          try {
            unlinkSync(zipPath);
          } catch {
            // ignore
          }
          try {
            unlinkSync(pyScriptPath);
          } catch {
            // ignore
          }
        }
      },
    },
  },
});