const express = require('express');
const multer = require('multer');
const path = require('path');
const fs = require('fs');
const { v4: uuidv4 } = require('uuid');
const { PDFDocument } = require('pdf-lib');
const { WebSocketServer } = require('ws');
const http = require('http');

const app = express();
const server = http.createServer(app);
const PORT = process.env.PORT || 3000;

// Trust the reverse proxy (Nginx) so req.ip returns the real user IP
app.set('trust proxy', 1);

// ── Authentication ─────────────────────────────────────────

let currentAppPassword = process.env.APP_PASSWORD || null;
const ADMIN_PASSWORD = process.env.ADMIN_PASSWORD || null;

// Generate unique session tokens every time the server starts
const USER_TOKEN = uuidv4(); 
const ADMIN_TOKEN = uuidv4(); 

function requireAuth(req, res, next) {
  if (!currentAppPassword && !ADMIN_PASSWORD) {
    req.userRole = 'admin';
    return next();
  }

  const token = req.headers['x-auth-token'];
  
  if (ADMIN_PASSWORD && token === ADMIN_TOKEN) {
    req.userRole = 'admin';
    return next();
  }
  
  if (currentAppPassword && token === USER_TOKEN) {
    req.userRole = 'user';
    return next();
  }
  
  res.status(401).json({ error: 'Unauthorized' });
}

function requireAdmin(req, res, next) {
  if (req.userRole === 'admin') return next();
  res.status(403).json({ error: 'Forbidden: Admins only' });
}

// ── Brute Force Protection ─────────────────────────────────

const loginAttempts = new Map();

function checkBruteForce(req, res, next) {
  const ip = req.ip;
  const record = loginAttempts.get(ip);
  if (record && record.lockoutUntil && Date.now() < record.lockoutUntil) {
    const remaining = Math.ceil((record.lockoutUntil - Date.now()) / 1000);
    return res.status(429).json({ error: `Too many attempts. Try again in ${remaining} seconds.` });
  }
  next();
}

function handleFailedLogin(ip) {
  let record = loginAttempts.get(ip) || { attempts: 0, lockoutUntil: 0 };
  record.attempts += 1;
  
  if (record.attempts >= 5) {
    // 5th attempt = 1 min, 6th = 2 mins, 7th = 4 mins, etc.
    const lockMinutes = Math.pow(2, record.attempts - 5);
    record.lockoutUntil = Date.now() + (lockMinutes * 60 * 1000);
  }
  loginAttempts.set(ip, record);
  return record;
}

function handleSuccessfulLogin(ip) {
  loginAttempts.delete(ip);
}

// ── Upload Storage ─────────────────────────────────────────

const UPLOAD_DIR = path.join(__dirname, 'uploads');
if (!fs.existsSync(UPLOAD_DIR)) fs.mkdirSync(UPLOAD_DIR, { recursive: true });

const storage = multer.diskStorage({
  destination: (req, file, cb) => cb(null, UPLOAD_DIR),
  filename: (req, file, cb) => {
    const sessionId = uuidv4();
    const ext = path.extname(file.originalname);
    cb(null, `${sessionId}${ext}`);
  }
});

const upload = multer({
  storage,
  limits: { fileSize: 500 * 1024 * 1024 }, // 500 MB
  fileFilter: (req, file, cb) => {
    if (file.mimetype === 'application/pdf') cb(null, true);
    else cb(new Error('Only PDF files are allowed'));
  }
});

// ── Session Tracking ───────────────────────────────────────
// Map sessionId -> { filePath, originalName, connectedClients, lastAccess }

const sessions = new Map();

function touchSession(sessionId) {
  const s = sessions.get(sessionId);
  if (s) s.lastAccess = Date.now();
}

function cleanupSession(sessionId) {
  const session = sessions.get(sessionId);
  if (!session) return;

  // Only delete if no clients are connected
  if (session.connectedClients > 0) return;

  try {
    if (fs.existsSync(session.filePath)) {
      fs.unlinkSync(session.filePath);
      console.log(`[cleanup] Deleted: ${session.originalName} (${sessionId})`);
    }
  } catch (e) {
    console.warn(`[cleanup] Failed to delete ${sessionId}:`, e.message);
  }
  sessions.delete(sessionId);
}

// Periodic cleanup: delete sessions idle for > 10 minutes
setInterval(() => {
  const now = Date.now();
  for (const [id, session] of sessions) {
    if (session.connectedClients <= 0 && now - session.lastAccess > 10 * 60 * 1000) {
      cleanupSession(id);
    }
  }
}, 60 * 1000); // check every minute

// ── WebSocket for Connection Tracking ──────────────────────

const wss = new WebSocketServer({ server });

wss.on('connection', (ws) => {
  let wsSessionId = null;

  ws.on('message', (msg) => {
    try {
      const data = JSON.parse(msg);
      if (data.type === 'register' && data.sessionId) {
        // Enforce auth on websocket
        if (currentAppPassword || ADMIN_PASSWORD) {
          const isAdmin = ADMIN_PASSWORD && data.token === ADMIN_TOKEN;
          const isUser = currentAppPassword && data.token === USER_TOKEN;
          if (!isAdmin && !isUser) return;
        }

        wsSessionId = data.sessionId;
        const session = sessions.get(wsSessionId);
        if (session) {
          session.connectedClients++;
          touchSession(wsSessionId);
          console.log(`[ws] Client connected to session ${wsSessionId} (${session.connectedClients} clients)`);
        }
      }
    } catch (e) { /* ignore */ }
  });

  ws.on('close', () => {
    if (wsSessionId) {
      const session = sessions.get(wsSessionId);
      if (session) {
        session.connectedClients--;
        console.log(`[ws] Client disconnected from session ${wsSessionId} (${session.connectedClients} remaining)`);
        // Schedule cleanup after a short delay (in case of page refresh)
        setTimeout(() => {
          const s = sessions.get(wsSessionId);
          if (s && s.connectedClients <= 0) {
            cleanupSession(wsSessionId);
          }
        }, 30 * 1000); // 30 second grace period for refreshes
      }
    }
  });
});

// ── Static Files ───────────────────────────────────────────

app.use(express.static(path.join(__dirname, 'public')));
app.use(express.json());

// ── API: Auth ──────────────────────────────────────────────

app.get('/api/check-auth', requireAuth, (req, res) => {
  res.json({ ok: true, role: req.userRole });
});

app.post('/api/login', checkBruteForce, (req, res) => {
  if (!currentAppPassword && !ADMIN_PASSWORD) {
    return res.json({ token: 'no-auth', role: 'admin' });
  }

  const { password } = req.body;
  const ip = req.ip;
  
  if (ADMIN_PASSWORD && password === ADMIN_PASSWORD) {
    handleSuccessfulLogin(ip);
    return res.json({ token: ADMIN_TOKEN, role: 'admin' });
  }
  
  if (currentAppPassword && password === currentAppPassword) {
    handleSuccessfulLogin(ip);
    return res.json({ token: USER_TOKEN, role: 'user' });
  }
  
  const record = handleFailedLogin(ip);
  
  if (record.lockoutUntil > Date.now()) {
    const remaining = Math.ceil((record.lockoutUntil - Date.now()) / 1000);
    return res.status(429).json({ error: `Too many attempts. Locked out for ${remaining} seconds.` });
  }
  
  res.status(401).json({ error: 'Invalid password' });
});

// ── API: Admin ─────────────────────────────────────────────

app.get('/api/admin/stats', requireAuth, requireAdmin, (req, res) => {
  const mem = process.memoryUsage();
  
  const blockedIps = [];
  for (const [ip, record] of loginAttempts) {
    if (record.lockoutUntil > Date.now()) {
      blockedIps.push({ ip, attempts: record.attempts, lockoutRemaining: Math.ceil((record.lockoutUntil - Date.now()) / 1000) });
    }
  }

  res.json({
    uptime: process.uptime(),
    memoryUsedMB: Math.round(mem.rss / 1024 / 1024),
    activeSessions: sessions.size,
    loginAttemptsTracked: loginAttempts.size,
    blockedIps
  });
});

app.post('/api/admin/clear-cache', requireAuth, requireAdmin, (req, res) => {
  let deletedCount = 0;
  for (const [id, session] of sessions) {
    try {
      if (fs.existsSync(session.filePath)) {
        fs.unlinkSync(session.filePath);
        deletedCount++;
      }
    } catch (e) { /* ignore */ }
  }
  sessions.clear();
  res.json({ success: true, deletedCount });
});

app.post('/api/admin/clear-logins', requireAuth, requireAdmin, (req, res) => {
  loginAttempts.clear();
  res.json({ success: true });
});

app.post('/api/admin/change-password', requireAuth, requireAdmin, (req, res) => {
  const { newPassword } = req.body;
  currentAppPassword = newPassword || null;
  res.json({ success: true });
});

// ── API: Upload PDF ────────────────────────────────────────

app.post('/api/upload', requireAuth, upload.single('pdf'), (req, res) => {
  if (!req.file) return res.status(400).json({ error: 'No file uploaded' });

  const sessionId = path.basename(req.file.filename, path.extname(req.file.filename));

  sessions.set(sessionId, {
    filePath: req.file.path,
    originalName: req.file.originalname,
    connectedClients: 0,
    lastAccess: Date.now()
  });

  console.log(`[upload] ${req.file.originalname} -> session ${sessionId}`);

  res.json({
    sessionId,
    originalName: req.file.originalname,
    fileSize: req.file.size
  });
});

// ── API: Serve PDF for client-side rendering ───────────────

app.get('/api/pdf/:sessionId', requireAuth, (req, res) => {
  const session = sessions.get(req.params.sessionId);
  if (!session || !fs.existsSync(session.filePath)) {
    return res.status(404).json({ error: 'Session not found' });
  }
  touchSession(req.params.sessionId);
  res.setHeader('Content-Type', 'application/pdf');
  res.sendFile(session.filePath);
});

// ── API: Split PDF ─────────────────────────────────────────

app.post('/api/split', requireAuth, async (req, res) => {
  const { sessionId, pageNumbers } = req.body;

  const session = sessions.get(sessionId);
  if (!session || !fs.existsSync(session.filePath)) {
    return res.status(404).json({ error: 'Session not found' });
  }
  touchSession(sessionId);

  try {
    const inputBytes = fs.readFileSync(session.filePath);
    const srcDoc = await PDFDocument.load(inputBytes, {
      ignoreEncryption: true,
      updateMetadata: false
    });
    const outDoc = await PDFDocument.create();

    const sorted = [...pageNumbers].sort((a, b) => a - b);
    const indices = sorted.map(p => p - 1).filter(i => i >= 0 && i < srcDoc.getPageCount());

    const copiedPages = await outDoc.copyPages(srcDoc, indices);
    copiedPages.forEach(page => outDoc.addPage(page));

    const outBytes = await outDoc.save();

    const outputName = session.originalName.replace(/\.pdf$/i, '') + '_split.pdf';
    res.setHeader('Content-Type', 'application/pdf');
    res.setHeader('Content-Disposition', `attachment; filename="${outputName}"`);
    res.send(Buffer.from(outBytes));
  } catch (err) {
    console.error('[split] Error:', err);
    res.status(500).json({ error: err.message });
  }
});

// ── Start Server ───────────────────────────────────────────

server.listen(PORT, '0.0.0.0', () => {
  console.log(`PDF Splitter running at http://0.0.0.0:${PORT}`);
});