tasq/supabase/functions/extract_topology/index.ts

// extract_topology — Supabase Edge Function
// Reads a network-imports file from storage, sends it to Gemini with a strict
// JSON-schema prompt, and writes the parsed draft back to network_imports.
//
// Auth: caller must be admin or it_staff. Mirrors teams_crud auth pattern.
// Gemini model selection: 2.5 Flash by default; 2.5 Pro for PDFs >5 pages.
//
// Env vars required:
//   SUPABASE_URL, SUPABASE_ANON_KEY, SUPABASE_SERVICE_ROLE_KEY, GEMINI_API_KEY

import { serve } from "https://deno.land/std@0.203.0/http/server.ts";
import { createClient } from "https://esm.sh/@supabase/supabase-js@2";

const CORS_HEADERS: Record<string, string> = {
  "Access-Control-Allow-Origin": "*",
  "Access-Control-Allow-Methods": "POST, OPTIONS",
  "Access-Control-Allow-Headers":
    "Content-Type, Authorization, apikey, x-client-info, x-requested-with, Origin, Accept",
  "Access-Control-Max-Age": "3600",
};

const jsonResponse = (body: unknown, status = 200) =>
  new Response(JSON.stringify(body), {
    status,
    headers: { "Content-Type": "application/json", ...CORS_HEADERS },
  });

// Strict response schema we send to Gemini. Items get confidence scores so the
// review UI can rank "definitely a switch" above "might be a switch".
const RESPONSE_SCHEMA = {
  type: "object",
  properties: {
    sites: {
      type: "array",
      items: {
        type: "object",
        properties: {
          name: { type: "string" },
          address: { type: "string" },
        },
        required: ["name"],
      },
    },
    devices: {
      type: "array",
      items: {
        type: "object",
        properties: {
          name: { type: "string" },
          kind: {
            type: "string",
            enum: ["router", "switch", "ap", "firewall", "server", "endpoint", "patch_panel", "other"],
          },
          vendor: { type: "string" },
          model: { type: "string" },
          mgmt_ip: { type: "string" },
          site_name: { type: "string" },
          location_label: { type: "string" },
          confidence: { type: "string", enum: ["high", "medium", "low"] },
        },
        required: ["name", "kind"],
      },
    },
    ports: {
      type: "array",
      items: {
        type: "object",
        properties: {
          device_name: { type: "string" },
          port_number: { type: "string" },
          port_kind: {
            type: "string",
            enum: ["rj45", "sfp", "sfp_plus", "console", "power", "wireless", "virtual", "other"],
          },
          access_vlan: { type: "integer" },
          is_trunk: { type: "boolean" },
        },
        required: ["device_name", "port_number"],
      },
    },
    links: {
      type: "array",
      items: {
        type: "object",
        properties: {
          device_a: { type: "string" },
          port_a: { type: "string" },
          device_b: { type: "string" },
          port_b: { type: "string" },
          link_kind: {
            type: "string",
            enum: ["copper", "fiber", "wireless", "virtual", "unknown"],
          },
          confidence: { type: "string", enum: ["high", "medium", "low"] },
        },
        required: ["device_a", "device_b"],
      },
    },
    vlans: {
      type: "array",
      items: {
        type: "object",
        properties: {
          vlan_id: { type: "integer" },
          name: { type: "string" },
          description: { type: "string" },
        },
        required: ["vlan_id", "name"],
      },
    },
    notes: { type: "string" },
  },
};

const EXTRACTION_PROMPT = `You are a network-documentation reader. Look at the attached file (which may be a network diagram in PDF, an image of a diagram, a photo of a whiteboard, a Visio export, or a spreadsheet/notes) and extract every network device, port, link, VLAN, and site you can identify.

Rules:
- Output ONLY structured JSON conforming to the provided schema. No prose.
- For devices, classify kind precisely (switch / router / ap / etc.). If you genuinely cannot tell, use "other".
- For ports and links, give the exact labels as they appear ("Gi0/1", "Port 24", "uplink", etc.) — do not normalize.
- For each device and each link, attach a "confidence": "high" if the document clearly states it, "medium" if it's reasonably implied, "low" if it's a guess.
- If you cannot extract anything useful, return all arrays as empty and put a brief explanation in "notes".
- Do not invent items. Better to return less with high confidence than more with low.
- VLANs: only include if a VLAN ID number is visible.
- Sites: only include if a site/building name is clearly indicated.`;

interface ImportRow {
  id: string;
  storage_path: string;
  file_kind: string;
  status: string;
  created_by: string;
}

function pickGeminiModel(fileKind: string, byteLength: number): string {
  // Cheap default; escalate to Pro for big PDFs where layout reasoning matters.
  if (fileKind === "pdf" && byteLength > 2 * 1024 * 1024) {
    return "gemini-2.5-pro";
  }
  return "gemini-2.5-flash";
}

function mimeForFileKind(fileKind: string): string {
  switch (fileKind) {
    case "pdf":
      return "application/pdf";
    case "image":
      // Default to PNG; Gemini accepts most image MIMEs and will sniff if needed.
      return "image/png";
    case "vsdx":
      return "application/vnd.ms-visio.drawing";
    case "csv":
      return "text/csv";
    case "xlsx":
      return "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet";
    case "docx":
      return "application/vnd.openxmlformats-officedocument.wordprocessingml.document";
    default:
      return "application/octet-stream";
  }
}

function bytesToBase64(bytes: Uint8Array): string {
  // Deno has no btoa for binary; build base64 in chunks to avoid stack overflow.
  let binary = "";
  const chunkSize = 0x8000;
  for (let i = 0; i < bytes.length; i += chunkSize) {
    const slice = bytes.subarray(i, i + chunkSize);
    binary += String.fromCharCode.apply(null, Array.from(slice));
  }
  return btoa(binary);
}

serve(async (req) => {
  if (req.method === "OPTIONS") {
    return new Response(null, { status: 204, headers: CORS_HEADERS });
  }
  if (req.method !== "POST") {
    return jsonResponse({ error: "Method not allowed" }, 405);
  }

  const supabaseUrl = Deno.env.get("SUPABASE_URL") ?? "";
  const anonKey = Deno.env.get("SUPABASE_ANON_KEY") ?? "";
  const serviceKey = Deno.env.get("SUPABASE_SERVICE_ROLE_KEY") ?? "";
  const geminiKey = Deno.env.get("GEMINI_API_KEY") ?? "";
  if (!supabaseUrl || !anonKey || !serviceKey) {
    return jsonResponse({ error: "Missing Supabase env config" }, 500);
  }
  if (!geminiKey) {
    return jsonResponse({ error: "GEMINI_API_KEY not configured" }, 500);
  }

  // Auth caller
  const authHeader = req.headers.get("Authorization") ?? "";
  const token = authHeader.replace("Bearer ", "").trim();
  if (!token) return jsonResponse({ error: "Missing access token" }, 401);

  const authClient = createClient(supabaseUrl, anonKey, {
    global: { headers: { Authorization: `Bearer ${token}` } },
  });
  const { data: authData, error: authError } = await authClient.auth.getUser();
  if (authError || !authData?.user) return jsonResponse({ error: "Unauthorized" }, 401);

  const adminClient = createClient(supabaseUrl, serviceKey);

  // Caller must be admin or it_staff
  const { data: profile, error: profileError } = await adminClient
    .from("profiles")
    .select("role")
    .eq("id", authData.user.id)
    .maybeSingle();
  const role = (profile?.role ?? "").toString().toLowerCase();
  if (profileError || (role !== "admin" && role !== "it_staff")) {
    return jsonResponse({ error: "Forbidden" }, 403);
  }

  // Body: { import_id: string }
  let body: { import_id?: string } = {};
  try {
    body = await req.json();
  } catch (_e) {
    return jsonResponse({ error: "Invalid JSON body" }, 400);
  }
  const importId = body.import_id;
  if (!importId || typeof importId !== "string") {
    return jsonResponse({ error: "Missing import_id" }, 400);
  }

  // Load the import row
  const { data: imp, error: impErr } = await adminClient
    .from("network_imports")
    .select("id, storage_path, file_kind, status, created_by")
    .eq("id", importId)
    .maybeSingle<ImportRow>();
  if (impErr || !imp) {
    return jsonResponse({ error: "Import not found" }, 404);
  }
  if (imp.status === "applied" || imp.status === "discarded") {
    return jsonResponse({ error: `Import already ${imp.status}` }, 409);
  }

  // Mark extracting
  await adminClient
    .from("network_imports")
    .update({ status: "extracting", error_message: null })
    .eq("id", importId);

  try {
    // Download from storage
    const { data: fileData, error: dlErr } = await adminClient
      .storage
      .from("network-imports")
      .download(imp.storage_path);
    if (dlErr || !fileData) {
      throw new Error(dlErr?.message ?? "Failed to download import file");
    }

    const arrayBuf = await fileData.arrayBuffer();
    const bytes = new Uint8Array(arrayBuf);
    const base64 = bytesToBase64(bytes);
    const mime = mimeForFileKind(imp.file_kind);
    const model = pickGeminiModel(imp.file_kind, bytes.length);

    // Call Gemini
    const geminiUrl =
      `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent?key=${geminiKey}`;
    const geminiBody = {
      contents: [
        {
          role: "user",
          parts: [
            { text: EXTRACTION_PROMPT },
            { inlineData: { mimeType: mime, data: base64 } },
          ],
        },
      ],
      generationConfig: {
        responseMimeType: "application/json",
        responseSchema: RESPONSE_SCHEMA,
        temperature: 0.1,
      },
    };

    const geminiRes = await fetch(geminiUrl, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(geminiBody),
    });

    if (!geminiRes.ok) {
      const errBody = await geminiRes.text();
      throw new Error(`Gemini ${geminiRes.status}: ${errBody.slice(0, 500)}`);
    }

    const geminiJson = await geminiRes.json();
    const textOut: string =
      geminiJson?.candidates?.[0]?.content?.parts?.[0]?.text ?? "";

    let parsed: unknown;
    try {
      parsed = JSON.parse(textOut);
    } catch (_e) {
      throw new Error("Gemini returned non-JSON output");
    }

    // Minimum-quality gate: must have at least one device, or status=failed.
    const deviceCount = Array.isArray((parsed as { devices?: unknown[] })?.devices)
      ? (parsed as { devices: unknown[] }).devices.length
      : 0;

    if (deviceCount === 0) {
      await adminClient
        .from("network_imports")
        .update({
          status: "failed",
          ai_result: parsed as Record<string, unknown>,
          error_message: "No devices extracted from document. Start from blank canvas.",
        })
        .eq("id", importId);
      return jsonResponse(
        { status: "failed", reason: "no_devices", ai_result: parsed },
        200,
      );
    }

    // Store + mark for review
    await adminClient
      .from("network_imports")
      .update({
        status: "review",
        ai_result: parsed as Record<string, unknown>,
        error_message: null,
      })
      .eq("id", importId);

    return jsonResponse({ status: "review", ai_result: parsed });
  } catch (err) {
    const msg = (err as Error).message ?? String(err);
    await adminClient
      .from("network_imports")
      .update({ status: "failed", error_message: msg })
      .eq("id", importId);
    return jsonResponse({ error: msg }, 500);
  }
});