const express = require('express');
const admin = require('firebase-admin');
const jwt = require('jsonwebtoken');
const { v4: uuidv4 } = require('uuid');
const axios = require('axios');
const bodyParser = require('body-parser');
const cors = require('cors'); 

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

// ---------------------------------------------------------
// 1. STATE MANAGEMENT
// ---------------------------------------------------------
const tempKeys = new Map();
const activeSessions = new Map();

// --- GLOBAL FIREBASE SERVICES ---
let db = null;
let firestore = null; // Added
let storage = null;   // Added

// ---------------------------------------------------------
// 2. FIREBASE INITIALIZATION
// ---------------------------------------------------------
try {
    if (process.env.FIREBASE_SERVICE_ACCOUNT_JSON) {
        const serviceAccount = JSON.parse(process.env.FIREBASE_SERVICE_ACCOUNT_JSON);
        
        // Define your bucket name here (or via ENV). 
        // Based on your previous context, it is likely:
        const bucketName = process.env.FIREBASE_STORAGE_BUCKET;
      
        if (admin.apps.length === 0) {
            admin.initializeApp({
                credential: admin.credential.cert(serviceAccount),
                databaseURL: process.env.FIREBASE_DB_URL,
                storageBucket: bucketName // Required for Storage deletion
            });
        }
        
        // Initialize ALL services
        db = admin.database();
        firestore = admin.firestore();
        storage = admin.storage();

        console.log("🔥 Firebase Connected (RTDB, Firestore, Storage)");
    } else {
        console.warn("⚠️ Memory-Only mode (Firebase credentials missing).");
    }
} catch (e) {
    console.error("Firebase Init Error:", e);
}

// ---------------------------------------------------------
// 3. MIDDLEWARE
// ---------------------------------------------------------
const verifyFirebaseUser = async (req, res, next) => {
    const debugMode = process.env.DEBUG_NO_AUTH === 'true'; 

    if (debugMode) {
        req.user = { uid: "user_dev_01" };
        return next();
    }

    const authHeader = req.headers.authorization;
    if (!authHeader || !authHeader.startsWith('Bearer ')) {
        return res.status(401).json({ error: 'Missing Bearer token' });
    }

    const idToken = authHeader.split('Bearer ')[1];

    try {
        if (admin.apps.length > 0) {
            const decodedToken = await admin.auth().verifyIdToken(idToken);
            req.user = decodedToken;
            next();
        } else {
            req.user = { uid: "memory_user" };
            next();
        }
    } catch (error) {
        return res.status(403).json({ error: 'Unauthorized', details: error.message });
    }
};

async function getSessionSecret(uid, projectId) {
    const cacheKey = `${uid}:${projectId}`;
    if (activeSessions.has(cacheKey)) {
        const session = activeSessions.get(cacheKey);
        session.lastAccessed = Date.now();
        return session.secret;
    }
    if (db) {
        try {
            const snapshot = await db.ref(`plugin_oauth/${uid}/${projectId}`).once('value');
            if (snapshot.exists()) {
                const secret = snapshot.val();
                activeSessions.set(cacheKey, { secret, lastAccessed: Date.now() });
                console.log(`💧 Hydrated secret for ${cacheKey} from DB`);
                return secret;
            }
        } catch (err) {
            console.error("DB Read Error:", err);
        }
    }
    return null;
}

// ---------------------------------------------------------
// 4. ENDPOINTS
// ---------------------------------------------------------

app.post('/key', verifyFirebaseUser, (req, res) => {
    const { projectId } = req.body;
    if (!projectId) return res.status(400).json({ error: 'projectId required' });

    const key = `key_${uuidv4().replace(/-/g, '')}`;
    
    tempKeys.set(key, {
        uid: req.user.uid,
        projectId: projectId,
        createdAt: Date.now()
    });

    console.log(`🔑 Generated Key for user ${req.user.uid}: ${key}`);
    res.json({ key, expiresIn: 300 });
});

app.post('/redeem', async (req, res) => {
    const { key } = req.body;
    
    if (!key || !tempKeys.has(key)) {
        return res.status(404).json({ error: 'Invalid or expired key' });
    }

    const data = tempKeys.get(key);
    const sessionSecret = uuidv4(); 
    
    const token = jwt.sign(
        { uid: data.uid, projectId: data.projectId }, 
        sessionSecret, 
        { expiresIn: '3d' } 
    );

    const cacheKey = `${data.uid}:${data.projectId}`;

    activeSessions.set(cacheKey, { secret: sessionSecret, lastAccessed: Date.now() });

    if (db) {
        await db.ref(`plugin_oauth/${data.uid}/${data.projectId}`).set(sessionSecret);
    }

    tempKeys.delete(key);
    console.log(`🚀 Redeemed JWT for ${cacheKey}`);
    res.json({ token });
});

app.post('/verify', async (req, res) => {
    const { token } = req.body;
    if (!token) return res.status(400).json({ valid: false, error: 'Token required' });

    const decoded = jwt.decode(token);
    if (!decoded || !decoded.uid || !decoded.projectId) {
        return res.status(401).json({ valid: false, error: 'Malformed token' });
    }

    const secret = await getSessionSecret(decoded.uid, decoded.projectId);

    if (!secret) {
        return res.status(401).json({ valid: false, error: 'Session revoked' });
    }

    try {
        jwt.verify(token, secret);
        const threeDaysInSeconds = 3 * 24 * 60 * 60;
        const nowInSeconds = Math.floor(Date.now() / 1000);
        if (decoded.iat && (nowInSeconds - decoded.iat > threeDaysInSeconds)) {
             return res.status(403).json({ valid: false, error: 'Expired' });
        }

        return res.json({ valid: true });
    } catch (err) {
        return res.status(403).json({ valid: false, error: 'Invalid signature' });
    }
});

// ---------------------------------------------------------
// PROXY ENDPOINTS
// ---------------------------------------------------------

app.post('/feedback', async (req, res) => {
    const { token, ...pluginPayload } = req.body;

    if (!token) return res.status(400).json({ error: 'Token required' });

    const decoded = jwt.decode(token);
    if (!decoded || !decoded.uid || !decoded.projectId) {
        return res.status(401).json({ error: 'Malformed token' });
    }

    const secret = await getSessionSecret(decoded.uid, decoded.projectId);
    if (!secret) return res.status(404).json({ error: 'Session revoked' });

    try {
        jwt.verify(token, secret); 

        const externalBase = process.env.EXTERNAL_SERVER_URL || 'http://localhost:7860';
        const targetUrl = externalBase.replace(/\/$/, '') + '/project/feedback';

        console.log(`📨 Forwarding PLUGIN feedback for ${decoded.projectId} (${decoded.uid})`);

        const response = await axios.post(targetUrl, {
            userId: decoded.uid,       
            projectId: decoded.projectId, 
            ...pluginPayload           
        });

        return res.json({ success: true, externalResponse: response.data });

    } catch (err) {
        console.error("Feedback Forward Error:", err.message);
        if (err.response) {
            return res.status(err.response.status).json(err.response.data);
        }
        return res.status(502).json({ error: 'Failed to forward feedback to Main AI server' });
    }
});

app.post('/feedback2', verifyFirebaseUser, async (req, res) => {
    const { projectId, prompt, images, ...otherPayload } = req.body; 
    const userId = req.user.uid;

    if (!projectId || !prompt) {
        return res.status(400).json({ error: 'Missing projectId or prompt' });
    }

    if (images && images.length > 0) {
        console.log(`📸 Received ${images.length} image(s) from Dashboard.`);
    }

    const externalBase = process.env.EXTERNAL_SERVER_URL || 'http://localhost:7860';
    const targetUrl = externalBase.replace(/\/$/, '') + '/project/feedback';

    try {
        const response = await axios.post(targetUrl, {
            userId: userId,
            projectId: projectId,
            prompt: prompt,
            images: images || [], 
            ...otherPayload
        });

        return res.json({ success: true, externalResponse: response.data });
    } catch (err) {
        console.error("Forward Error:", err.message);
        return res.status(502).json({ error: 'Failed to forward' });
    }
});

app.post('/poll', async (req, res) => {
    const { token } = req.body;

    if (!token) return res.status(400).json({ error: 'Token required' });

    const decoded = jwt.decode(token);
    if (!decoded || !decoded.uid || !decoded.projectId) {
        return res.status(401).json({ error: 'Malformed token' });
    }

    const secret = await getSessionSecret(decoded.uid, decoded.projectId);
    if (!secret) return res.status(404).json({ error: 'Session revoked or not found' });

    try {
        const verifiedData = jwt.verify(token, secret);

        const threeDaysInSeconds = 3 * 24 * 60 * 60;
        const nowInSeconds = Math.floor(Date.now() / 1000);
        if (verifiedData.iat && (nowInSeconds - verifiedData.iat > threeDaysInSeconds)) {
             return res.status(403).json({ error: 'Token expired (older than 3 days)' });
        }
        
        const externalBase = process.env.EXTERNAL_SERVER_URL || 'http://localhost:7860';
        const targetUrl = externalBase.replace(/\/$/, '') + '/project/ping'; 
      
        try {
            const response = await axios.post(targetUrl, {
              projectId: verifiedData.projectId,
              userId: verifiedData.uid
            });
          
            return res.json(response.data); 
        } catch (extError) {
            console.error("Poll Forward Error:", extError.message);
            return res.status(502).json({ error: 'External server error' });
        }

    } catch (err) {
        if (err.name === 'TokenExpiredError') {
            return res.status(403).json({ error: 'Token has expired' });
        }
        return res.status(403).json({ error: 'Invalid Token Signature' });
    }
});

// ---------------------------------------------------------
// MANAGEMENT ENDPOINTS
// ---------------------------------------------------------

app.post('/project/delete', verifyFirebaseUser, async (req, res) => {
    const { projectId } = req.body;
    const userId = req.user.uid;

    if (!projectId) return res.status(400).json({ error: "Missing Project ID" });

    console.log(`🗑️ Deleting Project: ${projectId} requested by ${userId}`);

    try {
        // 1. Verify Ownership
        // We check if the project info exists for this user
        const projectRef = db.ref(`projects/${projectId}/info`);
        const snapshot = await projectRef.once('value');
        if (snapshot.exists()) {
            const data = snapshot.val();
            if (data.userId !== userId) {
                return res.status(403).json({ error: "Unauthorized" });
            }
        }

        const promises = [];

        // 2. Delete from Realtime Database
        promises.push(db.ref(`projects/${projectId}`).remove());
        promises.push(db.ref(`plugin_oauth/${userId}/${projectId}`).remove());

        // 3. Delete from Firestore
        if (firestore) {
            promises.push(firestore.collection('projects').doc(projectId).delete());
        } else {
            console.warn("Skipping Firestore delete (not initialized)");
        }

        // 4. Delete from Storage (Recursive)
        if (storage) {
            const bucket = storage.bucket();
            promises.push(bucket.deleteFiles({ prefix: `${projectId}/` }));
        } else {
            console.warn("Skipping Storage delete (not initialized)");
        }

        // 5. Clear from Memory (Manually, since StateManager isn't imported here)
        activeSessions.delete(`${userId}:${projectId}`);
        // Also iterate tempKeys if needed, though they expire quickly anyway
        for (const [key, val] of tempKeys.entries()) {
            if (val.projectId === projectId) tempKeys.delete(key);
        }
        
        await Promise.all(promises);

        console.log(`✅ Project ${projectId} deleted successfully.`);
        res.json({ success: true });

    } catch (err) {
        console.error("Delete Error:", err);
        res.status(500).json({ error: "Failed to delete project resources" });
    }
});

app.get('/cleanup', (req, res) => {
    const THRESHOLD = 1000 * 60 * 60; 
    const now = Date.now();
    let cleanedCount = 0;

    for (const [key, value] of activeSessions.entries()) {
        if (now - value.lastAccessed > THRESHOLD) {
            activeSessions.delete(key);
            cleanedCount++;
        }
    }
    for (const [key, value] of tempKeys.entries()) {
        if (now - value.createdAt > (1000 * 60 * 4)) {
            tempKeys.delete(key);
        }
    }
    res.json({ message: `Cleaned ${cleanedCount} cached sessions from memory.` });
});

app.post('/nullify', verifyFirebaseUser, async (req, res) => {
    const { projectId } = req.body;
    if (!projectId) return res.status(400).json({ error: 'projectId required' });

    const cacheKey = `${req.user.uid}:${projectId}`;
    const existedInMemory = activeSessions.delete(cacheKey);

    let deletedTempKeys = 0;
    for (const [tKey, tData] of tempKeys.entries()) {
        if (tData.uid === req.user.uid && tData.projectId === projectId) {
            tempKeys.delete(tKey);
            deletedTempKeys++;
        }
    }

    if (db) {
        try {
            await db.ref(`plugin_oauth/${req.user.uid}/${projectId}`).remove();
        } catch (e) {
            return res.status(500).json({ error: 'Database error during nullify' });
        }
    }

    console.log(`☢️ NULLIFIED session for ${cacheKey}.`);
    res.json({ 
        success: true, 
        message: 'Session purged.',
        wasCached: existedInMemory,
        tempKeysRemoved: deletedTempKeys
    });
});

app.get('/', (req, res) => {
    res.send('Plugin Auth Proxy Running');
});

const PORT = process.env.PORT || 7860;
app.listen(PORT, () => {
    console.log(`🚀 Auth Proxy running on http://localhost:${PORT}`);
});