import { WebSocketServer, WebSocket } from 'ws';
import * as pty from 'node-pty';
import { exec } from 'child_process';
import util from 'util';
import path from 'path';
import fs from 'fs';

const execAsync = util.promisify(exec);
const wss = new WebSocketServer({ port: 3001 });

let configPaths = {
  flusmPath: '',
  apiPath: '',
  usmTagger: '',
  configFilePath: '',
  usmCmd: '',
  appCmdTemplate: ''
};

// --- STRUCTURES ---
interface ChildState {
  pid: number;
  ppid: number;
  command: string;
  cpu: number;
  memory: number;
  minFlt: number;
  majFlt: number;
  status: 'Running' | 'Dead';
}

interface TerminalSession {
  pid: number;
  ptyProcess: pty.IPty;
  type: 'USM' | 'APP';
  command: string;
  startTime: number;
  childrenHistory: Map<number, ChildState>;
}

const sessions: Map<number, TerminalSession> = new Map();

// --- BUFFERS (ANTI-FREEZE) ---
const pendingLogs = new Map<number, string>();

// Buffer d'agrégation pour les applications
interface AppAggregator {
  pid: number;
  malloc_count: number;
  total_bytes_allocated: number;
  mmap_count: number;
  mmap_bytes_allocated: number;
  munmap_count: number;
  munmap_bytes_freed: number;
  last_allocation_size: number;
  last_action: string;
  free_count: number;
  free_bytes_freed: number;
  madvise_count: number;
  mprotect_count: number;
  mlock_count: number;
  locked_bytes: number; // Pour garder la trace de la quantité de RAM verrouillée !
}
const appStatsBuffer = new Map<number, AppAggregator>();


function broadcast(data: string) {
  wss.clients.forEach((client) => {
    if (client.readyState === WebSocket.OPEN) {
      client.send(data);
    }
  });
}

// --- LECTURE DES STATS KERNEL (/proc/usm_stats) ---
async function getKernelStats() {
    try {
        // Lecture non bloquante du fichier proc
        const data = await fs.promises.readFile('/proc/usm_stats', 'utf8');
        const stats: Record<string, number> = {};
        
        data.split('\n').forEach(line => {
            const parts = line.split(':');
            if (parts.length === 2) {
                stats[parts[0].trim()] = parseInt(parts[1].trim());
            }
        });
        return stats;
    } catch (e) {
        // Le module n'est peut-être pas chargé
        return null;
    }
}

// --- BOUCLE D'ENVOI RAPIDE (FLUSH & STATS) ---
setInterval(async () => {
  // 1. Envoi Logs (Bufferisés)
  if (pendingLogs.size > 0) {
    pendingLogs.forEach((content, pid) => {
      broadcast(JSON.stringify({ type: 'LOG', pid, content }));
    });
    pendingLogs.clear();
  }

  // 2. Envoi Stats Kernel (Toutes les 1s approx, ici on fait rapide mais léger)
  // On peut le faire un peu moins souvent que les logs si on veut, 
  // mais atomic_read est très rapide.
  const kStats = await getKernelStats();
  if (kStats) {
      broadcast(JSON.stringify({ 
          type: 'KERNEL_STATS', 
          timestamp: Date.now(),
          data: kStats 
      }));
  }
}, 100); // 10Hz refresh rate pour l'UI

// --- BOUCLE DE STATS PROCESSUS (Lourde : 1Hz) ---
setInterval(sendProcessStats, 1000);

console.log(`Serveur USM launched`);

// --- HANDLER WEBSOCKET ---
wss.on('connection', (ws: WebSocket, req: any) => {
  
  // --- SÉCURITÉ : CROSS-SITE WEBSOCKET HIJACKING (CSWSH) ---
  const origin = req.headers.origin;
  
  // Liste blanche des origines autorisées (Dev + Production)
  const allowedOrigins = [
    'http://localhost:3000', 
    'http://127.0.0.1:3000',
    'https://usm.imag.fr',    // <-- Obligatoire pour ton site en production !
    'http://usm.imag.fr'
  ];
  
  // Si l'origine est présente et n'est pas dans la liste blanche, on rejette !
  if (origin && !allowedOrigins.includes(origin)) {
    console.warn(`[SECURITY] Connection refused from malicious/unknown origin: ${origin}`);
    ws.close(4003, 'Forbidden Origin');
    return;
  }
  ws.on('message', (message: string) => {
    try {
      const data = JSON.parse(message.toString());

  if (data.action === 'SET_PATHS') {
        configPaths = {
          flusmPath: data.paths.flusmPath,
          apiPath: data.paths.apiPath,
          usmTagger: data.paths.usmTagger,
          configFilePath: data.paths.configFilePath,
          usmCmd: `sudo ${data.paths.apiPath}/project-2 ${data.paths.configFilePath}`,
          appCmdTemplate: data.paths.appCmdTemplate || `sudo -E USM=0 LD_PRELOAD="{{TAGGER}} /lib/x86_64-linux-gnu/libc.so.6" {{COMMAND}}`
        };
        console.log('Configuration updated:', configPaths);
        return;
  }

    if (data.action === 'START_ENGINE') {
      const existing = Array.from(sessions.values()).find(s => s.type === 'USM');
      // Utilisation de configPaths.usmCmd et configPaths.apiPath
      if (!existing && configPaths.usmCmd) {
        spawnTerminal(ws, configPaths.usmCmd, configPaths.apiPath, 'USM');
      }
    }

    if (data.action === 'SPAWN') {
      const userCmd = data.cmd.trim();
      // On construit la commande finale à partir du template modifiable
      let fullCmd = configPaths.appCmdTemplate
                      .replace('{{TAGGER}}', configPaths.usmTagger)
                      .replace('{{COMMAND}}', userCmd);
                      
      // si le client a envoyé un template vide ou cassé, on utilise un fallback
      if (!fullCmd || !fullCmd.includes(userCmd)) {
          fullCmd = `sudo -E USM=0 LD_PRELOAD="${configPaths.usmTagger} /lib/x86_64-linux-gnu/libc.so.6" ${userCmd}`;
      }

      spawnTerminal(ws, fullCmd, configPaths.flusmPath, 'APP');
    }
    
    if (data.action === 'GET_CONFIG') {
      fs.readFile(configPaths.configFilePath, 'utf8', (err, content) => {
          if (!err) ws.send(JSON.stringify({ type: 'CONFIG_CONTENT', content: content }));
      });
    }

    if (data.action === 'SAVE_CONFIG') {
      if (!data.content) return;
      fs.writeFile(configPaths.configFilePath, data.content, 'utf8', (err) => {
        if (!err) {
          ws.send(JSON.stringify({ type: 'LOG', pid: 0, level: 'success', content: `\r\n[CONFIG] Sauvegardé !\r\n` }));
        }
      });
    }
    if (data.action === 'INPUT') {
        const session = sessions.get(data.pid);
        if (session) {
            // Écrit les touches directement dans le processus terminal
            session.ptyProcess.write(data.data);
        }
    }

    if (data.action === 'KILL') {
        const session = sessions.get(data.pid);
        if (session) {
            session.ptyProcess.kill();
            sessions.delete(data.pid);
        }
    }

    } catch (e) { console.error(e); }
  });
});

// --- TERMINAL ---
function spawnTerminal(ws: WebSocket, command: string, cwd: string, type: 'USM' | 'APP') {
  console.log(`[${type}] Spawn attempt: ${command} (Folder: ${cwd})`);  
  try {
      const ptyProcess = pty.spawn('bash', ['-c', command], {
        name: 'xterm-color',
        cols: 80,
        rows: 30,
        cwd: cwd,
        env: process.env as any
      });

      const pid = ptyProcess.pid;
      console.log(`[DEBUG] Process ${type} successfully started with PID: ${pid}`);
      
      sessions.set(pid, { 
          pid, ptyProcess, type, command, 
          startTime: Date.now(),
          childrenHistory: new Map()
      });

      let incomingBuffer = ""; 
      let flushTimeout: NodeJS.Timeout | null = null; // NOUVEAU: Timer pour les lignes coupées

      ptyProcess.onData((chunk) => {
        if (flushTimeout) clearTimeout(flushTimeout); // On annule l'affichage si la suite du texte arrive vite !
        incomingBuffer += chunk;

        let eolIndex;
        while ((eolIndex = incomingBuffer.indexOf('\n')) >= 0) {
            // On extrait la ligne propre
            let line = incomingBuffer.slice(0, eolIndex);
            if (line.endsWith('\r')) line = line.slice(0, -1);
            incomingBuffer = incomingBuffer.slice(eolIndex + 1);

            // REGEX MAGIQUE : On cherche "JSON_STATS:" suivi d'un JSON valide, peu importe les caractères invisibles autour !
            const match = line.match(/JSON_STATS:\s*(\{.*\})/);
            
            if (match) {
                try {
                    const parsed = JSON.parse(match[1]); // On parse uniquement le contenu propre

                    if (parsed.type === 'APP_MALLOC' || parsed.type === 'APP_MMAP' || parsed.type === 'APP_MUNMAP' || parsed.type === 'APP_FREE' || parsed.type === 'APP_MADVISE' || parsed.type === 'APP_MPROTECT' || parsed.type === 'APP_MLOCK') {
                        const parsedPid = parsed.pid;
                        const size = parsed.size || 0;
                        
                        let stats = appStatsBuffer.get(parsedPid);
                        if (!stats) {
                            stats = {
                                pid: parsedPid, malloc_count: 0, total_bytes_allocated: 0,
                                mmap_count: 0, mmap_bytes_allocated: 0,
                                munmap_count: 0, munmap_bytes_freed: 0,
                                free_count: 0, free_bytes_freed: 0,
                                madvise_count: 0, mprotect_count: 0, mlock_count: 0, locked_bytes: 0,
                                last_allocation_size: 0, last_action: '-'
                            };
                            appStatsBuffer.set(parsedPid, stats);
                        }

                        if (parsed.type === 'APP_MALLOC') {
                            stats.malloc_count++;
                            stats.total_bytes_allocated += size;
                        } else if (parsed.type === 'APP_MMAP') {
                            stats.mmap_count++;
                            stats.mmap_bytes_allocated += size;
                        } else if (parsed.type === 'APP_MUNMAP') {
                            stats.munmap_count++;
                            stats.munmap_bytes_freed += size;
                        } else if (parsed.type === 'APP_FREE') { 
                            stats.free_count++;
                            stats.free_bytes_freed += size;
                        } else if (parsed.type === 'APP_MADVISE') { 
                            stats.madvise_count++;
                        } else if (parsed.type === 'APP_MPROTECT') { 
                            stats.mprotect_count++;
                        } else if (parsed.type === 'APP_MLOCK') { 
                            stats.mlock_count++;
                            stats.locked_bytes += size;
                        }
                        
                        stats.last_allocation_size = size;
                        stats.last_action = parsed.type.replace('APP_', '');
                    } else {
                        broadcast(JSON.stringify({ type: 'APP_TELEMETRY', data: parsed }));
                    }
                } catch (e) {
                    // SILENCE : Si on a matché mais que le JSON est profondément corrompu, on l'ignore silencieusement.
                }
            } else {
                // Si ce n'est PAS un log USM, on l'affiche dans le terminal du site
                const currentLog = pendingLogs.get(pid) || "";
                pendingLogs.set(pid, currentLog + line + "\n");
            }
        }
        
        // GESTION INTELLIGENTE DES LIGNES INCOMPLÈTES (Le fameux "Debouncing")
        if (incomingBuffer.length > 0) {
            flushTimeout = setTimeout(() => {
                const currentLog = pendingLogs.get(pid) || "";
                pendingLogs.set(pid, currentLog + incomingBuffer);
                incomingBuffer = ""; 
            }, 50); // On attend 50ms. Si c'est un JSON coupé en deux, la suite arrivera bien avant 50ms !
        }
      });

      ptyProcess.onExit(({ exitCode, signal }) => {
        // ALERTE SI LE PROCESSUS CRASH
        console.log(`[ALERT] Process ${pid} (${type}) stopped on its own! Code=${exitCode}, Signal=${signal}`);
        const s = sessions.get(pid);
        if(s) {
            s.childrenHistory.forEach(c => c.status = 'Dead');
            sessions.delete(pid);
        }
        broadcast(JSON.stringify({ type: 'EXIT', pid, code: exitCode }));
        sendProcessStats();
      });

      sendProcessStats();
  } catch (err) {
    console.error(`[FATAL ERROR] Unable to launch process ${type}:`, err);  }
}
// --- HELPER : ARBRE DES PIDS ---
async function getAllDescendants(rootPid: number): Promise<number[]> {
    const allPids = new Set<number>();
    const queue = [rootPid];
  
    while (queue.length > 0) {
      const currentPid = queue.shift()!;
      if (allPids.has(currentPid)) continue;
      allPids.add(currentPid);
  
      try {
        const { stdout } = await execAsync(`ps --ppid ${currentPid} -o pid --no-headers`);
        const children = stdout.trim().split(/\s+/).map(p => parseInt(p)).filter(p => !isNaN(p));
        children.forEach(child => queue.push(child));
      } catch (e) { /* Pas d'enfants */ }
    }
    return Array.from(allPids);
  }
  
  // --- HELPER : MISE À JOUR DONNÉES PS ---
  async function updateSessionTree(session: TerminalSession) {
    try {
      const familyPids = await getAllDescendants(session.pid);
      if (familyPids.length === 0) return;
  
      const pidListString = familyPids.join(',');
  
      // PS avec les colonnes MIN_FLT et MAJ_FLT
      const { stdout } = await execAsync(`ps -L -p ${pidListString} -o lwp,ppid,pcpu,rss,min_flt,maj_flt,args --no-headers`);
      
      const lines = stdout.trim().split('\n');
      const currentPids = new Set<number>();
  
      lines.forEach(line => {
        const parts = line.trim().split(/\s+/);
        
        if (parts.length >= 7) {
          const lwp = parseInt(parts[0]); 
          const ppid = parseInt(parts[1]);
          const cpu = parseFloat(parts[2]);
          const mem = parseInt(parts[3]) / 1024;
          const minFlt = parseInt(parts[4]);
          const majFlt = parseInt(parts[5]);
          const cmd = parts.slice(6).join(' '); 
  
          if (lwp === session.pid) return; // Ignorer le wrapper bash
  
          currentPids.add(lwp);
          
          session.childrenHistory.set(lwp, {
              pid: lwp, ppid, command: cmd, 
              cpu, memory: mem,
              minFlt, majFlt,
              status: 'Running'
          });
        }
      });
  
      // Marquer les absents comme Dead
      session.childrenHistory.forEach((child, pid) => {
          if (!currentPids.has(pid)) {
              child.status = 'Dead';
              child.cpu = 0; child.memory = 0;
          }
      });
  
    } catch (error) { /* Erreur silencieuse */ }
  }
  
  // --- FONCTION PRINCIPALE DE STATS (PROCESS ONLY) ---
  async function sendProcessStats() {
    const statsList = [];
    // Note: On garde les détails (pmap/fd) si besoin, mais simplifié ici pour l'exemple
    // Si vous avez besoin de pmap/files, réintégrez les helpers getProcessMemoryMap/getOpenFiles ici
  
    for (const session of sessions.values()) {
      await updateSessionTree(session);
  
      let totalCpu = 0; let totalMem = 0; let totalMinFlt = 0; let totalMajFlt = 0;
      const childrenArray: any[] = [];
      
      session.childrenHistory.forEach(child => {
          if (child.status === 'Running') {
               totalCpu += child.cpu; totalMem += child.memory;
          }
          totalMinFlt += child.minFlt || 0; 
          totalMajFlt += child.majFlt || 0;
          childrenArray.push(child);
      });
  
      statsList.push({
        pid: session.pid, name: session.command, type: session.type, status: 'Running',
        cpu: totalCpu.toFixed(1), memory: totalMem.toFixed(1),
        minFlt: totalMinFlt, majFlt: totalMajFlt,
        children: childrenArray.sort((a, b) => a.pid - b.pid)
      });
    }
    
    // Envoi des paquets
    broadcast(JSON.stringify({ type: 'STATS', data: statsList }));
    appStatsBuffer.forEach((stats, pid) => {
      broadcast(JSON.stringify({
          type: 'APP_TELEMETRY',
          data: {
              type: 'APP_AGGREGATED',
              ...stats
          }
      }));
    });

  }