// server.js
// ESM Node (package.json should contain "type": "module")
// Install: npm install express body-parser cors @google/genai
// Run: GEMINI_API_KEY=your_key node server.js

import express from "express";
import bodyParser from "body-parser";
import cors from "cors";
import { GoogleGenAI } from "@google/genai";

const PORT = process.env.PORT ? Number(process.env.PORT) : 7860;
if (!process.env.GEMINI_API_KEY) {
  console.error("GEMINI_API_KEY environment variable is required.");
  process.exit(1);
}

const ai = new GoogleGenAI({ apiKey: process.env.GEMINI_API_KEY });

const app = express();
app.use(cors());
app.use(bodyParser.json({ limit: "2mb" }));

const projects = new Map();
const simpleId = (prefix = "") => prefix + Math.random().toString(36).slice(2, 9);

/**
 * STRICT VOCABULARY - Matched to Plugin
 */
const VALID_ACTIONS = [
  "create_part",
  "create_model",
  "create_meshpart",
  "create_script",
  "attach_script",
  "edit_script",
  "set_property",
  "move_to",
  "rotate_to",
  "set_parent",
  "delete",
  "weld",
  "list_workspace",
  "finish"
];

const ACTION_SYNONYMS = {
  "create_complex_mesh": "create_meshpart",
  "load_mesh": "create_meshpart",
  "move": "move_to",
  "position": "move_to",
  "rotate": "rotate_to",
  "set_rotation": "rotate_to",
  "add_script": "create_script",
  "insert_script": "create_script",
  "modify_script": "edit_script",
  "append_script": "attach_script",
  "group": "create_model", 
  "parent": "set_parent",
  "set_parent_to": "set_parent",
  "join": "weld", 
  "union": "create_model", 
  "create_union": "create_model",
  "create_rig": "create_model", 
  "add_bone": "create_part" 
};

function validateAndPrepareTasks(arr) {
  if (!Array.isArray(arr)) throw new Error("Model output must be a JSON array of tasks.");
  return arr.map((t) => {
    if (!t || typeof t !== "object") throw new Error("Each task must be an object.");
    if (!t.action) throw new Error("Each task must have an action field.");

    let actionRaw = String(t.action).trim();
    const actionLower = actionRaw.toLowerCase();

    let canonicalAction = null;
    if (VALID_ACTIONS.includes(actionRaw)) canonicalAction = actionRaw;
    if (!canonicalAction) {
      canonicalAction = VALID_ACTIONS.find(a => a.toLowerCase() === actionLower) || null;
    }
    if (!canonicalAction && ACTION_SYNONYMS[actionLower]) {
      canonicalAction = ACTION_SYNONYMS[actionLower];
    }

    if (!canonicalAction) {
      throw new Error(`Invalid action '${actionRaw}'. Supported: ${VALID_ACTIONS.join(", ")}`);
    }

    const id = t.id || simpleId("cmd_");
    const params = t.params && typeof t.params === "object" ? t.params : {};
    return { id, action: canonicalAction, originalAction: actionRaw, params, target_ref: t.target_ref || null };
  });
}

async function generateTasks(project) {
  const historySample = (project.history || []).slice(-20).map(h => ({
    commandId: h.commandId,
    status: h.status,
    message: h.message ? (typeof h.message === "string" ? (h.message.length > 120 ? h.message.slice(0, 117) + "..." : h.message) : null) : null,
    originalAction: h.originalAction || null,
    target_ref: h.target_ref || null,
  }));
  const refsKeys = Object.keys(project._refs || {}).slice(0, 200);

  const stateSummary = {
    projectId: project.id,
    refsSummary: refsKeys,
    recentHistory: historySample
  };

  const systemPrompt = `
You are a precise Roblox Studio Assistant.
Allowed actions: ${VALID_ACTIONS.join(", ")}.

RULES:
1. **References:** When creating an object, ALWAYS assign a unique "target_ref" (e.g., "ref_wall_1").
2. **Chaining:** When moving/modifying later, use "target_ref": "ref_wall_1".
3. **Hierarchy:** To put a Part inside a Model: create_model (ref_grp) -> create_part (ref_part) -> set_parent (target_ref: ref_part, params: { parent_ref: ref_grp }).
4. **Math:** Use numbers. Position {x,y,z}. Rotation {x,y,z} (degrees).
5. **Output:** Return ONLY a JSON Array.

EXISTING STATE:
${JSON.stringify(stateSummary, null, 2)}

User Request:
${JSON.stringify({ description: project.description })}
`;

  // Using 1.5-pro as the closest valid "Pro" model. 
  // If you strictly need "gemini-2.5-pro", allow the string, but note it might 404 if not enabled on your account.
  const resp = await ai.models.generateContent({
    model: "gemini-2.5-pro", 
    contents: systemPrompt,
    generationConfig: { responseMimeType: "application/json" }
  });

  // === FIX: Robust Response Text Extraction ===
  let text = "";
  if (resp && typeof resp.text === 'string') {
     // SDK returned text property directly
     text = resp.text; 
  } else if (resp && typeof resp.text === 'function') {
     // SDK returned text method
     text = resp.text();
  } else if (resp?.response?.text) {
     // Nested response object
     text = typeof resp.response.text === 'function' ? resp.response.text() : resp.response.text;
  } else {
     // Fallback for raw data structure
     text = resp?.data?.candidates?.[0]?.content?.parts?.[0]?.text || "";
  }
  // ============================================

  let arr;
  try {
    // Sanitize markdown if still present
    const cleanText = text.replace(/```json/g, "").replace(/```/g, "").trim();
    arr = JSON.parse(cleanText);
  } catch (err) {
    // Fallback extraction
    const start = text.indexOf("[");
    const end = text.lastIndexOf("]");
    if (start >= 0 && end > start) {
      arr = JSON.parse(text.slice(start, end + 1));
    } else {
      console.error("AI Output was:", text);
      throw new Error("Could not parse JSON from model output.");
    }
  }

  return validateAndPrepareTasks(arr);
}

// --- ENDPOINTS ---

app.post("/projects", async (req, res) => {
  try {
    const projectId = req.body.projectId || simpleId("proj_");
    const project = {
      id: projectId,
      description: req.body.description || "New Project",
      status: "queued",
      tasks: [],
      commandQueue: [],
      pendingResults: new Map(),
      running: false,
      history: [], 
      _refs: {}    
    };

    projects.set(projectId, project);

    runProject(project).catch((err) => {
      project.status = "error";
      console.error("runProject error:", err);
    });

    return res.status(202).json({ projectId, status: "accepted" });
  } catch (err) {
    return res.status(500).json({ error: err.message });
  }
});

app.post("/projects/:projectId/prompt", async (req, res) => {
  const project = projects.get(req.params.projectId);
  if (!project) return res.status(404).json({ error: "project not found" });
  
  const prompt = req.body.prompt;
  project.description += `\nUser Update: ${prompt}`;

  try {
    const newTasks = await generateTasks({ ...project, description: prompt });
    project.tasks.push(...newTasks);
    project.commandQueue.push(...newTasks);
    return res.json({ appended: newTasks.length, tasks: newTasks });
  } catch (err) {
    return res.status(500).json({ error: err.message });
  }
});

app.get("/projects/:projectId/next", (req, res) => {
  const project = projects.get(req.params.projectId);
  if (!project) return res.status(404).json({ error: "project not found" });

  if (project.commandQueue.length > 0) {
    return res.json(project.commandQueue.shift());
  }
  return res.status(204).end();
});

app.post("/projects/:projectId/result", (req, res) => {
  const project = projects.get(req.params.projectId);
  if (!project) return res.status(404).json({ error: "project not found" });

  const { commandId, status, message, extra, target_ref } = req.body;
  
  // History recording
  const entry = {
    commandId,
    status,
    message: message || null,
    extra: extra || null,
    target_ref: target_ref || null,
    ts: Date.now()
  };
  project.history = project.history || [];
  project.history.push(entry);

  if (target_ref && status === "ok") {
    project._refs = project._refs || {};
    project._refs[target_ref] = { commandId, ts: Date.now() };
  }

  const waiter = project.pendingResults.get(commandId);
  if (waiter) {
    project.pendingResults.delete(commandId);
    waiter.resolve({ status, message, target_ref });
    return res.json({ ok: true });
  }

  return res.json({ ok: true });
});

app.get("/status/:projectId", (req, res) => {
  const project = projects.get(req.params.projectId);
  if (!project) return res.status(404).json({ error: "project not found" });
  return res.json({
    id: project.id,
    status: project.status,
    queued: project.commandQueue.length,
    refs: project._refs || {}
  });
});

async function runProject(project) {
  if (project.running) return;
  project.running = true;
  project.status = "running";

  try {
    if (!project.tasks || project.tasks.length === 0) {
      project.tasks = await generateTasks(project);
    }
    project.commandQueue.push(...project.tasks);

    // Blocking wait logic
    for (const task of project.tasks) {
      const result = await new Promise((resolve, reject) => {
        const timeoutMs = 3 * 60 * 1000; 
        let timedOut = false;
        const timeout = setTimeout(() => {
          timedOut = true;
          project.pendingResults.delete(task.id);
          reject(new Error("Command timed out: " + task.id));
        }, timeoutMs);

        project.pendingResults.set(task.id, {
          resolve: (payload) => {
            if (timedOut) return;
            clearTimeout(timeout);
            resolve(payload);
          },
          reject: (err) => {
            if (timedOut) return;
            clearTimeout(timeout);
            reject(err);
          }
        });
      });

      if (result.status !== "ok") {
        project.status = "failed";
        throw new Error(`Command ${task.id} failed: ${result.message}`);
      }
    }
    project.status = "complete";
  } catch (err) {
    console.error("Project run error:", err);
    if (project.status !== "failed") project.status = "error";
  } finally {
    project.running = false;
  }
}

app.listen(PORT, () => {
  console.log(`BloxBuddy backend listening on port ${PORT}`);
});