server.js

import { pipeline } from '@huggingface/transformers';
import http from 'http';

// Configuration
const PORT = 7860;
const MODEL_NAME = 'Xenova/codegen-2B-mono'; // ← Upgraded: 2B param model, much stronger code generation
let generator;

// Initialize the model on startup
async function loadModel() {
    console.log("Loading coding model...");
    generator = await pipeline('text-generation', MODEL_NAME);
    console.log("Model loaded successfully!");
}

const server = http.createServer(async (req, res) => {
    // 1. Add CORS headers so external websites can talk to this API
    res.setHeader('Access-Control-Allow-Origin', '*');
    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
    res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
    
    // Handle preflight requests for the API
    if (req.method === 'OPTIONS') {
        res.writeHead(200);
        return res.end();
    }

    res.setHeader('Content-Type', 'application/json');

    // Clean the URL (Removes Hugging Face's ?__theme=light additions)
    const pathname = req.url.split('?')[0];

    // 2. Your requested Status Check
    if (pathname === '/' && req.method === 'GET') {
        res.writeHead(200);
        res.end(JSON.stringify({ "status": " Backend is running" }));
        return;
    }

    // 3. Code Generation Endpoint
    if (pathname === '/generate' && req.method === 'POST') {
        let body = '';
        req.on('data', chunk => { body += chunk.toString(); });
        req.on('end', async () => {
            try {
                const { prompt } = JSON.parse(body);
                if (!generator) {
                    res.writeHead(503);
                    return res.end(JSON.stringify({ error: "Model is still loading..." }));
                }
                // Generate code based on the prompt
                const output = await generator(prompt, { 
                    max_new_tokens: 100,
                    temperature: 0.7 
                });
                res.writeHead(200);
                res.end(JSON.stringify({ result: output[0].generated_text }));
            } catch (err) {
                res.writeHead(400);
                res.end(JSON.stringify({ error: "Invalid JSON or request" }));
            }
        });
        return;
    }

    // Default 404
    res.writeHead(404);
    res.end(JSON.stringify({ error: "Not Found", requested_path: pathname }));
});

// Start everything
loadModel().then(() => {
    server.listen(PORT, '0.0.0.0', () => {
        console.log(`Server running at http://0.0.0.0:${PORT}`);
    });
});