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

// Configuration
const PORT = 7860;
const MODEL_NAME = 'Xenova/codegen-350M-mono'; // Lightweight, powerful coding model

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) => {
    // Set headers for JSON response
    res.setHeader('Content-Type', 'application/json');

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

    // 2. Code Generation Endpoint
    if (req.url === '/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" }));
});

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