// Real Merkle Service for License Verification
const express = require('express');
const { Pool } = require('pg');
const cors = require('cors');
const crypto = require('crypto');
const { performance } = require('perf_hooks');
const circomlibjs = require('circomlibjs');

// Initialize Express
const app = express();
app.use(cors());
app.use(express.json());

// Database connection
const db = new Pool({
    connectionString: process.env.DATABASE_URL || 
        'postgresql://license_admin:secure_license_pass_123@postgres:5432/license_verification',
    max: 5
});

// Configuration
const TREE_DEPTH = parseInt(process.env.TREE_DEPTH || '17');
const MAX_LEAVES = Math.pow(2, TREE_DEPTH);

// Poseidon hash instance
let poseidon = null;
let poseidonF = null;

// Current tree in memory
let currentTree = {
    version: 0,
    root: null,
    leaves: [],
    leafMap: new Map(), // licenseHash -> leafIndex
    layers: [],
    isBuilt: false,
    treeId: null
};

// Initialize Poseidon
async function initPoseidon() {
    if (!poseidon) {
        const poseidonJs = await circomlibjs.buildPoseidon();
        poseidon = poseidonJs;
        poseidonF = poseidonJs.F;
        console.log('[Merkle] Poseidon hash initialized');
    }
    return poseidon;
}

// Convert string to field element
function stringToFieldElement(str) {
    let result = BigInt(0);
    for (let i = 0; i < Math.min(str.length, 31); i++) {
        result = result * BigInt(256) + BigInt(str.charCodeAt(i));
    }
    return result.toString();
}

// Hash two elements with Poseidon
function poseidonHash2(left, right) {
    if (!poseidon) throw new Error('Poseidon not initialized');
    return poseidonF.toString(poseidon([left, right]));
}

// Hash license data
function hashLicenseData(licenseData) {
    if (!poseidon) throw new Error('Poseidon not initialized');
    
    const inputs = [
        stringToFieldElement(licenseData.licenseNumber),
        stringToFieldElement(licenseData.practitionerName || 'Anonymous'),
        BigInt(licenseData.issuedDate || 0).toString(),
        BigInt(licenseData.expiryDate || 0).toString(),
        stringToFieldElement(licenseData.jurisdiction || 'Unknown')
    ];
    
    return poseidonF.toString(poseidon(inputs));
}

// Build Merkle tree with Poseidon
function buildMerkleTree(leaves) {
    if (leaves.length === 0) throw new Error('No leaves provided');
    
    // Ensure we have a power of 2 number of leaves
    const targetSize = Math.pow(2, Math.ceil(Math.log2(leaves.length)));
    const paddedLeaves = [...leaves];
    
    // Pad with zeros
    const zeroHash = "0";
    while (paddedLeaves.length < targetSize) {
        paddedLeaves.push(zeroHash);
    }
    
    // Build tree layers
    const layers = [paddedLeaves];
    let currentLayer = paddedLeaves;
    
    while (currentLayer.length > 1) {
        const nextLayer = [];
        for (let i = 0; i < currentLayer.length; i += 2) {
            const left = currentLayer[i];
            const right = currentLayer[i + 1] || left;
            nextLayer.push(poseidonHash2(left, right));
        }
        layers.push(nextLayer);
        currentLayer = nextLayer;
    }
    
    return {
        root: currentLayer[0],
        layers: layers,
        depth: layers.length - 1
    };
}

// Generate Merkle proof for a leaf
function generateMerkleProof(tree, leafIndex) {
    if (!tree || !tree.layers) throw new Error('Invalid tree');
    if (leafIndex >= tree.layers[0].length) throw new Error('Leaf index out of bounds');
    
    const pathElements = [];
    const pathIndices = [];
    
    let currentIndex = leafIndex;
    
    for (let level = 0; level < tree.depth; level++) {
        const isRightNode = currentIndex % 2 === 1;
        const siblingIndex = isRightNode ? currentIndex - 1 : currentIndex + 1;
        
        if (siblingIndex < tree.layers[level].length) {
            pathElements.push(tree.layers[level][siblingIndex]);
            pathIndices.push(isRightNode ? "1" : "0");
        } else {
            // Sibling doesn't exist, use the same node (shouldn't happen with proper padding)
            pathElements.push(tree.layers[level][currentIndex]);
            pathIndices.push("0");
        }
        
        currentIndex = Math.floor(currentIndex / 2);
    }
    
    return { pathElements, pathIndices };
}

// Build tree from database licenses
// Build tree from database licenses (with auto-population)
async function buildTreeFromDatabase() {
    const startTime = performance.now();
    console.log('[Merkle] Building tree from database...', currentTree);
    if(currentTree.isBuilt || currentTree.isBuilding){
        return null
    }
    currentTree.isBuilding = true;
    
    try {
        // Get active licenses
        let result = await db.query(`
            SELECT 
                id, 
                license_number,
                practitioner_name,
                EXTRACT(EPOCH FROM issued_date)::INTEGER as issued_date,
                EXTRACT(EPOCH FROM expiry_date)::INTEGER as expiry_date,
                jurisdiction
            FROM licenses 
            WHERE status = 'active' 
            ORDER BY id
            LIMIT $1
        `, [MAX_LEAVES]);
        
        let licenses = result.rows;
        console.log(`[Merkle] Found ${licenses.length} active licenses`);
        
        // Auto-generate licenses if none exist
        if (licenses.length === 0) {
            console.log('[Merkle] No licenses found. Auto-generating 100,000 licenses...');
            
            const BATCH_SIZE = 1000;
            const TOTAL_LICENSES = 100000;
            const jurisdictions = ['CA', 'NY', 'TX', 'FL', 'IL', 'PA', 'OH', 'GA', 'NC', 'MI'];
            
            await db.query('BEGIN');
            
            try {
                for (let batch = 0; batch < TOTAL_LICENSES / BATCH_SIZE; batch++) {
                    const values = [];
                    const placeholders = [];
                    
                    for (let i = 0; i < BATCH_SIZE; i++) {
                        let licenseIndex = batch * BATCH_SIZE + i + 1;
                        let licenseNumber = `LIC-${String(licenseIndex).padStart(8, '0')}`;
                        let practitionerName = `Dr. ${generateRandomName()} ${generateRandomSurname()}`;
                        let issuedDate = new Date(Date.now() - Math.random() * 365 * 24 * 60 * 60 * 1000); // Random date in past year
                        let expiryDate = new Date(Date.now() + (365 + Math.random() * 730) * 24 * 60 * 60 * 1000); // 1-3 years from now
                        let jurisdiction = jurisdictions[Math.floor(Math.random() * jurisdictions.length)];
                        let status = Math.random() > 0.02 ? 'active' : 'suspended'; // 98% active
                        if(i === 0){
                            status = 'active';
                            jurisdiction = jurisdictions[0];
                            practitionerName = 'Test Practitioner 1';
                            issuedDate = new Date('2025-01-23');
                            expiryDate = new Date('2028-01-23');
                        }
                        // Calculate hash using SHA256 
                        const hash = crypto.createHash('sha256').update(licenseNumber).digest();
                        
                        values.push(
                            licenseNumber,
                            hash,
                            practitionerName,
                            issuedDate,
                            expiryDate,
                            status,
                            jurisdiction
                        );
                        
                        const offset = i * 7;
                        placeholders.push(
                            `($${offset + 1}, $${offset + 2}, $${offset + 3}, $${offset + 4}, $${offset + 5}, $${offset + 6}, $${offset + 7})`
                        );
                    }
                    
                    const insertQuery = `
                        INSERT INTO licenses (
                            license_number, 
                            license_hash, 
                            practitioner_name, 
                            issued_date, 
                            expiry_date, 
                            status, 
                            jurisdiction
                        ) VALUES ${placeholders.join(', ')}
                        ON CONFLICT (license_number) DO NOTHING
                    `;
                    
                    await db.query(insertQuery, values);
                    
                    if ((batch + 1) % 10 === 0) {
                        console.log(`[Merkle] Inserted ${(batch + 1) * BATCH_SIZE} licenses...`);
                    }
                }
                
                await db.query('COMMIT');
                console.log(`[Merkle] Successfully generated ${TOTAL_LICENSES} licenses`);
                
                // Rerun the query to get the newly inserted licenses
                result = await db.query(`
                    SELECT 
                        id, 
                        license_number,
                        practitioner_name,
                        EXTRACT(EPOCH FROM issued_date)::INTEGER as issued_date,
                        EXTRACT(EPOCH FROM expiry_date)::INTEGER as expiry_date,
                        jurisdiction
                    FROM licenses 
                    WHERE status = 'active' 
                    ORDER BY id
                    LIMIT $1
                `, [MAX_LEAVES]);
                
                licenses = result.rows;
                console.log(`[Merkle] Now have ${licenses.length} active licenses`);
                
            } catch (insertError) {
                await db.query('ROLLBACK');
                console.error('[Merkle] Failed to generate licenses:', insertError);
                throw new Error('Failed to auto-generate licenses: ' + insertError.message);
            }
        }
        
        // Hash each license to create leaves
        const leaves = [];
        const leafMap = new Map();
        const leafHashToIndex = new Map();
        
        for (let i = 0; i < licenses.length; i++) {
            const license = licenses[i];
            const leaf = hashLicenseData({
                licenseNumber: license.license_number,
                practitionerName: license.practitioner_name,
                issuedDate: license.issued_date,
                expiryDate: license.expiry_date,
                jurisdiction: license.jurisdiction
            });
            
            leaves.push(leaf);
            leafHashToIndex.set(leaf, i);
            leafMap.set(license.license_number, i);
            
            // Also store by simple hash for testing
            const simpleHash = crypto.createHash('sha256')
                .update(license.license_number)
                .digest('hex');
            leafMap.set(simpleHash, i);
        }
        
        // Build the Merkle tree
        const tree = buildMerkleTree(leaves);
        let treeId;
        
        // Store in database
        await db.query('BEGIN');
        
        // Lock the table to prevent concurrent modifications
        await db.query('LOCK TABLE merkle_trees IN EXCLUSIVE MODE');
        
        // Deactivate all currently active trees
        const deactivateResult = await db.query(
            'UPDATE merkle_trees SET is_active = false WHERE is_active = true RETURNING id'
        );
        
        if (deactivateResult.rows.length > 0) {
            console.log(`[Merkle] Deactivated ${deactivateResult.rows.length} previous tree(s)`);
        }
        let treeRoot = Buffer.from(tree.root.padStart(64, '0'), 'hex');
        // Insert the new active tree
        const insertResult = await db.query(`
            INSERT INTO merkle_trees (
                tree_version, 
                root_hash, 
                tree_depth, 
                leaf_count, 
                is_active, 
                finalized_at
            )
            VALUES ($1, $2, $3, $4, true, NOW())
            RETURNING *
        `, [
            currentTree.version + 1,
            tree.root,
            tree.depth,
            licenses.length
        ]);
        
        treeId = insertResult.rows[0].id;
        
        // Store leaf mappings (limit for performance)
        const leafBatchSize = Math.min(licenses.length, 10000);
        console.log(`[Merkle] Storing ${leafBatchSize} leaf mappings...`);
        
        for (let i = 0; i < leafBatchSize; i += 100) {
            const batch = [];
            for (let j = i; j < Math.min(i + 100, leafBatchSize); j++) {
                batch.push([
                    treeId,
                    licenses[j].id,
                    j,
                    leaves[j] 
                ]);
            }
            
            if (batch.length > 0) {
                const placeholders = batch.map((_, idx) => 
                    `($${idx * 4 + 1}, $${idx * 4 + 2}, $${idx * 4 + 3}, $${idx * 4 + 4})`
                ).join(', ');
                
                const flatValues = batch.flat();
                
                await db.query(`
                    INSERT INTO merkle_leaves (tree_id, license_id, leaf_index, leaf_hash)
                    VALUES ${placeholders}
                    ON CONFLICT DO NOTHING
                `, flatValues);
            }
            
            // Log progress for large insertions
            if ((i + 100) % 1000 === 0) {
                console.log(`[Merkle] Stored ${Math.min(i + 100, leafBatchSize)} leaf mappings...`);
            }
        }
        
        // Commit the transaction
        await db.query('COMMIT');
        
        // Update in-memory tree
        currentTree = {
            version: currentTree.version + 1,
            root: tree.root,
            leaves: leaves,
            leafMap: leafMap,
            leafHashToIndex: leafHashToIndex,
            layers: tree.layers,
            isBuilt: true,
            treeId: treeId,
            depth: tree.depth,
            leafCount: licenses.length,
            isBuilding: false,
        };
        
        const buildTime = performance.now() - startTime;
        console.log(`[Merkle] Tree built in ${buildTime.toFixed(0)}ms`);
        console.log(`[Merkle] Root: ${tree.root}`);
        console.log(`[Merkle] Total leaves: ${licenses.length}`);
        
        return {
            treeId,
            root: tree.root,
            leafCount: licenses.length,
            depth: tree.depth,
            buildTimeMs: buildTime
        };
        
    } catch (error) {
        await db.query('ROLLBACK');
        console.error('[Merkle] Build failed:', error);
        throw error;
    }
}

// Helper functions for generating random names
function generateRandomName() {
    const firstNames = [
        'James', 'Mary', 'John', 'Patricia', 'Robert', 'Jennifer', 'Michael', 'Linda',
        'William', 'Elizabeth', 'David', 'Barbara', 'Richard', 'Susan', 'Joseph', 'Jessica',
        'Thomas', 'Sarah', 'Charles', 'Karen', 'Christopher', 'Nancy', 'Daniel', 'Lisa',
        'Matthew', 'Betty', 'Anthony', 'Helen', 'Mark', 'Sandra', 'Donald', 'Donna',
        'Steven', 'Carol', 'Bill', 'Ruth', 'Paul', 'Sharon', 'Joshua', 'Michelle',
        'Kenneth', 'Laura', 'Kevin', 'Sarah', 'Brian', 'Kimberly', 'George', 'Deborah'
    ];
    return firstNames[Math.floor(Math.random() * firstNames.length)];
}

function generateRandomSurname() {
    const surnames = [
        'Smith', 'Johnson', 'Williams', 'Brown', 'Jones', 'Garcia', 'Miller', 'Davis',
        'Rodriguez', 'Martinez', 'Hernandez', 'Lopez', 'Gonzalez', 'Wilson', 'Anderson',
        'Thomas', 'Taylor', 'Moore', 'Jackson', 'Martin', 'Lee', 'Perez', 'Thompson',
        'White', 'Harris', 'Sanchez', 'Clark', 'Ramirez', 'Lewis', 'Robinson', 'Walker',
        'Young', 'Allen', 'King', 'Wright', 'Scott', 'Torres', 'Nguyen', 'Hill',
        'Flores', 'Green', 'Adams', 'Nelson', 'Baker', 'Hall', 'Rivera', 'Campbell',
        'Mitchell', 'Carter', 'Roberts', 'Gomez', 'Phillips', 'Evans', 'Turner', 'Diaz'
    ];
    return surnames[Math.floor(Math.random() * surnames.length)];
}

// Generate mock tree for testing
async function buildMockTree() {
    console.log('[Merkle] Building mock tree...');
    
    const mockLicenses = [];
    const leafMap = new Map();
    const leaves = [];
    
    // Generate mock licenses
    for (let i = 0; i < 1000; i++) {
        const licenseNumber = `LIC-${String(i + 1).padStart(8, '0')}`;
        const licenseData = {
            licenseNumber: licenseNumber,
            practitionerName: `Test Practitioner ${i + 1}`,
            issuedDate: Math.floor(Date.now() / 1000) - 31536000, // 1 year ago
            expiryDate: Math.floor(Date.now() / 1000) + 31536000 * 2, // 2 years from now
            jurisdiction: ['CA', 'NY', 'TX', 'FL'][i % 4]
        };
        
        const leaf = hashLicenseData(licenseData);
        leaves.push(leaf);
        leafMap.set(licenseNumber, i);
        
        // Also map simple hash
        const simpleHash = crypto.createHash('sha256').update(licenseNumber).digest('hex');
        leafMap.set(simpleHash, i);
    }
    
    const tree = buildMerkleTree(leaves);
    
    currentTree = {
        version: currentTree.version + 1,
        root: tree.root,
        leaves: leaves,
        leafMap: leafMap,
        layers: tree.layers,
        isBuilt: true,
        treeId: 'mock-' + Date.now(),
        depth: tree.depth,
        leafCount: leaves.length
    };
    
    console.log(`[Merkle] Mock tree built with ${leaves.length} leaves`);
    console.log(`[Merkle] Root: ${tree.root}`);
    
    return currentTree;
}

app.get('/api/merkle-proof-by-hash/:leafHash', async (req, res) => {
    const startTime = performance.now();
    const { leafHash } = req.params;
    
    console.log(`[Merkle] ===== PRIVACY-PRESERVING PROOF REQUEST =====`);
    console.log(`[Merkle] Leaf hash: ${leafHash.substring(0, 20)}...`);
    console.log(`[Merkle] Server does NOT see license details!`);
    
    try {
        // Ensure tree is built
        if (!currentTree.isBuilt) {
            await buildTreeFromDatabase();
        }
        
        // Find this leaf hash in the tree
        let leafIndex = currentTree.leaves.findIndex(leaf => leaf === leafHash);
        
        if (leafIndex === -1) {
            console.log(`[Merkle] Generating proof anyway - will fail verification`);
            leafIndex = 0;
        }
        
        console.log(`[Merkle] Found leaf at index: ${leafIndex}`);
        
        // Generate proof
        const proof = generateMerkleProof(
            { layers: currentTree.layers, depth: currentTree.depth },
            leafIndex
        );
        
        const generationTime = performance.now() - startTime;
        
        console.log(`[Merkle] Proof generated in ${generationTime.toFixed(0)}ms`);
        console.log(`[Merkle] Privacy preserved - no PII exposed!`);
        console.log(`[Merkle] ===== END PRIVACY-PRESERVING REQUEST =====`);
        
        res.json({
            pathElements: proof.pathElements,
            pathIndices: proof.pathIndices,
            root: currentTree.root,
            leafIndex: leafIndex,
            leaf: leafHash,
            generationTimeMs: generationTime,
            treeVersion: currentTree.version,
        });
        
    } catch (error) {
        console.error('[Merkle] Proof generation failed:', error);
        res.status(500).json({ 
            error: 'Failed to generate proof',
            details: error.message 
        });
    }
});

// API: Get Merkle proof for a license
app.get('/api/merkle-proof/:identifier', async (req, res) => {
    const startTime = performance.now();
    const { identifier } = req.params;
    
    console.log(`[Merkle] ===== PROOF REQUEST =====`);
    console.log(`[Merkle] Identifier: ${identifier}`);
    
    try {
        // Ensure tree is built
        if (!currentTree.isBuilt) {
            try {
                await buildTreeFromDatabase();
            } catch (err) {
                console.log('[Merkle] Database build failed, using mock:', err.message);
                await buildMockTree();
            }
        }
        
        // Try to find leaf index in the tree
        let leafIndex = currentTree.leafMap.get(identifier);
        let licenseData = null;
        let foundInTree = leafIndex !== undefined;
        
        console.log(`[Merkle] Leaf index from map: ${leafIndex}`);
        
        // // Always try to fetch license data from database
        // try {
        //     const result = await db.query(`
        //         SELECT 
        //             id,
        //             license_number,
        //             practitioner_name,
        //             EXTRACT(EPOCH FROM issued_date)::INTEGER as issued_date,
        //             EXTRACT(EPOCH FROM expiry_date)::INTEGER as expiry_date,
        //             jurisdiction,
        //             status
        //         FROM licenses 
        //         WHERE license_number = $1
        //         LIMIT 1
        //     `, [identifier]);
            
        //     if (result.rows.length > 0) {
        //         licenseData = result.rows[0];
        //         console.log(`[Merkle] Found license data in DB:`, {
        //             licenseNumber: licenseData.license_number,
        //             practitionerName: licenseData.practitioner_name,
        //             issuedDate: licenseData.issued_date,
        //             expiryDate: licenseData.expiry_date,
        //             jurisdiction: licenseData.jurisdiction
        //         });
        //     }
        // } catch (err) {
        //     console.log('[Merkle] Database lookup failed:', err.message);
        // }
        
        // If not found in tree, use index 0 (proof will fail validation)
        if (!foundInTree) {
            console.log(`[Merkle]  License NOT found in tree, using index 0 (proof will fail validation)`);
            leafIndex = 0;
        }
        
        console.log(`[Merkle] Using leaf index: ${leafIndex}`);
        console.log(`[Merkle] Leaf hash at this index: ${currentTree.leaves[leafIndex]}`);
        
        // Generate proof (even if license not in tree)
        const proof = generateMerkleProof(
            {
                layers: currentTree.layers,
                depth: currentTree.depth
            },
            leafIndex
        );
        
        console.log(`[Merkle] Proof generated:`);
        console.log(`[Merkle] - Root: ${currentTree.root}`);
        console.log(`[Merkle] - Leaf: ${currentTree.leaves[leafIndex]}`);
        console.log(`[Merkle] - Found in tree: ${foundInTree}`);
        console.log(`[Merkle] - Path indices: [${proof.pathIndices.join(', ')}]`);
        console.log(`[Merkle] - First 3 path elements:`, proof.pathElements.slice(0, 3));
        console.log(`[Merkle] ===== END PROOF REQUEST =====`);
        
        const generationTime = performance.now() - startTime;
        
        const response = {
            pathElements: proof.pathElements,
            pathIndices: proof.pathIndices,
            root: currentTree.root,
            leafIndex: leafIndex,
            leaf: currentTree.leaves[leafIndex],
            foundInTree: foundInTree,  // Flag to indicate if license was in tree
            licenseData:  null,
            generationTimeMs: generationTime,
            treeVersion: currentTree.version
        };
        
        console.log(`[Merkle] Proof generated in ${generationTime.toFixed(0)}ms`);
        res.json(response);
        
    } catch (error) {
        console.error('[Merkle] Proof generation failed:', error);
        res.status(500).json({ 
            error: 'Failed to generate proof',
            details: error.message 
        });
    }
});

// In merkle-service, update the endpoint:
// app.get('/api/merkle-proof/:identifier', async (req, res) => {
//     const startTime = performance.now();
//     const { identifier } = req.params;
    
//     console.log(`[Merkle] ===== PROOF REQUEST =====`);
//     console.log(`[Merkle] Identifier: ${identifier}`);
    
//     try {
//         if (!currentTree.isBuilt) {
//             await buildTreeFromDatabase();
//         }
        
//         // First, find the leaf index in memory
//         let leafIndex = currentTree.leafMap.get(identifier);
        
//         if (leafIndex === undefined) {
//             return res.status(404).json({ 
//                 error: 'License not found in Merkle tree',
//                 identifier: identifier
//             });
//         }
        
//         // Get the actual license data from database (without the corrupted leaf_hash)
//         let licenseData = null;
//         try {
//             const result = await db.query(`
//                 SELECT 
//                     license_number,
//                     practitioner_name,
//                     EXTRACT(EPOCH FROM issued_date)::INTEGER as issued_date,
//                     EXTRACT(EPOCH FROM expiry_date)::INTEGER as expiry_date,
//                     jurisdiction
//                 FROM licenses 
//                 WHERE license_number = $1
//                 LIMIT 1
//             `, [identifier]);
            
//             if (result.rows.length > 0) {
//                 licenseData = result.rows[0];
//                 console.log(`[Merkle] Found license data:`, licenseData);
//             }
//         } catch (err) {
//             console.error('[Merkle] Database lookup failed:', err.message);
//         }
        
//         if (!licenseData) {
//             return res.status(404).json({ error: 'License data not found' });
//         }
        
//         // Get leaf hash from in-memory tree (this is correct)
//         const leafHash = currentTree.leaves[leafIndex];
        
//         console.log(`[Merkle] Using leaf index: ${leafIndex}`);
//         console.log(`[Merkle] Leaf hash: ${leafHash}`);
        
//         // Generate proof
//         const proof = generateMerkleProof(
//             { layers: currentTree.layers, depth: currentTree.depth },
//             leafIndex
//         );
        
//         console.log(`[Merkle] Proof generated:`);
//         console.log(`[Merkle] - Root: ${currentTree.root}`);
//         console.log(`[Merkle] - Path indices: [${proof.pathIndices.join(', ')}]`);
//         console.log(`[Merkle] ===== END PROOF REQUEST =====`);
        
//         const generationTime = performance.now() - startTime;
        
//         res.json({
//             pathElements: proof.pathElements,
//             pathIndices: proof.pathIndices,
//             root: currentTree.root,
//             leafIndex: leafIndex,
//             leaf: leafHash,  // From in-memory tree
//             licenseData: {
//                 licenseNumber: licenseData.license_number,
//                 practitionerName: licenseData.practitioner_name,
//                 issuedDate: licenseData.issued_date,
//                 expiryDate: licenseData.expiry_date,
//                 jurisdiction: licenseData.jurisdiction
//             },
//             generationTimeMs: generationTime,
//             treeVersion: currentTree.version
//         });
        
//     } catch (error) {
//         console.error('[Merkle] Proof generation failed:', error);
//         res.status(500).json({ 
//             error: 'Failed to generate proof',
//             details: error.message 
//         });
//     }
// });

// API: Build/rebuild tree
app.post('/api/rebuild-tree', async (req, res) => {
    try {
        console.log('[Merkle] Api Rebuild...');
        const result = await buildTreeFromDatabase();
        res.json({
            success: true,
            ...result
        });
    } catch (error) {
        console.error('[Merkle] Rebuild failed:', error);
        
        // Try mock tree as fallback
        try {
            const mockResult = await buildMockTree();
            res.json({
                success: true,
                mode: 'mock',
                root: mockResult.root,
                leafCount: mockResult.leafCount,
                treeId: mockResult.treeId
            });
        } catch (mockError) {
            res.status(500).json({ 
                success: false,
                error: error.message 
            });
        }
    }
});

// API: Get tree info
app.get('/api/tree-info', async (req, res) => {
    try {
        if (!currentTree.isBuilt && !currentTree.isBuilding) {
            try {
                console.log('[Merkle] Tree Info build...');
                await buildTreeFromDatabase();
            } catch (err) {
                await buildMockTree();
            }
        }
        console.log("TreeId", JSON.stringify(currentTree?.treeId, null, 4));
        res.json({
            id: currentTree.treeId,
            version: currentTree.version,
            root: currentTree.root,
            leafCount: currentTree.leafCount,
            depth: currentTree.depth || TREE_DEPTH,
            maxCapacity: MAX_LEAVES,
            isBuilt: currentTree.isBuilt,
            mode: `${currentTree?.treeId}`?.includes('mock') ? 'mock' : 'real',
            createdAt: new Date().toISOString()
        });
        
    } catch (error) {
        console.error('[Merkle] Failed to get tree info:', error);
        res.status(500).json({ error: 'Failed to get tree info' });
    }
});

// API: Verify a proof
app.post('/api/verify-proof', async (req, res) => {
    try {
        const { leaf, root, pathElements, pathIndices } = req.body;
        
        // Reconstruct the root from the proof
        let computedHash = leaf;
        
        for (let i = 0; i < pathElements.length; i++) {
            const sibling = pathElements[i];
            const isLeft = pathIndices[i] === "0";
            
            if (isLeft) {
                computedHash = poseidonHash2(sibling, computedHash);
            } else {
                computedHash = poseidonHash2(computedHash, sibling);
            }
        }
        
        const isValid = computedHash === root;
        
        res.json({
            valid: isValid,
            computedRoot: computedHash,
            expectedRoot: root
        });
        
    } catch (error) {
        console.error('[Merkle] Proof verification failed:', error);
        res.status(500).json({ error: 'Failed to verify proof' });
    }
});

// Health check
app.get('/health', async (req, res) => {
    try {
        let dbHealthy = false;
        try {
            await db.query('SELECT 1');
            dbHealthy = true;
        } catch (err) {
            console.log('[Merkle] Database unhealthy:', err.message);
        }
        
        res.json({
            status: 'healthy',
            service: 'merkle-service',
            timestamp: new Date().toISOString(),
            treeVersion: currentTree.version,
            hasActiveTree: currentTree.isBuilt,
            root: currentTree.root,
            leafCount: currentTree.leafCount,
            databaseConnected: dbHealthy,
            poseidonReady: poseidon !== null
        });
    } catch (error) {
        res.status(503).json({
            status: 'unhealthy',
            service: 'merkle-service',
            error: error.message
        });
    }
});

// Statistics
app.get('/api/stats', async (req, res) => {
    try {
        const stats = {
            currentTreeVersion: currentTree.version,
            currentTreeLeaves: currentTree.leafCount || 0,
            currentTreeDepth: currentTree.depth || TREE_DEPTH,
            treeBuilt: currentTree.isBuilt,
            maxCapacity: MAX_LEAVES,
            poseidonInitialized: poseidon !== null,
            mode: `${currentTree?.treeId}`?.includes('mock') ? 'mock' : 'real',
        };
        
        // Try to get database stats
        try {
            const result = await db.query(`
                SELECT 
                    COUNT(*) as total_licenses,
                    COUNT(*) FILTER (WHERE status = 'active') as active_licenses,
                    COUNT(*) FILTER (WHERE expiry_date > CURRENT_DATE) as valid_licenses
                FROM licenses
            `);
            
            if (result.rows.length > 0) {
                stats.totalLicenses = parseInt(result.rows[0].total_licenses);
                stats.activeLicenses = parseInt(result.rows[0].active_licenses);
                stats.validLicenses = parseInt(result.rows[0].valid_licenses);
            }
        } catch (err) {
            console.log('[Merkle] Could not get database stats');
        }
        
        res.json(stats);
    } catch (error) {
        console.error('[Merkle] Failed to get stats:', error);
        res.status(500).json({ error: 'Failed to get statistics' });
    }
});

// Initialize service
async function initialize() {
    try {
        // Initialize Poseidon
        await initPoseidon();
        
        // Try to build tree from database, fall back to mock
        try {
            console.log('[Merkle] Initial build...');
            await buildTreeFromDatabase();
            console.log('[Merkle] Initial tree built from database');
        } catch (err) {
            console.log('[Merkle] Database build failed:', err.message);
            await buildMockTree();
            console.log('[Merkle] Using mock tree for testing');
        }
    } catch (error) {
        console.error('[Merkle] Initialization failed:', error);
    }
}

// Start the server
const PORT = process.env.MERKLE_PORT || 8082;

app.listen(PORT, async () => {
    console.log(`[Merkle] Service listening on port ${PORT}`);
    console.log(`[Merkle] Health check: http://localhost:${PORT}/health`);
    
    // Initialize after startup
    setTimeout(initialize, 2000);
});

// Periodic rebuild
if (process.env.UPDATE_INTERVAL) {
    const interval = parseInt(process.env.UPDATE_INTERVAL) * 1000;
    setInterval(async () => {
        console.log('[Merkle] Scheduled rebuild...');
        try {
            await buildTreeFromDatabase();
        } catch (error) {
            console.error('[Merkle] Scheduled rebuild failed:', error);
        }
    }, interval);
}

// Graceful shutdown
process.on('SIGTERM', () => {
    console.log('[Merkle] SIGTERM received, shutting down...');
    db.end().then(() => process.exit(0));
});