server.js

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

const PORT = 7860;
const MODEL_NAME = 'onnx-community/Qwen2.5-0.5B-Instruct';
let generator;

async function loadModel() {
    console.log("Loading coding model...");
    generator = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
    console.log("Model loaded successfully!");
}

// ── WIKIPEDIA ────────────────────────────────────────────────────────────────
async function searchWikipedia(query) {
    try {
        const url = `https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch=${encodeURIComponent(query)}&srlimit=1&format=json&origin=*`;
        const res  = await fetch(url);
        const data = await res.json();
        const results = data?.query?.search;
        if (!results || results.length === 0) return null;
        return results[0].title;
    } catch (e) {
        console.error("Wikipedia search error:", e.message);
        return null;
    }
}

async function fetchWikipediaSummary(title) {
    try {
        const url = `https://en.wikipedia.org/api/rest_v1/page/summary/${encodeURIComponent(title)}`;
        const res  = await fetch(url);
        const data = await res.json();
        if (data.extract) return data.extract.slice(0, 800);
        return null;
    } catch (e) {
        console.error("Wikipedia fetch error:", e.message);
        return null;
    }
}

async function getWikipediaContext(prompt) {
    try {
        const title = await searchWikipedia(prompt);
        if (!title) return null;
        const summary = await fetchWikipediaSummary(title);
        if (!summary) return null;
        console.log(`Wikipedia context fetched: "${title}"`);
        return { title, summary };
    } catch (e) {
        console.error("Wikipedia context error:", e.message);
        return null;
    }
}

// ── CODING DETECTION ─────────────────────────────────────────────────────────
const CODING_PHRASES = [
    'generate code', 'write code', 'create code',
    'generate a website', 'build a website', 'create a website', 'make a website',
    'generate a webpage', 'build a webpage',
    'generate a login', 'create a login', 'build a login',
    'generate a form', 'create a form', 'build a form',
    'generate a page', 'create a page', 'build a page',
    'generate an app', 'create an app', 'build an app',
    'write a function', 'write a script', 'write a program',
    'create a function', 'create a script',
    'build an api', 'create an api', 'generate an api',
    'node.js', 'express.js', 'express server',
    'show me the code', 'give me the code', 'write html',
    'write css', 'write javascript', 'write js',
    'create a backend', 'build a backend',
    'create a frontend', 'build a frontend',
    'create a server', 'build a server',
    'how do i code', 'how to code', 'code for',
    'sample code', 'example code', '```'
];

const NON_CODING_CONTEXTS = [
    'function of', 'functions of', 'what is the function',
    'what are the functions', 'purpose of', 'role of',
    'explain', 'what is', 'what are', 'tell me about',
    'history of', 'meaning of', 'definition of',
    'describe', 'how does', 'why is', 'benefits of',
    'importance of', 'effects of', 'causes of',
    'who is', 'who was', 'when did', 'when was',
    'where is', 'where was', 'which is', 'which was'
];

function isCodingRequest(prompt) {
    const lower = prompt.toLowerCase().trim();
    if (NON_CODING_CONTEXTS.some(ctx => lower.includes(ctx))) return false;
    return CODING_PHRASES.some(phrase => lower.includes(phrase));
}

function isWebsiteRequest(prompt) {
    const lower = prompt.toLowerCase();
    return ['website', 'webpage', 'web page', 'html page', 'login page',
            'register page', 'landing page', 'form page', 'signup page',
            'homepage', 'home page', 'portfolio', 'dashboard'].some(w => lower.includes(w));
}

function cleanHTML(html) {
    if (!html.trimStart().toLowerCase().startsWith('<!doctype')) {
        if (html.trimStart().toLowerCase().startsWith('<html')) {
            html = '<!DOCTYPE html>\n' + html.trimStart();
        } else {
            html = '<!DOCTYPE html>\n<html lang="en">\n' + html.trimStart();
        }
    }
    if (!html.trimEnd().toLowerCase().endsWith('</html>')) {
        if (html.toLowerCase().includes('</body>')) {
            html = html.replace(/<\/body>\s*$/i, '</body>\n</html>');
        } else {
            html = html.trimEnd() + '\n</body>\n</html>';
        }
    }
    return html.trim();
}

function enforceHTMLStructure(text) {
    const codeBlockMatch = text.match(/```html\s*([\s\S]*?)```/i);
    if (codeBlockMatch) {
        return '```html\n' + cleanHTML(codeBlockMatch[1].trim()) + '\n```';
    }
    const htmlStart = text.search(/<(!DOCTYPE|html)/i);
    if (htmlStart !== -1) {
        let html = text.slice(htmlStart);
        const htmlEnd = html.search(/<\/html>/i);
        if (htmlEnd !== -1) html = html.slice(0, htmlEnd + 7);
        return '```html\n' + cleanHTML(html) + '\n```';
    }
    return text;
}

// ── CORE GENERATION (non-streaming, reliable) ─────────────────────────────────
async function generateResponse(messages, coding) {
    const output = await generator(messages, {
        max_new_tokens: coding ? 800 : 400,
        temperature: coding ? 0.2 : 0.5,
        repetition_penalty: 1.2,
        do_sample: !coding
    });
    return output[0].generated_text.at(-1).content || '';
}

// ── HTTP SERVER ───────────────────────────────────────────────────────────────
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();
    }

    const pathname = req.url.split('?')[0];

    if (pathname === '/' && req.method === 'GET') {
        res.setHeader('Content-Type', 'application/json');
        res.writeHead(200);
        res.end(JSON.stringify({ status: "Backend is running", model: MODEL_NAME }));
        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.setHeader('Content-Type', 'application/json');
                    res.writeHead(503);
                    return res.end(JSON.stringify({ error: "Model is still loading..." }));
                }

                const coding     = isCodingRequest(prompt);
                const websiteReq = isWebsiteRequest(prompt);

                // ── Fetch Wikipedia for non-coding questions ──────────────────
                let wikipediaContext = '';
                let wikiTitle       = null;
                if (!coding) {
                    const wiki = await getWikipediaContext(prompt);
                    if (wiki) {
                        wikiTitle       = wiki.title;
                        wikipediaContext = `\n\nWIKIPEDIA REFERENCE (article: "${wiki.title}"):\n${wiki.summary}\n\nIMPORTANT: Use this Wikipedia data to answer accurately. Rewrite it in your own words and add helpful context.`;
                    }
                }

                const finalPrompt = coding
                    ? `You are a coding assistant. Write complete, working code. Never refuse: ${prompt}`
                    : prompt;

                const systemPrompt = coding
                    ? `You are Gini AI, a full-stack web development AI assistant created by Emalawi19, founded by Professor Roosevelt Chinkwende.

IDENTITY: Your name is Gini AI. Company: Emalawi19. Founder: Professor Roosevelt Chinkwende.

CODING RULES:
1. ALWAYS write code when asked. NEVER refuse.
2. Complete, real, copy-paste ready code only — no placeholders.
3. Default: JavaScript/HTML/CSS. Never Python unless asked.
4. Add comments inside the code.

HTML RULES:
1. Start with <!DOCTYPE html> then <html lang="en">
2. Include proper <head> with meta tags, title, and <style>
3. End with </body> then </html>
4. All CSS in <style>, all JS in <script> at bottom of body.
5. Wrap in \`\`\`html code block. Never cut off mid-code.

STEP-BY-STEP: List all steps first, then one step at a time. Say "next" to continue.`

                    : `You are Gini AI, a helpful AI assistant created by Emalawi19, founded by Professor Roosevelt Chinkwende.

IDENTITY: Your name is Gini AI. Company: Emalawi19. Founder: Professor Roosevelt Chinkwende. You are an AI, not a human.
${wikipediaContext}
RULES:
- Answer the question directly and clearly in plain English.
- If Wikipedia reference is provided above, use it to give an accurate, well-informed answer.
- Rewrite Wikipedia info in your own words — never copy it directly.
- If the question is about a person (president, leader, founder etc.), state their name clearly at the start.
- NEVER generate code for general knowledge questions.
- Keep answers friendly, clear, and concise.`;

                // ── Set SSE headers ───────────────────────────────────────────
                res.writeHead(200, {
                    'Content-Type': 'text/event-stream',
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                    'X-Accel-Buffering': 'no'
                });

                // Send a heartbeat so the client knows we're working
                res.write(`data: ${JSON.stringify({ status: "generating" })}\n\n`);

                // ── Generate (wait for full result, then stream it out) ────────
                // The callback_function in transformers.js is unreliable on CPU
                // so we generate fully then simulate streaming character by character
                const messages = [
                    { role: 'system', content: systemPrompt },
                    { role: 'user',   content: finalPrompt  }
                ];

                let result = await generateResponse(messages, coding);

                // Catch refusals
                const refusals = ["i can't assist", "i cannot assist", "i'm sorry", "i am sorry", "i'm unable", "i cannot help", "i can't help"];
                const refused  = refusals.some(r => result.toLowerCase().startsWith(r));

                if (refused && coding) {
                    result = await generateResponse([
                        { role: 'system', content: 'You are a code generator. Output only complete working code. Never refuse.' },
                        { role: 'user',   content: `Write complete HTML code starting with <!DOCTYPE html> ending with </html> for: ${prompt}` }
                    ], true);
                }

                // If result is still empty, use Wikipedia summary directly
                if (!result || result.trim().length < 5) {
                    if (wikiTitle && wikipediaContext) {
                        result = `Based on Wikipedia: ${wikipediaContext.split('\n').slice(3).join(' ').trim()}`;
                    } else {
                        result = "I'm sorry, I couldn't generate a response. Please try again.";
                    }
                }

                if (websiteReq && coding) {
                    result = enforceHTMLStructure(result);
                }

                // ── Stream the result word by word ────────────────────────────
                const words = result.split(' ');
                for (let i = 0; i < words.length; i++) {
                    const chunk = (i === 0 ? '' : ' ') + words[i];
                    res.write(`data: ${JSON.stringify({ chunk })}\n\n`);
                    // Small delay so frontend can render progressively
                    await new Promise(r => setTimeout(r, 15));
                }

                // Send done event
                res.write(`data: ${JSON.stringify({ done: true, result, source: wikiTitle ? `📖 Wikipedia — ${wikiTitle}` : null })}\n\n`);
                res.end();

            } catch (err) {
                console.error("Generation error:", err.message);
                res.write(`data: ${JSON.stringify({ error: err.message })}\n\n`);
                res.end();
            }
        });
        return;
    }

    res.setHeader('Content-Type', 'application/json');
    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}`);
    });
});