app.js

import express from 'express';
import bodyParser from 'body-parser';
import cors from 'cors';
import { StateManager, initDB } from './stateManager.js'; 
import { AIEngine } from './aiEngine.js';
import fs from 'fs';
import crypto from "crypto";
import dotenv from 'dotenv';

dotenv.config();

// Initialize Database
initDB();
const supabase = StateManager.getSupabaseClient(); 

const app = express();
const PORT = process.env.PORT || 7860;

app.use(cors());
app.use(express.json({ limit: '50mb' })); 

// --- BILLING CONSTANTS ---
const MIN_BASIC_REQUIRED = 100; 
const MIN_DIAMOND_REQUIRED = 100; 

// --- STATUS PHASES ---
const WORKER_PHASES = [
    "Worker: Sorting...",
    "Worker: Analyzing Request...",
    "Worker: Reading Context...",
    "Worker: Checking Hierarchy...",
    "Worker: Planning Logic...",
    "Worker: Preparing Environment...",
    "Worker: Thinking..." 
];

const PM_PHASES = [
    "Manager: Sorting...",
    "Manager: Reviewing Request...",
    "Manager: Analyzing Project Structure...",
    "Manager: Consulting Guidelines...",
    "Manager: Formulating Strategy...",
    "Manager: Delegating Tasks...",
    "Manager: Thinking..." 
];

// --- HELPER FUNCTIONS ---
function startStatusLoop(projectId, type = 'worker') {
    const phases = type === 'pm' ? PM_PHASES : WORKER_PHASES;
    let index = 0;
    
    // Set the initial state immediately
    StateManager.setStatus(projectId, phases[0]);

    const interval = setInterval(() => {
        if (index < phases.length - 1) {
            index++;
            StateManager.setStatus(projectId, phases[index]);
        } else {
            // Stop the loop once we reach the last entry
            clearInterval(interval);
        }
    }, 3500); // 3.5 second delay as requested

    // Return a cleanup function in case the component unmounts early
    return () => clearInterval(interval);
};

/*
function startStatusLoop(projectId, type = 'worker') {
    const phases = type === 'pm' ? PM_PHASES : WORKER_PHASES;
    let index = 0;
    
    StateManager.setStatus(projectId, phases[0]);

    const interval = setInterval(() => {
        index = (index + 1) % phases.length; 
        StateManager.setStatus(projectId, phases[index]);
    }, 2000); // Slightly faster updates for better UI feedback

    return () => clearInterval(interval);
} */

function extractWorkerPrompt(text) {
    if (!text || typeof text !== 'string') return null;
    const match = text.match(/WORKER_PROMPT:\s*(.*)/s);
    return match ? match[1].trim() : null;
}

function extractPMQuestion(text) {
    if (!text || typeof text !== 'string') return null;
    const match = text.match(/\[ASK_PM:\s*(.*?)\]/s); 
    return match ? match[1].trim() : null;
}

function extractImagePrompt(text) {
    if (!text || typeof text !== 'string') return null;
    const match = text.match(/\[GENERATE_IMAGE:\s*(.*?)\]/s);
    return match ? match[1].trim() : null;
}

function extractRouteToPM(text) {
    if (!text || typeof text !== 'string') return null;
    const match = text.match(/\[ROUTE_TO_PM:\s*(.*?)\]/s);
    return match ? match[1].trim() : null;
}

function formatContext({ hierarchyContext, scriptContext, logContext }) {
    let out = "";
    if (scriptContext) out += `\n[TARGET SCRIPT]: ${scriptContext.targetName}\n[SOURCE PREVIEW]: ${scriptContext.scriptSource}`;
    if (logContext) out += `\n[LAST LOGS]: ${logContext.logs}`;
    if (hierarchyContext) out += `\n[Hierarchy Context]: ${hierarchyContext}`;
    return out;
}

const validateRequest = (req, res, next) => {
    if (req.path.includes('/admin/cleanup')) return next();
    const { userId } = req.body;
    if (!userId && (req.path.includes('/onboarding') || req.path.includes('/new') || req.path.includes('/project'))) {
        return res.status(400).json({ error: "Missing userId" });
    }
    next();
};

// --- BILLING LOGIC ---

async function checkMinimumCredits(userId, type = 'basic') {
    if (!supabase) return;
    const { data } = await supabase.from('users').select('credits').eq('id', userId).single();
    if (!data) return; 
    
    const credits = data.credits?.[type] || 0;
    const req = type === 'diamond' ? MIN_DIAMOND_REQUIRED : MIN_BASIC_REQUIRED;
    
    if (credits < req) {
        throw new Error(`Insufficient ${type} credits. Required: ${req}, Available: ${credits}`);
    }
}

async function deductUserCredits(userId, amount, type = 'basic') {
    const deduction = Math.floor(parseInt(amount, 10));
    if (!supabase || !deduction || deduction <= 0) return;

    try {
        const { data } = await supabase.from('users').select('credits').eq('id', userId).single();
        if (!data) return;
        
        const currentCredits = data.credits || { basic: 0, diamond: 0 };
        const currentVal = currentCredits[type] || 0;
        const newVal = Math.max(0, currentVal - deduction);
        
        const updatedCredits = { ...currentCredits, [type]: newVal };
        
        await supabase.from('users').update({ credits: updatedCredits }).eq('id', userId);
        console.log(`[Credits] User ${userId} | -${deduction} ${type} | Remaining: ${newVal}`);
    } catch (err) { 
        console.error("Deduction failed", err); 
    }
}

// --- CORE LOGIC: Background Initialization ---

async function runBackgroundInitialization(projectId, userId, description) {
    if (StateManager.isLocked(projectId)) {
        console.log(`[Init Guard] Project ${projectId} is already initializing. Skipping.`);
        return;
    }

    StateManager.lock(projectId);
    console.log(`[Background] Starting initialization for ${projectId}`);
    
    let diamondUsage = 0; 
    let basicUsage = 0;   

    try {
        await StateManager.getProject(projectId);

        await StateManager.updateProject(projectId, { 
            status: "working", 
            gdd: "",
            failureCount: 0 
        });

        const pmHistory = [];
        
        // --- Step 1: Generate GDD (PM) ---
        let stopStatus = startStatusLoop(projectId, 'pm');
        const gddPrompt = `Create a comprehensive Game Design Document (GDD) for: ${description}`;
        const gddResult = await AIEngine.callPM(pmHistory, gddPrompt);
        stopStatus();
        
        if (!gddResult?.text) throw new Error("PM failed to generate GDD");
        
        const gddText = gddResult.text;
        diamondUsage += (gddResult.usage?.totalTokenCount || 0);

        await StateManager.addHistory(projectId, 'pm', 'user', gddPrompt);
        await StateManager.addHistory(projectId, 'pm', 'model', gddText);

        pmHistory.push({ role: 'user', parts: [{ text: gddPrompt }] });
        pmHistory.push({ role: 'model', parts: [{ text: gddText }] });

        // --- Step 2: Generate Tasks (PM) ---
        stopStatus = startStatusLoop(projectId, 'pm');
        const taskPrompt = "Based on the GDD, generate the first technical milestone.\nOutput format:\nTASK_NAME: <Name>\nWORKER_PROMPT: <Specific, isolated instructions for the worker>";
        const taskResult = await AIEngine.callPM(pmHistory, taskPrompt);
        stopStatus();

        if (!taskResult?.text) throw new Error("PM failed to generate Task");

        const taskText = taskResult.text;
        diamondUsage += (taskResult.usage?.totalTokenCount || 0);

        await StateManager.addHistory(projectId, 'pm', 'user', taskPrompt);
        await StateManager.addHistory(projectId, 'pm', 'model', taskText);

        // --- Step 3: Initialize Worker ---
        const initialWorkerInstruction = extractWorkerPrompt(taskText) || `Initialize structure for: ${description}`;
        const initialWorkerPrompt = `CONTEXT: New Project. \nINSTRUCTION: ${initialWorkerInstruction}`;
        
        stopStatus = startStatusLoop(projectId, 'worker');
        
        let workerTextAccumulated = "";

        const workerResult = await AIEngine.callWorkerStream(
            [], 
            initialWorkerPrompt, 
            (thought) => {
                stopStatus(); 
                StateManager.setStatus(projectId, "Worker: Thinking...");
                StateManager.appendSnapshotOnly(projectId, thought);
            },
            (chunk) => {
                stopStatus(); 
                StateManager.setStatus(projectId, "Worker: Coding...");
                workerTextAccumulated += chunk;
                StateManager.appendStream(projectId, chunk);
            }
        );
        stopStatus();
        
        basicUsage += (workerResult.usage?.totalTokenCount || 0);

        if (!workerTextAccumulated && !workerResult.text) throw new Error("Worker failed to initialize");
        const finalWorkerText = workerResult.text || workerTextAccumulated;

        await StateManager.addHistory(projectId, 'worker', 'user', initialWorkerPrompt);
        await StateManager.addHistory(projectId, 'worker', 'model', finalWorkerText);

        await StateManager.updateProject(projectId, {
            gdd: gddText,
            status: "idle" 
        });

        await processAndQueueResponse(projectId, finalWorkerText, userId); 
        console.log(`[Background] Init complete for ${projectId}`);

    } catch (err) {
        console.error(`[Background] Init Error for ${projectId}:`, err.message);
        await StateManager.updateProject(projectId, { status: "error" });
    } finally {
        StateManager.unlock(projectId);
        if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
        if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
    }
}

// --- CORE LOGIC: Async Feedback Loop ---

async function runAsyncFeedback(projectId, userId, fullInput, images = []) {
    let diamondUsage = 0;
    let basicUsage = 0;

    try {
        console.log(`[${projectId}] Feedback received. Persisting state immediately.`);
        
        // 1. Instant Persistence: Saves status and message BEFORE processing
        // This ensures if user reloads, the state is preserved
        await Promise.all([
            StateManager.updateProject(projectId, { status: "working" }),
            StateManager.addHistory(projectId, 'worker', 'user', fullInput)
        ]);

        const project = await StateManager.getProject(projectId);
        
        StateManager.clearSnapshot(projectId);
        let stopStatus = startStatusLoop(projectId, 'worker');
        
        let firstTurnResponse = ""; 
        let thoughtText = "";
        
        // Clone history to prevent mutation issues
        // Note: project.workerHistory includes the message we just saved above
        const currentWorkerHistory = [...(project.workerHistory || [])];

        // 2. First Turn (Worker Execution)
        const workerResult = await AIEngine.callWorkerStream(
            currentWorkerHistory, 
            fullInput, 
            (thought) => {
                stopStatus();
                thoughtText += thought;
                StateManager.setStatus(projectId, "Worker: Thinking...");
                StateManager.appendSnapshotOnly(projectId, thought);
            },
            (chunk) => {
                stopStatus();
                firstTurnResponse += chunk;
                StateManager.setStatus(projectId, "Worker: Coding...");
                StateManager.appendStream(projectId, chunk);
            },
            images 
        );
        stopStatus();
        basicUsage += (workerResult.usage?.totalTokenCount || 0);
        firstTurnResponse = workerResult.text || firstTurnResponse;

        // 3. Analyze Output (Delegation/Questions)
        const routeTask = extractRouteToPM(firstTurnResponse);
        const pmQuestion = extractPMQuestion(firstTurnResponse);
        let finalResponseToSave = firstTurnResponse;

        if (routeTask || pmQuestion) {
            // Save preliminary response to DB so it doesn't get lost
            await StateManager.addHistory(projectId, 'worker', 'model', firstTurnResponse);
            currentWorkerHistory.push({ role: 'model', parts: [{ text: firstTurnResponse }] });

            let pmPrompt = "";
            let pmContextPrefix = "";

            if (routeTask) {
                console.log(`[${projectId}] Routing task to PM...`);
                pmPrompt = `[ROUTED TASK]: ${routeTask}\nWrite the backend/logic. Then use WORKER_PROMPT: <task> to delegate.`;
                pmContextPrefix = "[PM DELEGATION]:";
            } else {
                console.log(`[${projectId}] Worker consulting PM...`);
                pmPrompt = `[WORKER QUESTION]: ${pmQuestion}\nProvide a technical answer or code snippet.`;
                pmContextPrefix = "[PM ANSWER]:";
            }

            // 4. Run PM (Diamond Logic)
            StateManager.clearSnapshot(projectId); 
            stopStatus = startStatusLoop(projectId, 'pm');
            
            let pmResponseText = "";
            
            const pmResult = await AIEngine.callPMStream(
                project.pmHistory,
                pmPrompt,
                (thought) => {
                    stopStatus();
                    StateManager.setStatus(projectId, "Manager: Thinking...");
                    StateManager.appendSnapshotOnly(projectId, thought);
                },
                (chunk) => {
                    stopStatus();
                    StateManager.setStatus(projectId, "Manager: Architecture...");
                    pmResponseText += chunk;
                    StateManager.appendSnapshotOnly(projectId, chunk); 
                }
            );
            stopStatus();
            diamondUsage += (pmResult.usage?.totalTokenCount || 0);
            pmResponseText = pmResult.text || pmResponseText;
            
            await StateManager.addHistory(projectId, 'pm', 'user', pmPrompt);
            await StateManager.addHistory(projectId, 'pm', 'model', pmResponseText);

            // --- EXECUTE PM CODE ---
            // If the PM wrote code (e.g. ServerScriptService logic), execute it immediately
            await processAndQueueResponse(projectId, pmResponseText, userId);

            // 5. Run Worker Continuation
            const nextInstruction = extractWorkerPrompt(pmResponseText) || pmResponseText;
            const workerContinuationPrompt = `${pmContextPrefix} ${nextInstruction}\n\nBased on this, continue the task and output the code.`;

            // Add system prompt to history
            await StateManager.addHistory(projectId, 'worker', 'user', workerContinuationPrompt);
            currentWorkerHistory.push({ role: 'user', parts: [{ text: workerContinuationPrompt }] });

            StateManager.clearSnapshot(projectId);
            stopStatus = startStatusLoop(projectId, 'worker');
            
            let secondTurnResponse = "";
            
            const secondWorkerResult = await AIEngine.callWorkerStream(
                currentWorkerHistory, 
                workerContinuationPrompt,
                (thought) => {
                    stopStatus();
                    StateManager.setStatus(projectId, "Worker: Applying Fix...");
                    StateManager.appendSnapshotOnly(projectId, thought);
                },
                (chunk) => {
                    stopStatus();
                    StateManager.setStatus(projectId, "Worker: Finalizing...");
                    secondTurnResponse += chunk;
                    StateManager.appendStream(projectId, chunk); 
                }
            );
            stopStatus();
            
            // Deduct usage for second turn
            basicUsage += (secondWorkerResult.usage?.totalTokenCount || 0);
            secondTurnResponse = secondWorkerResult.text || secondTurnResponse;

            await StateManager.addHistory(projectId, 'worker', 'model', secondTurnResponse);
            finalResponseToSave = secondTurnResponse;

        } else {
            // No delegation, just save the first response
            await StateManager.addHistory(projectId, 'worker', 'model', firstTurnResponse);
        }

        StateManager.setStatus(projectId, "Idle");
        
        // 6. Final Execution & Cleanup
        await processAndQueueResponse(projectId, finalResponseToSave, userId);
        
        // Ensure status is idle (this also updates timestamps)
        await StateManager.updateProject(projectId, { status: "idle" });

    } catch (err) {
        console.error("Async Feedback Error:", err);
        StateManager.setStatus(projectId, "Error: " + err.message);
        await StateManager.updateProject(projectId, { status: "error" });
    } finally {
        // Billing
        if (diamondUsage > 0) await deductUserCredits(userId, diamondUsage, 'diamond');
        if (basicUsage > 0) await deductUserCredits(userId, basicUsage, 'basic');
    }
}

async function processAndQueueResponse(projectId, rawResponse, userId) {
    if (!rawResponse) return;
    
    // Image Generation
    const imgPrompt = extractImagePrompt(rawResponse);
    if (imgPrompt) {
        try {
            const imgResult = await AIEngine.generateImage(imgPrompt);
            if (imgResult?.image) {
                if (userId) await deductUserCredits(userId, 1000, 'basic'); // Flat fee
                await StateManager.queueCommand(projectId, { type: "CREATE_ASSET", payload: imgResult.image });
            }
        } catch (e) {
            console.error("Image gen failed", e);
        }
    }
    
    // Code Execution
    // Regex looks for lua/luau code blocks to execute
    const codeMatch = rawResponse.match(/```(?:lua|luau)?([\s\S]*?)```/i);
    
    if (codeMatch) {
        await StateManager.queueCommand(projectId, { type: "EXECUTE", payload: codeMatch[1].trim() });
    } else {
        // If there's no code block but response is non-empty, maybe check if it's pure code?
        // For safety, we usually only execute if wrapped in blocks. 
        // But if you want to support raw text execution:
        // await StateManager.queueCommand(projectId, { type: "EXECUTE", payload: rawResponse });
    }
}

// --- EXPRESS ROUTES ---

app.post('/onboarding/analyze', validateRequest, async (req, res) => {
    const { description, userId } = req.body;
    if (!description) return res.status(400).json({ error: "Description required" });

    try {
       // await checkMinimumCredits(userId, 'basic');
       try {
            await checkMinimumCredits(userId, 'basic');
        } catch (creditErr) {
            return res.json({ success: false, insufficient: true, error: creditErr.message });
        }
      
      const result = await AIEngine.generateEntryQuestions(description);
     
        if (result.usage?.totalTokenCount > 0) {
            await deductUserCredits(userId, result.usage.totalTokenCount, 'basic');
        }
      
        if (result.status === "REJECTED") {
            return res.json({ rejected: true, reason: result.reason || "Idea violates TOS." });
        }
      
        res.json({ questions: result.questions });
    } catch (err) {
        res.status(500).json({ error: err.message });
    }
});

app.post('/onboarding/create', validateRequest, async (req, res) => {
    const { userId, description, answers } = req.body;
    
    try {
        
      try {
            await checkMinimumCredits(userId, 'basic');
            await checkMinimumCredits(userId, 'diamond'); 
      } catch (creditErr) {
            return res.json({ success: false, insufficient: true, error: creditErr.message });
        }

        const randomHex = (n = 6) => crypto.randomBytes(Math.ceil(n/2)).toString("hex").slice(0, n);
        const projectId = `proj_${Date.now()}_${randomHex(7)}`;

        const gradingResult = await AIEngine.gradeProject(description, answers);
        
        if (gradingResult.usage?.totalTokenCount) {
             await deductUserCredits(userId, gradingResult.usage.totalTokenCount, 'basic');
        }

        const isFailure = gradingResult.feasibility < 30 || gradingResult.rating === 'F';
        
        if (!isFailure) {
            const { error: insertError } = await supabase.from('projects').insert({
                id: projectId,
                user_id: userId,
                info: {
                    title: gradingResult.title || "Untitled",
                    stats: gradingResult,
                    description,
                    answers,
                    status: "initializing"
                }
            });

            if (insertError) {
                console.error("Failed to insert project:", insertError.message);
                return res.status(500).json({ error: "Database save failed" });
            }

            runBackgroundInitialization(projectId, userId, description);
        }

        res.json({ 
            status: 200,
            success: !isFailure, 
            projectId, 
            stats: gradingResult, 
            title: gradingResult.title || "Untitled"
        });

    } catch (err) {
        console.error("Create Error:", err);
        res.status(500).json({ error: err.message });
    }
});

app.post('/project/feedback', async (req, res) => {
    const { userId, projectId, prompt, hierarchyContext, scriptContext, logContext, images } = req.body;
    
    try {
        const project = await StateManager.getProject(projectId);
        if (!project || project.userId !== userId) return res.status(403).json({ error: "Auth Error" });

        // --- FIXED CREDIT CHECK ---
        // Catch credit errors here to return a soft 200 OK so frontend banner shows
        try {
            await checkMinimumCredits(userId, 'basic');
        } catch (creditErr) {
            return res.json({ success: false, insufficient: true, error: creditErr.message });
        }
        
        const context = formatContext({ hierarchyContext, scriptContext, logContext });
        const fullInput = `USER: ${prompt || "Automatic Feedback"}${context}`;
        
        // Trigger Async
        runAsyncFeedback(projectId, userId, fullInput, images || []);

        res.json({ success: true, message: "Processing started" });

    } catch (err) {
        console.error("Feedback Error:", err);
        await StateManager.updateProject(projectId, { status: "error" });
        res.status(500).json({ error: "Feedback process failed" });
    }
});

/*
app.post('/project/ping', async (req, res) => {
    const { projectId, userId, isFrontend } = req.body;
    if (!projectId || !userId) return res.status(400).json({ error: "Missing IDs" });
    
    const project = await StateManager.getProject(projectId);
    if (!project || project.userId !== userId) return res.json({ action: "IDLE" });

    if (isFrontend) {
        return res.json({ 
            status: StateManager.getStatus(projectId), 
            snapshot: StateManager.getSnapshot(projectId) 
        });
    }

    const command = await StateManager.popCommand(projectId);
    const streamData = StateManager.popStream(projectId);

    let response = { action: "IDLE", stream: streamData || null };

    if (streamData && streamData.length > 0) {
        response.action = "STREAM_APPEND";
    }

    if (command) {
        response.action = command.type; 
        response.target = command.payload;
        response.code = command.type === 'EXECUTE' ? command.payload : null;
    }

    res.json(response);
});
*/

app.post('/project/ping', async (req, res) => {
    const { projectId, userId, isFrontend } = req.body;
    if (!projectId || !userId) return res.status(400).json({ error: "Missing IDs" });
    
    const project = await StateManager.getProject(projectId);
    if (!project || project.userId !== userId) return res.json({ action: "IDLE" });

    // 1. BASE RESPONSE (Visuals)
    // Both Frontend and Plugin need this to show "Thinking..." or text generation
    const response = {
        action: "IDLE",
        status: StateManager.getStatus(projectId),
        snapshot: StateManager.getSnapshot(projectId)
    };

    // 2. EXECUTOR LOGIC (Plugin Only)
    // Only pop commands if this is NOT the passive frontend viewer
    if (!isFrontend) {
        const command = await StateManager.popCommand(projectId);
        
        // We don't need 'streamBuffers' anymore since we use 'snapshot' for visuals.
        // But if you have legacy stream logic, we can leave it or remove it. 
        // For this hybrid model, snapshot is king.

        if (command) {
            response.action = command.type; 
            response.target = command.payload;
            response.code = command.type === 'EXECUTE' ? command.payload : null;
        }
    }

    res.json(response);
});


app.post('/human/override', validateRequest, async (req, res) => {
    const { projectId, instruction, userId } = req.body;
    try {
        await checkMinimumCredits(userId, 'basic');
        const project = await StateManager.getProject(projectId);
        
        res.json({ success: true }); 

        const overrideMsg = `[SYSTEM OVERRIDE]: ${instruction}`;
        
        let overrideText = "";
        let stopStatus = startStatusLoop(projectId, 'worker');

        const workerResult = await AIEngine.callWorkerStream(
            project.workerHistory, 
            overrideMsg,
            (thought) => {
                stopStatus();
                StateManager.setStatus(projectId, "Worker: Thinking...");
                StateManager.appendSnapshotOnly(projectId, thought);
            },
            (chunk) => {
                stopStatus();
                StateManager.setStatus(projectId, "Worker: Coding...");
                overrideText += chunk;
                StateManager.appendStream(projectId, chunk);
            }
        );
        stopStatus();
        
        await StateManager.addHistory(projectId, 'worker', 'user', overrideMsg);
        await StateManager.addHistory(projectId, 'worker', 'model', workerResult.text || overrideText);
        await processAndQueueResponse(projectId, workerResult.text || overrideText, userId);
        
        const usage = workerResult.usage?.totalTokenCount || 0;
        if (usage > 0) await deductUserCredits(userId, usage, 'basic');

    } catch (err) {
        console.error("Override failed", err);
    }
});

app.get('/admin/cleanup', async (req, res) => {
    try {
        const removed = StateManager.cleanupMemory ? StateManager.cleanupMemory() : 0;
        res.json({ success: true, removedCount: removed });
    } catch (err) {
        res.status(500).json({ error: "Cleanup failed" });
    }
});

app.listen(PORT, () => {
    console.log(`AI Backend Running on ${PORT} (Supabase + Billing Edition)`);
});