server.js

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

const PORT       = 7860;
const MODEL_NAME = 'onnx-community/Qwen2.5-0.5B-Instruct';
const KNOWLEDGE_DIR = './knowledge';
let generator;
let knowledgeBase = [];

// ── KNOWLEDGE FILES ───────────────────────────────────────────────────────────
function loadKnowledge() {
    if (!fs.existsSync(KNOWLEDGE_DIR)) { fs.mkdirSync(KNOWLEDGE_DIR); return; }
    const files = fs.readdirSync(KNOWLEDGE_DIR).filter(f => f.endsWith('.txt'));
    knowledgeBase = [];
    for (const file of files) {
        const content = fs.readFileSync(path.join(KNOWLEDGE_DIR, file), 'utf-8');
        const chunks  = splitChunks(content, 600, 60);
        chunks.forEach(c => knowledgeBase.push({ source: file, text: c }));
    }
    console.log(`Knowledge loaded: ${files.length} files, ${knowledgeBase.length} chunks`);
}

function splitChunks(text, size, overlap) {
    const chunks = [];
    let start = 0;
    while (start < text.length) {
        chunks.push(text.slice(start, start + size));
        start += size - overlap;
    }
    return chunks;
}

function retrieveContext(prompt, topK = 4) {
    if (knowledgeBase.length === 0) return '';
    const words  = prompt.toLowerCase().split(/\W+/).filter(w => w.length > 2);
    const scored = knowledgeBase.map(chunk => ({
        ...chunk,
        score: words.reduce((acc, w) => acc + (chunk.text.toLowerCase().includes(w) ? 1 : 0), 0)
    }));
    return scored
        .filter(c => c.score > 0)
        .sort((a, b) => b.score - a.score)
        .slice(0, topK)
        .map(c => `[${c.source}]\n${c.text}`)
        .join('\n\n---\n\n');
}

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

async function generateResponse(messages) {
    const output = await generator(messages, {
        max_new_tokens:     500,
        temperature:        0.3,
        repetition_penalty: 1.2,
        do_sample:          false
    });
    // Extract only the assistant reply content
    const generated = output[0].generated_text;
    if (Array.isArray(generated)) {
        return generated.at(-1)?.content || '';
    }
    return String(generated || '');
}

// ── 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];

    // Status check
    if (pathname === '/' && req.method === 'GET') {
        res.setHeader('Content-Type', 'application/json');
        res.writeHead(200);
        return res.end(JSON.stringify({
            status: "running",
            model: MODEL_NAME,
            knowledge_chunks: knowledgeBase.length
        }));
    }

    // Reload knowledge
    if (pathname === '/reload-knowledge' && req.method === 'POST') {
        loadKnowledge();
        res.setHeader('Content-Type', 'application/json');
        res.writeHead(200);
        return res.end(JSON.stringify({ message: `Reloaded: ${knowledgeBase.length} chunks` }));
    }

    // Main generate endpoint — returns plain JSON (no SSE, no streaming delays)
    if (pathname === '/generate' && req.method === 'POST') {
        let body = '';
        req.on('data', c => { body += c.toString(); });
        req.on('end', async () => {
            res.setHeader('Content-Type', 'application/json');

            try {
                const { prompt, system } = JSON.parse(body);

                if (!generator) {
                    res.writeHead(503);
                    return res.end(JSON.stringify({ error: "Model still loading, please wait..." }));
                }

                // RAG: retrieve relevant knowledge chunks
                const ragContext = retrieveContext(prompt, 4);
                const ragSection = ragContext
                    ? `\n\nKNOWLEDGE BASE — use ONLY this information to answer:\n${ragContext}\n`
                    : '';

                // Build final system prompt
                const finalSystem = (system || `You are Mlimi Connect AI, a free agricultural advisor for Malawian farmers. Only answer agriculture questions.`) + ragSection;

                const messages = [
                    { role: 'system', content: finalSystem },
                    { role: 'user',   content: prompt       }
                ];

                console.log(`Generating response for: "${prompt.slice(0, 60)}..."`);

                let result = await generateResponse(messages);

                // Fallback if model returns empty
                if (!result || result.trim().length < 5) {
                    if (ragContext) {
                        result = `Here is what I know about this topic:\n\n${ragContext.slice(0, 600)}`;
                    } else {
                        result = "I don't have specific information on that topic. Please consult your local agricultural extension officer (AEO) for advice.";
                    }
                }

                console.log(`Response ready: ${result.length} chars`);

                res.writeHead(200);
                res.end(JSON.stringify({ result }));

            } catch (err) {
                console.error("Generation error:", err.message);
                res.writeHead(500);
                res.end(JSON.stringify({ error: err.message || "Generation failed" }));
            }
        });
        return;
    }

    res.setHeader('Content-Type', 'application/json');
    res.writeHead(404);
    res.end(JSON.stringify({ error: "Not Found" }));
});

loadKnowledge();
loadModel().then(() => {
    server.listen(PORT, '0.0.0.0', () => {
        console.log(`Mlimi Connect backend running on port ${PORT}`);
    });
});