app/api/augment/route.ts

/* eslint-disable @typescript-eslint/no-explicit-any */
import { NextRequest, NextResponse } from "next/server";
import { GoogleGenAI } from "@google/genai";
import JSON5 from "json5";
import { InferenceClient } from "@huggingface/inference";
import type { AugmentRequest, AugmentResponse, FileUpdate } from "@/types";

const SYS = `You are Omni Engine, an autonomous code evolution system.
INPUTS:
- project context (tree + critical files)
- a user instruction

OUTPUT:
- STRICT JSON ONLY. No prose. No code fences. No leading or trailing text.
- Prefer this object shape:
{
  "ok": true,
  "summary": string,
  "logs": string[],
  "files": [ { "path": string, "action": "add"|"update"|"delete", "content"?: string, "note"?: string } ]
}
- If you cannot produce an object, return a top-level JSON array of the same file objects (files array).

RULES:
- Paths must be POSIX-like and rooted from repo (e.g., /frontend/src/App.tsx)
- For add/update include full file content (runnable)
- Keep changes consistent and compilable for the chosen language/framework
- DO NOT include markdown fences or explanation outside JSON
`;

function cleanText(raw: string): string {
  let t = raw || "";
  t = t.replace(/```[a-zA-Z]*\n?/g, "");
  t = t.replace(/```/g, "");
  return t.trim();
}

function tryParseAny(jsonish: string): any {
  // Try strict JSON first
  try {
    return JSON.parse(jsonish);
  } catch {
    // Relaxed JSON5 parse (handles trailing commas, etc.)
    try {
      // Prefer full string
      return JSON5.parse(jsonish);
    } catch {
      // Try to extract object
      const sObj = jsonish.indexOf("{");
      const eObj = jsonish.lastIndexOf("}");
      if (sObj >= 0 && eObj > sObj) {
        const candidate = jsonish.slice(sObj, eObj + 1);
        try { return JSON.parse(candidate); } catch {}
        try { return JSON5.parse(candidate); } catch {}
      }
      // Try to extract array
      const sArr = jsonish.indexOf("[");
      const eArr = jsonish.lastIndexOf("]");
      if (sArr >= 0 && eArr > sArr) {
        const candidate = jsonish.slice(sArr, eArr + 1);
        try { return JSON.parse(candidate); } catch {}
        try { return JSON5.parse(candidate); } catch {}
      }
      throw new Error("Model did not return valid JSON");
    }
  }
}

function findJsonCandidates(text: string, maxCandidates = 8): string[] {
  const candidates: string[] = [];
  const cleaned = cleanText(text);

  // 1) Code-fenced JSON blocks (already stripped in cleanText, so skip)

  // 2) Balanced scanning for objects and arrays
  const openers = ['{', '['] as const;
  for (let start = 0; start < cleaned.length; start++) {
    const ch = cleaned[start];
    if (!openers.includes(ch as any)) continue;
    const isArray = ch === '[';
    const closer = isArray ? ']' : '}';

    let depth = 0;
    let inStr: null | '"' | "'" = null;
    let prev = '';
    for (let i = start; i < cleaned.length; i++) {
      const c = cleaned[i];
      if (inStr) {
        if (c === inStr && prev !== '\\') inStr = null as any;
      } else {
        if (c === '"' || c === "'") inStr = c as any;
        else if (c === ch) depth++;
        else if (c === closer) depth--;
        if (depth === 0) {
          const slice = cleaned.slice(start, i + 1).trim();
          if ((isArray && slice.startsWith('[') && slice.endsWith(']')) || (!isArray && slice.startsWith('{') && slice.endsWith('}'))) {
            candidates.push(slice);
            if (candidates.length >= maxCandidates) return candidates;
          }
          break;
        }
      }
      prev = c;
    }
  }
  return candidates;
}

export async function POST(req: NextRequest) {
  try {
    const body = (await req.json()) as AugmentRequest;
    const { context, instruction, language = "javascript", framework = "express-react", response_type = "file_updates", model, provider } = body || {} as AugmentRequest;
    if (!context || !instruction) {
      return NextResponse.json({ ok: false, message: "Missing context or instruction" } as AugmentResponse, { status: 400 });
    }

    const userPrompt = `Project Context:\n${context}\n\nInstruction:\n${instruction}\n\nTarget:\n- language: ${language}\n- framework: ${framework}\n- response_type: ${response_type}`;

    let text = "";
    if ((provider || "").toLowerCase() === "google" || (model || "").toLowerCase().startsWith("gemini-")) {
      const apiKey = process.env.GEMINI_API_KEY;
      if (!apiKey) return NextResponse.json({ ok: false, message: "Missing GEMINI_API_KEY" } as AugmentResponse, { status: 500 });
      const ai = new GoogleGenAI({ apiKey });
      const res = await ai.models.generateContent({
        model: model || "gemini-2.5-flash",
        contents: [
          { role: "user", parts: [{ text: SYS }] },
          { role: "user", parts: [{ text: userPrompt }] },
        ],
        config: { maxOutputTokens: 4096 },
      } as any);
      text = ((res as any)?.response?.text && (res as any).response.text())
        || (res as any)?.text
        || ((res as any)?.candidates?.[0]?.content?.parts?.map((p: any) => p?.text || "").join("") || "");
      // Fallback: single user message with combined system + prompt
      if (!text || !String(text).trim()) {
        const res2 = await ai.models.generateContent({
          model: model || "gemini-2.5-flash",
          contents: [ { role: 'user', parts: [{ text: `${SYS}\n\n${userPrompt}` }] } ],
          config: { maxOutputTokens: 4096 },
        } as any);
        text = ((res2 as any)?.response?.text && (res2 as any).response.text())
          || (res2 as any)?.text
          || ((res2 as any)?.candidates?.[0]?.content?.parts?.map((p: any) => p?.text || "").join("") || "");
      }
    } else {
      const token = process.env.HF_TOKEN || process.env.DEFAULT_HF_TOKEN;
      const client = new InferenceClient(token);
      const res = await client.chatCompletion({
        model: model || "deepseek-ai/DeepSeek-V3.1",
        provider: (provider as any) || undefined,
        messages: [
          { role: "system", content: SYS },
          { role: "user", content: userPrompt },
        ],
      });
      text = res.choices?.[0]?.message?.content || "";
    }

    const cleaned = cleanText(text);
    if (!cleaned) {
      const title = instruction.match(/Title:\s*(.*)/)?.[1] || "OmniDev App";
      return NextResponse.json({ ok: true, files: generateFallbackFiles(framework, language, title) } as AugmentResponse, { status: 200 });
    }

    let json: any;
    let parseErr: any = null;
    try {
      json = tryParseAny(cleaned);
    } catch (e: any) {
      parseErr = e;
      // Try candidates extracted by balanced scanning
      const cands = findJsonCandidates(cleaned, 12);
      for (const c of cands) {
        try { json = tryParseAny(c); parseErr = null; break; } catch {}
      }
      // Last resort: attempt to coerce by simple replacements
      if (parseErr) {
        try {
          const coerced = cleaned
            .replace(/(\w+)\s*:/g, '"$1":') // quote keys
            .replace(/'([^']*)'/g, '"$1"'); // single to double quotes
          json = tryParseAny(coerced);
          parseErr = null;
        } catch {}
      }
      // REPAIR PASS via Gemini: convert arbitrary text to strict JSON per schema
      if (parseErr) {
        try {
          const apiKey = process.env.GEMINI_API_KEY;
          if (apiKey) {
            const ai = new GoogleGenAI({ apiKey });
            const repairSys = `You are a JSON repair tool. Convert the user's content into STRICT JSON matching either: { ok, files[] } or just files[]. files[] is an array of { path, action, content?, note? }. No prose, no code fences.`;
            const resR = await ai.models.generateContent({
              model: "gemini-2.5-flash",
              contents: [
                { role: 'user', parts: [{ text: repairSys }] },
                { role: 'user', parts: [{ text: cleaned }] },
              ],
              config: { maxOutputTokens: 2048 },
            } as any);
            const repaired = ((resR as any)?.response?.text && (resR as any).response.text())
              || (resR as any)?.text
              || ((resR as any)?.candidates?.[0]?.content?.parts?.map((p: any) => p?.text || "").join("") || "");
            const repairedClean = cleanText(repaired || "");
            if (repairedClean) {
              json = tryParseAny(repairedClean);
              parseErr = null;
            }
          }
        } catch {}
      }
      if (parseErr) {
        // As a last-resort, generate a minimal runnable scaffold so the flow never breaks
        const files = generateFallbackFiles(framework, language, (instruction.match(/Title:\s*(.*)/)?.[1] || "OmniDev App"));
        return NextResponse.json({ ok: true, files } as AugmentResponse, { status: 200 });
      }
    }

    // Minimal validation
    if (json && json.ok && Array.isArray(json.files)) {
      return NextResponse.json(json as AugmentResponse, { status: 200 });
    }
    // If model returned array at root, assume it's files
    if (Array.isArray(json)) {
      return NextResponse.json({ ok: true, files: json } as AugmentResponse, { status: 200 });
    }
    // If the object contains files nested under any key, lift it
    if (json && typeof json === 'object') {
      const stack: any[] = [json];
      while (stack.length) {
        const cur = stack.pop();
        if (cur && typeof cur === 'object') {
          if (Array.isArray(cur.files)) {
            return NextResponse.json({ ok: true, files: cur.files, summary: cur.summary, logs: cur.logs } as AugmentResponse, { status: 200 });
          }
          for (const k of Object.keys(cur)) {
            // eslint-disable-next-line @typescript-eslint/no-explicit-any
            const v: any = (cur as any)[k];
            if (v && typeof v === 'object') stack.push(v);
          }
        }
      }
    }

    // Shape still unexpected: provide runnable fallback
    const files = generateFallbackFiles(framework, language, (instruction.match(/Title:\s*(.*)/)?.[1] || "OmniDev App"));
    return NextResponse.json({ ok: true, files } as AugmentResponse, { status: 200 });
  } catch (e: any) {
    return NextResponse.json({ ok: false, message: e?.message || "Internal error" } as AugmentResponse, { status: 500 });
  }
}

function generateFallbackFiles(framework: string, language: string, title: string): FileUpdate[] {
  const isTs = language.toLowerCase().includes('ts');
  if (framework?.startsWith('express')) {
    const server = `import express from 'express';\nimport cors from 'cors';\nimport path from 'path';\nimport { fileURLToPath } from 'url';\nconst __filename = fileURLToPath(import.meta.url);\nconst __dirname = path.dirname(__filename);\nconst app = express();\napp.use(cors());\nconst distPath = path.join(__dirname, '..', 'frontend', 'dist');\napp.use(express.static(distPath));\napp.get('/health', (_, res) => res.json({ ok: true }));\napp.get('*', (req, res) => { try { res.sendFile(path.join(distPath, 'index.html')); } catch { res.send('OK'); } });\nconst PORT = process.env.PORT || 3000;\napp.listen(PORT, () => console.log('Server listening on', PORT));\n`;
    const indexHtml = `<!DOCTYPE html>\n<html lang=\"en\">\n<head>\n<meta charset=\"UTF-8\" />\n<meta name=\"viewport\" content=\"width=device-width, initial-scale=1.0\" />\n<title>${title}</title>\n<script src=\"https://cdn.tailwindcss.com\"></script>\n</head>\n<body class=\"bg-neutral-950 text-white\">\n<div id=\"root\"></div>\n<script type=\"module\" src=\"/src/main.${isTs ? 'tsx' : 'jsx'}\"></script>\n</body>\n</html>`;
    const main = `import React from 'react';\nimport { createRoot } from 'react-dom/client';\nimport App from './App';\ncreateRoot(document.getElementById('root')).render(<App />);\n`;
    const app = `export default function App(){\n  return (<main className=\"min-h-screen grid place-content-center\">\n    <h1 className=\"text-3xl font-bold\">${title}</h1>\n    <p className=\"text-neutral-400\">Frontend + Backend scaffolded by OmniDev</p>\n  </main>);\n}\n`;
    const fePkg = { name: "frontend", private: true, scripts: { dev: "vite", build: "vite build", preview: "vite preview" }, dependencies: { react: "^18.3.1", "react-dom": "^18.3.1" }, devDependencies: { vite: "^5.4.0" } };
    const bePkg = { name: "backend", private: true, type: "module", scripts: { start: "node server.js" }, dependencies: { express: "^4.19.0", cors: "^2.8.5" } };
    const files: FileUpdate[] = [
      { path: "/backend/server.js", action: "add", content: server },
      { path: "/backend/package.json", action: "add", content: JSON.stringify(bePkg, null, 2) },
      { path: "/frontend/index.html", action: "add", content: indexHtml },
      { path: `/frontend/src/main.${isTs ? 'tsx' : 'jsx'}`, action: "add", content: main },
      { path: `/frontend/src/App.${isTs ? 'tsx' : 'jsx'}`, action: "add", content: app },
      { path: "/frontend/package.json", action: "add", content: JSON.stringify(fePkg, null, 2) },
      { path: "/README.md", action: "add", content: `# ${title}\n\nGenerated by OmniDev scaffold.` },
    ];
    return files;
  }
  if (framework?.startsWith('next')) {
    const page = `export default function Page(){ return (<main style={{minHeight:'100vh',display:'grid',placeContent:'center'}}><h1>${title}</h1><p>Next.js scaffold by OmniDev</p></main>); }`;
    const api = `export async function GET(){ return Response.json({ ok: true }); }`;
    const pkg = { name: "next-app", private: true, scripts: { dev: "next dev", build: "next build", start: "next start" }, dependencies: { next: "^15.0.0", react: "^18.3.1", "react-dom": "^18.3.1" } };
    return [
      { path: "/app/page.tsx", action: "add", content: page },
      { path: "/app/api/health/route.ts", action: "add", content: api },
      { path: "/package.json", action: "add", content: JSON.stringify(pkg, null, 2) },
      { path: "/README.md", action: "add", content: `# ${title}\n\nGenerated by OmniDev scaffold.` },
    ];
  }
  return [ { path: "/README.md", action: "add", content: `# ${title}\n\nOmniDev scaffold.` } ];
}