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

// Configuration
const PORT = 7860;
// UPGRADED: Starcoder2-1b is a much more capable coding model
const MODEL_NAME = 'onnx-community/starcoder2-3b'; 

let generator;

// Initialize the model on startup
async function loadModel() {
    console.log("Loading upgraded coding model...");
    // We specify 'auto' for the device to let the server manage memory efficiently
    generator = await pipeline('text-generation', MODEL_NAME, { device: 'auto' });
    console.log("Model loaded successfully!");
}

const server = http.createServer(async (req, res) => {
    res.setHeader('Access-Control-Allow-Origin', '*');
    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
    res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
    
    if (req.method === 'OPTIONS') {
        res.writeHead(200);
        return res.end();
    }

    res.setHeader('Content-Type', 'application/json');
    const pathname = req.url.split('?')[0];

    if (pathname === '/' && req.method === 'GET') {
        res.writeHead(200);
        res.end(JSON.stringify({ "status": "Backend running with Starcoder2-1b" }));
        return;
    }

    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..." }));
                }

                // UPGRADED: Increased tokens and tuned temperature for better code quality
                const output = await generator(prompt, { 
                    max_new_tokens: 250, 
                    temperature: 0.2,
                    do_sample: true
                });

                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;
    }

    res.writeHead(404);
    res.end(JSON.stringify({ error: "Not Found" }));
});

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