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

// Configuration
const PORT = 7860;
// Instruction-following model: understands "generate a login page" and responds correctly
const MODEL_NAME = 'onnx-community/Qwen2.5-0.5B-Instruct';
let generator;

// Initialize the model on startup
async function loadModel() {
    console.log("Loading coding model...");
    generator = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
    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..." }));
                }

                // Use chat format so the model understands instructions properly
                const messages = [
                    { role: 'system', content: 'You are an expert JavaScript and web developer. When asked to generate code, respond with clean, complete, working code only. No explanations unless asked.' },
                    { role: 'user', content: prompt }
                ];

                const output = await generator(messages, {
                    max_new_tokens: 512,  // More tokens for complete code output
                    temperature: 0.7,
                    do_sample: true
                });

                // Extract only the assistant's reply
                const result = output[0].generated_text.at(-1).content;
                res.writeHead(200);
                res.end(JSON.stringify({ result }));
            } catch (err) {
                res.writeHead(400);
                res.end(JSON.stringify({ error: "Invalid JSON or request", detail: err.message }));
            }
        });
        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}`);
    });
});