// [TSM.ID].[11031972] -- All Rights Reserved. Proprietary & Confidential.
require('dotenv').config();
const { WebSocketServer } = require('ws');
const http = require('http');
const jwt = require('jsonwebtoken');
const fs = require('fs');
const path = require('path');
const crypto = require('crypto');
const { generateRegistrationOptions, verifyRegistrationResponse, generateAuthenticationOptions, verifyAuthenticationResponse } = require('@simplewebauthn/server');

const PORT = process.env.PORT || 4001;
const JWT_SECRET = process.env.JWT_SECRET;
if (!JWT_SECRET) { console.error('[3Z VIOLATION] JWT_SECRET env var MUST be set. Exiting.'); process.exit(1); }
const WEBAUTHN_RP_ID = process.env.WEBAUTHN_RP_ID || 'jumpa.id';
const WEBAUTHN_EXPECTED_ORIGIN = process.env.WEBAUTHN_EXPECTED_ORIGIN || 'https://jumpa.id';
const STATE_FILE = path.join(__dirname, 'iam_state.json');

// IAM Omni-Policy State
let iamState = {
  identities: {
    'AEGIS-SUPREME-974EFB44BE-X': {
      role: 'supreme_admin',
      tenant_id: 'SUPREME-MODE',
      packages: ['phase-1-to-75', 'aegis-synthetica', 'ouroboros-protocol']
    },
    'AEGIS-TENANT-80AA3EEBE8-X': {
      role: 'tenant',
      tenant_id: 'JUMPA.ID ENTERPRISE',
      packages: ['JUMPA.ID OMNI-ENGINE']
    }
  },
  policies: {
    'AEGIS-TENANT-80AA3EEBE8-X': {
      name: 'JUMPA.ID ENTERPRISE',
      mfa_mode: 'free', // 'free', 'strict', 'random'
      supreme_allowed: {
        biometric: true,
        optical: true,
        aegis: true
      },
      tenant_enabled: {
        biometric: false,
        optical: false,
        aegis: true
      }
    }
  }
};

const Redis = require('ioredis');

// Connect to Redis for live sync
const redisUrl = process.env.REDIS_URL || 'redis://127.0.0.1:6379';
const redis = new Redis(redisUrl);
redis.subscribe('AEGIS_IAM_STATE_CHANNEL', (err) => {
  if (err) console.error('[IAM GATEKEEPER] Failed to subscribe to Redis:', err);
  else console.log('[IAM GATEKEEPER] Subscribed to Redis channel AEGIS_IAM_STATE_CHANNEL');
});

redis.on('message', (channel, message) => {
  if (channel === 'AEGIS_IAM_STATE_CHANNEL') {
    try {
      const newState = JSON.parse(message);
      iamState = { ...iamState, ...newState };
      console.log('[IAM GATEKEEPER] Synced state from Neural Relay (Redis)');
      // Broadcast to connected WS clients
      const broadcastData = JSON.stringify({ type: 'SYNC_IAM_POLICY', payload: iamState });
      wss.clients.forEach(client => {
        if (client.readyState === 1) client.send(broadcastData);
      });
    } catch (e) {
      console.error('[IAM GATEKEEPER] Error parsing Redis message:', e);
    }
  }
});

const BILLING_MATRIX_URL = process.env.BILLING_MATRIX_URL || 'http://127.0.0.1:8082';

// Hydrate Policy from DuckDB / Billing Matrix
const hydrateState = async () => {
  try {
    const res = await fetch(`${BILLING_MATRIX_URL}/v1/iam/state`);
    const loadedState = await res.json();
    if (loadedState.identities && Object.keys(loadedState.identities).length > 0) {
      iamState.identities = loadedState.identities;
      iamState.policies = { ...iamState.policies, ...loadedState.policies };
      console.log('[IAM GATEKEEPER] Successfully hydrated policies from DuckDB.');
    } else {
      console.log('[IAM GATEKEEPER] DuckDB state empty, using defaults. Seeding DB...');
      saveStateToFile(); // Will actually post to DuckDB
    }
  } catch (err) {
    console.error('[IAM GATEKEEPER] Failed to fetch state from DuckDB.', err.message);
  }
};
hydrateState();

const saveStateToFile = () => {
  fetch(`${BILLING_MATRIX_URL}/v1/iam/state`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(iamState)
  }).catch(err => {
    console.error('[IAM GATEKEEPER] Failed to save state to DuckDB.', err.message);
  });
};


const server = http.createServer((req, res) => {
  res.writeHead(200, { 'Content-Type': 'text/plain', 'Access-Control-Allow-Origin': '*' });
  res.end('XCU IAM Aegis Gatekeeper + Policy Engine + Key Forge is Online.\n');
});

const wss = new WebSocketServer({ server });

let activeAuthSessions = {}; // { sessionId: { targetKeyId, sequence: [], currentIndex: 0, ws: websocket } }

wss.on('connection', (ws) => {
  console.log('[IAM GATEKEEPER] Quantum Client Connected.');
  
  // Send the current policy state and mapped keys immediately
  ws.send(JSON.stringify({
    type: 'SYNC_IAM_POLICY',
    payload: iamState
  }));

  const opticalInterval = setInterval(() => {
    if (ws.readyState === 1) {
      ws.send(JSON.stringify({
        type: 'OPTICAL_CHALLENGE',
        payload: `XCU-OPT-${Math.random().toString(36).substr(2, 9).toUpperCase()}`
      }));
    }
  }, 5000);

  ws.on('message', async (message) => {
    try {
      const data = JSON.parse(message);
      
      // ==========================================
      // 1. POLICY MUTATION HUB (DevSecOps)
      // ==========================================
      if (data.type === 'UPDATE_IAM_POLICY') {
        const { keyId, supreme_allowed, tenant_enabled, mfa_mode } = data.payload;
        if (iamState.policies[keyId]) {
          if (supreme_allowed) iamState.policies[keyId].supreme_allowed = supreme_allowed;
          if (tenant_enabled) iamState.policies[keyId].tenant_enabled = tenant_enabled;
          if (mfa_mode) iamState.policies[keyId].mfa_mode = mfa_mode;
          
          saveStateToFile();
          
          const broadcastData = JSON.stringify({ type: 'SYNC_IAM_POLICY', payload: iamState });
          wss.clients.forEach(client => {
            if (client.readyState === 1) client.send(broadcastData);
          });
        }
      }

      // ==========================================
      // 2. AEGIS KEY FORGE (Zero-Database Credential Swap)
      // ==========================================
      if (data.type === 'REGENERATE_AEGIS_KEY') {
        const { oldKeyId } = data.payload;
        const identity = iamState.identities[oldKeyId];
        
        if (identity) {
          // Create new chaotic string
          const newKeyString = `AEGIS-${identity.role === 'supreme_admin' ? 'SUPREME' : 'TENANT'}-${Math.random().toString(36).substr(2, 10).toUpperCase()}-X`;
          
          // Move identity to new Key
          iamState.identities[newKeyString] = identity;
          delete iamState.identities[oldKeyId];

          // If it's a tenant, move the policy as well
          if (iamState.policies[oldKeyId]) {
            iamState.policies[newKeyString] = iamState.policies[oldKeyId];
            delete iamState.policies[oldKeyId];
          }

          saveStateToFile();
          console.log(`[IAM FORGE] Key Rotated: ${oldKeyId} -> ${newKeyString}`);

          // Send the new key back exclusively to the requester so they can download it
          ws.send(JSON.stringify({
            type: 'NEW_KEYCARD_GENERATED',
            payload: { oldKeyId, newKeyString, role: identity.role }
          }));

          // Broadcast state to all
          const broadcastData = JSON.stringify({ type: 'SYNC_IAM_POLICY', payload: iamState });
          wss.clients.forEach(client => {
            if (client.readyState === 1) client.send(broadcastData);
          });
        }
      }

      // ==========================================
      // 2.5 WEBAUTHN HARDWARE BINDING (FASE 39)
      // ==========================================
      if (data.type === 'GENERATE_REGISTRATION_OPTIONS') {
        const { targetKeyId } = data.payload;
        const identity = iamState.identities[targetKeyId];
        if (identity) {
          const user = {
            id: Buffer.from(targetKeyId).toString('base64url'),
            name: targetKeyId,
            displayName: identity.tenant_id
          };
          
          const options = await generateRegistrationOptions({
            rpName: 'XCU Ultra Gatekeeper',
            rpID: WEBAUTHN_RP_ID,
            userID: user.id,
            userName: user.name,
            attestationType: 'none',
            authenticatorSelection: {
              residentKey: 'required',
              userVerification: 'preferred',
            }
          });
          
          iamState.identities[targetKeyId].currentChallenge = options.challenge;
          ws.send(JSON.stringify({ type: 'REGISTRATION_OPTIONS', payload: options }));
        }
      }

      if (data.type === 'VERIFY_REGISTRATION_RESPONSE') {
        const { targetKeyId, response } = data.payload;
        const identity = iamState.identities[targetKeyId];
        if (identity && identity.currentChallenge) {
          try {
            const verification = await verifyRegistrationResponse({
              response: response,
              expectedChallenge: identity.currentChallenge,
              expectedOrigin: WEBAUTHN_EXPECTED_ORIGIN,
              expectedRPID: WEBAUTHN_RP_ID,
            });
            
            if (verification.verified) {
              const { registrationInfo } = verification;
              iamState.identities[targetKeyId].webAuthn = {
                credentialID: Buffer.from(registrationInfo.credentialID).toString('base64url'),
                credentialPublicKey: Buffer.from(registrationInfo.credentialPublicKey).toString('base64url'),
                counter: registrationInfo.counter,
              };
              delete iamState.identities[targetKeyId].currentChallenge;
              saveStateToFile();
              ws.send(JSON.stringify({ type: 'REGISTRATION_SUCCESS', payload: 'Hardware Biometrik Berhasil Didaftarkan' }));
            }
          } catch (e) {
            ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'WebAuthn Registration Failed: ' + e.message }));
          }
        }
      }

      if (data.type === 'GENERATE_AUTHENTICATION_OPTIONS') {
        const options = await generateAuthenticationOptions({
          rpID: 'localhost',
          userVerification: 'required',
        });
        // Save challenge globally for verification later (since we don't know who is logging in yet)
        iamState.globalAuthChallenge = options.challenge;
        ws.send(JSON.stringify({ type: 'AUTHENTICATION_OPTIONS', payload: options }));
      }

      // ==========================================
      // 2.8 AIR-GAPPED MOBILE PAIRING (FASE 41)
      // ==========================================
      if (data.type === 'GENERATE_OPTICAL_PAIRING') {
        const { targetKeyId } = data.payload;
        if (iamState.identities[targetKeyId]) {
          const secret = crypto.randomBytes(32).toString('hex');
          iamState.identities[targetKeyId].pendingOpticalSecret = secret;
          ws.send(JSON.stringify({ 
            type: 'OPTICAL_PAIRING_QR', 
            payload: `PAIR|${targetKeyId}|${secret}` 
          }));
        }
      }

      if (data.type === 'CONFIRM_OPTICAL_PAIRING') {
        const { targetKeyId, secret } = data.payload;
        const identity = iamState.identities[targetKeyId];
        if (identity && identity.pendingOpticalSecret === secret) {
          identity.opticalSecret = secret;
          delete identity.pendingOpticalSecret;
          saveStateToFile();
          ws.send(JSON.stringify({ type: 'OPTICAL_PAIRING_SUCCESS', payload: 'Perangkat Kuantum Berhasil Dikawinkan.' }));
        } else {
          ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Gagal mengawinkan perangkat.' }));
        }
      }

      // ==========================================
      // 3. AEGIS AUTHENTICATION HUB (MFA ORCHESTRATOR)
      // ==========================================
      if (data.type === 'AEGIS_AUTH_REQUEST') {
        const { authType, keyPayload, sessionId } = data.payload;
        console.log(`[IAM GATEKEEPER] Auth Request: ${authType}`);

        let targetKeyId = null;

        if (authType === 'BIOMETRIC_PULSE') {
          const response = keyPayload; // keyPayload is the credential response from browser
          // Find which identity this credential belongs to
          targetKeyId = Object.keys(iamState.identities).find(k => 
            iamState.identities[k].webAuthn && 
            iamState.identities[k].webAuthn.credentialID === response.id
          );

          if (targetKeyId && iamState.identities[targetKeyId].webAuthn) {
            try {
              const verification = await verifyAuthenticationResponse({
                response: response,
                expectedChallenge: iamState.globalAuthChallenge,
                expectedOrigin: WEBAUTHN_EXPECTED_ORIGIN,
                expectedRPID: WEBAUTHN_RP_ID,
                authenticator: {
                  credentialID: Buffer.from(iamState.identities[targetKeyId].webAuthn.credentialID, 'base64url'),
                  credentialPublicKey: Buffer.from(iamState.identities[targetKeyId].webAuthn.credentialPublicKey, 'base64url'),
                  counter: iamState.identities[targetKeyId].webAuthn.counter,
                }
              });
              
              if (verification.verified) {
                const { authenticationInfo } = verification;
                iamState.identities[targetKeyId].webAuthn.counter = authenticationInfo.newCounter;
                saveStateToFile();
              } else {
                ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Tanda Tangan Biometrik Tidak Valid.' }));
                return;
              }
            } catch (e) {
              ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Gagal Memverifikasi Biometrik: ' + e.message }));
              return;
            }
          } else {
            ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Hardware Biometrik Belum Terdaftar.' }));
            return;
          }
        } else if (authType === 'AEGIS_KEYCARD') {
          targetKeyId = keyPayload;
        } else if (authType === 'OPTICAL_SYNC') {
          // keyPayload carries { targetKeyId, signature, challenge }
          targetKeyId = keyPayload.targetKeyId;
          const identity = iamState.identities[targetKeyId];
          
          if (!identity || !identity.opticalSecret) {
            ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Ponsel ini belum di-Pairing (Dikawinkan) di dalam Dasbor.' }));
            return;
          }

          // Verify HMAC-SHA256 signature
          const expectedSignature = crypto.createHmac('sha256', identity.opticalSecret)
                                          .update(keyPayload.challenge)
                                          .digest('hex');

          if (keyPayload.signature !== expectedSignature) {
            ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Tanda Tangan Kriptografi Ponsel DITOLAK!' }));
            return;
          }
        }

        if (targetKeyId && iamState.identities[targetKeyId]) {
          const identity = iamState.identities[targetKeyId];
          const policy = iamState.policies[targetKeyId];
          
          const typeMap = { 'BIOMETRIC_PULSE': 'biometric', 'OPTICAL_SYNC': 'optical', 'AEGIS_KEYCARD': 'aegis' };
          const methodKey = typeMap[authType];

          // === STRICT POLICY ENFORCEMENT CHECK ===
          if (identity.role === 'tenant' && policy) {
            if (!policy.supreme_allowed[methodKey]) {
              ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: `ACCESS DENIED: Protokol [${methodKey}] telah dicabut secara absolut oleh Supreme Admin.` }));
              return;
            }
            if (!policy.tenant_enabled[methodKey]) {
              ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: `ACCESS DENIED: Anda telah menonaktifkan protokol [${methodKey}] di Security Enclave Anda.` }));
              return;
            }
          }

          // === MFA ORCHESTRATION ===
          if (identity.role === 'tenant' && policy) {
            const mode = policy.mfa_mode || 'free';
            const activeLayers = ['biometric', 'optical', 'aegis'].filter(k => policy.supreme_allowed[k] && policy.tenant_enabled[k]);

            if (mode !== 'free' && activeLayers.length > 1) {
              // Multi-Factor required
              if (!sessionId) {
                // Initialize new session
                const newSessionId = `MFA-SESS-${Math.random().toString(36).substr(2, 9).toUpperCase()}`;
                let sequence = [...activeLayers];
                if (mode === 'random') {
                  sequence = sequence.sort(() => Math.random() - 0.5);
                }
                
                // If the first step in sequence is what they just submitted, advance. Otherwise, reject and tell them the sequence
                if (sequence[0] === methodKey) {
                  activeAuthSessions[newSessionId] = { targetKeyId, sequence, currentIndex: 1, ws };
                  ws.send(JSON.stringify({ 
                    type: 'MFA_STEP_REQUIRED', 
                    payload: { sessionId: newSessionId, nextStep: sequence[1], sequence, currentIndex: 1 } 
                  }));
                  return;
                } else {
                  ws.send(JSON.stringify({ 
                    type: 'MFA_SEQUENCE_STARTED', 
                    payload: { sessionId: newSessionId, sequence, currentIndex: 0, message: `Kebijakan Multi-Lapis Aktif. Anda wajib mematuhi urutan: ${sequence.join(' -> ').toUpperCase()}` } 
                  }));
                  return;
                }
              } else {
                // Advance existing session
                const session = activeAuthSessions[sessionId];
                if (!session || session.targetKeyId !== targetKeyId) {
                  ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Sesi MFA tidak valid atau telah usang.' }));
                  return;
                }
                if (session.sequence[session.currentIndex] !== methodKey) {
                  ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: `Pelanggaran Urutan! Langkah yang diharapkan: ${session.sequence[session.currentIndex].toUpperCase()}` }));
                  delete activeAuthSessions[sessionId];
                  return;
                }
                session.currentIndex++;
                if (session.currentIndex < session.sequence.length) {
                  ws.send(JSON.stringify({ 
                    type: 'MFA_STEP_REQUIRED', 
                    payload: { sessionId, nextStep: session.sequence[session.currentIndex], sequence: session.sequence, currentIndex: session.currentIndex } 
                  }));
                  return;
                }
                // If finished, delete session and grant access
                delete activeAuthSessions[sessionId];
              }
            }
          }

          // Issue JWT if single-layer, free mode, or MFA completed
          console.log('[IAM GATEKEEPER] Auth Successful. Issuing JWT.');
          const token = jwt.sign({
            role: identity.role,
            tenant_id: identity.tenant_id,
            used_gb: Math.random() * 50,
            quota_limit_gb: 500,
            packages: identity.packages
          }, JWT_SECRET, { expiresIn: '24h' });

          ws.send(JSON.stringify({ type: 'AUTH_SUCCESS', payload: { token } }));
        } else {
          console.warn('[IAM GATEKEEPER] Auth Failed. Invalid Key.');
          ws.send(JSON.stringify({ type: 'AUTH_FAILED', payload: 'Kriptografi Gagal Divalidasi. Kunci Fisik ini mungkin sudah Hangus (Revoked).' }));
        }
      }
    } catch (err) {
      console.error('[IAM GATEKEEPER] Payload parsing error:', err);
    }
  });

  ws.on('close', () => {
    clearInterval(opticalInterval);
  });
});

server.listen(PORT, () => {
  console.log(`[IAM GATEKEEPER] Listening on ws://localhost:${PORT}`);
});