src/app/api/screenshot/route.ts

import { NextRequest, NextResponse } from "next/server";
import { getSession } from "auth/server";

const getBackendUrl = () => {
  if (process.env.BACKEND_API_URL) return process.env.BACKEND_API_URL;
  const spaceId = process.env.SPACE_ID || "";
  if (spaceId.startsWith("shyota/")) {
    return "https://shyota-claude-code-backend.hf.space";
  }
  return "https://augment17-claude-code-backend.hf.space";
};

export async function GET(req: NextRequest) {
  const session = await getSession();
  if (!session?.user) {
    return new NextResponse("Unauthorized", { status: 401 });
  }

  const backendUrl = getBackendUrl();
  const apiKey = process.env.BACKEND_API_KEY || "";

  try {
    const res = await fetch(`${backendUrl}/api/workspace/latest-screenshot`, {
      headers: {
        "Authorization": `Bearer ${apiKey}`,
      },
      cache: "no-store",
    });

    if (!res.ok) {
      return new NextResponse("No screenshot found", { status: res.status });
    }

    const buffer = await res.arrayBuffer();
    return new NextResponse(buffer, {
      headers: {
        "Content-Type": "image/png",
        "Cache-Control": "no-cache",
      },
    });
  } catch (err: any) {
    return new NextResponse(err.message, { status: 500 });
  }
}