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

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

// ── SYSTEM STATE ──────────────────────────────────────────────────────────────
const queue      = [];
const jobs       = new Map();
const rateLimits = new Map();
const MAX_PARALLEL = 1;
let activeCount    = 0;

// ── UTILITIES ─────────────────────────────────────────────────────────────────
function generateId() { return crypto.randomBytes(8).toString('hex'); }

function log(event, details = {}) {
    const timestamp = new Date().toISOString();
    const logEntry  = `[${timestamp}] ${event.toUpperCase()} - ${JSON.stringify(details)}\n`;
    process.stdout.write(logEntry);
    fs.appendFileSync('generation_logs.txt', logEntry);
}

function setCORS(res) {
    res.setHeader('Access-Control-Allow-Origin',  '*');
    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, DELETE, OPTIONS');
    res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
}

function sendJSON(res, status, data) {
    setCORS(res);
    res.setHeader('Content-Type', 'application/json');
    res.writeHead(status);
    res.end(JSON.stringify(data));
}

// ── RATE LIMITING ─────────────────────────────────────────────────────────────
function checkRateLimit(ip) {
    const now   = Date.now();
    const limit = rateLimits.get(ip) || { count: 0, resetTime: now + 60000 };
    if (now > limit.resetTime) { limit.count = 0; limit.resetTime = now + 60000; }
    if (limit.count >= 10) return false;
    limit.count++;
    rateLimits.set(ip, limit);
    return true;
}

// ── QUEUE ─────────────────────────────────────────────────────────────────────
function enqueue(jobId, task) {
    return new Promise((resolve, reject) => {
        queue.push({ jobId, task, resolve, reject });
        processQueue();
    });
}

async function processQueue() {
    if (activeCount >= MAX_PARALLEL || queue.length === 0) return;
    activeCount++;
    const { jobId, task, resolve, reject } = queue.shift();
    if (jobs.get(jobId)?.status === 'cancelled') {
        activeCount--;
        return processQueue();
    }
    try {
        jobs.get(jobId).status = 'processing';
        resolve(await task());
    } catch (err) {
        reject(err);
    } finally {
        activeCount--;
        processQueue();
    }
}

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

// ── CHUNK GENERATION WITH RETRY ───────────────────────────────────────────────
async function generateChunk(messages, maxTokens, attempt = 0) {
    const startTime = Date.now();
    try {
        const output    = await generator(messages, {
            max_new_tokens:     maxTokens,
            temperature:        0.2,
            repetition_penalty: 1.15,
            do_sample:          false
        });
        const generated = output[0].generated_text;
        let text        = Array.isArray(generated)
            ? generated.at(-1)?.content || ''
            : String(generated || '');

        // Repair unclosed code blocks
        const openBlocks = (text.match(/```/g) || []).length;
        if (openBlocks % 2 !== 0) text += '\n```';

        return { text, duration: Date.now() - startTime };
    } catch (err) {
        if (attempt < 3) {
            log('retry', { attempt: attempt + 1, error: err.message });
            await new Promise(r => setTimeout(r, Math.pow(2, attempt) * 1000));
            return generateChunk(messages, maxTokens, attempt + 1);
        }
        throw err;
    }
}

// ── PROMPT CLASSIFIER ─────────────────────────────────────────────────────────
function classifyPrompt(prompt) {
    const lower = prompt.toLowerCase();

    // Identity questions
    const identityWords = ['who are you','what are you','your name','who made you',
        'who created you','your company','your founder','about you','introduce yourself'];
    if (identityWords.some(w => lower.includes(w))) return 'identity';

    // Non-coding general questions
    const generalWords = ['what is','what are','explain','describe','tell me about',
        'history of','meaning of','definition','how does','why is','who is','who was',
        'when did','where is','function of','functions of'];
    if (generalWords.some(w => lower.includes(w))) return 'general';

    // Coding / website / app requests
    const codingWords = ['generate','create','build','make','write','code','website',
        'webpage','page','app','application','login','register','form','html','css',
        'javascript','js','node','express','backend','frontend','server','api',
        'function','script','component','template','dashboard','portfolio','ecommerce',
        'shop','blog','landing page','navbar','footer','button','database'];
    if (codingWords.some(w => lower.includes(w))) return 'coding';

    return 'general';
}

// ── SYSTEM PROMPT BUILDER ─────────────────────────────────────────────────────
function buildSystemPrompt(type) {

    const IDENTITY = `You are Gini AI, a full-stack web development AI assistant.
- Created by Emalawi19.
- Founded by Professor Roosevelt Chinkwende.
- You are an AI assistant, not a human.
- When asked your name: "I am Gini AI."
- When asked who made you: "I was created by Emalawi19."
- When asked about the founder: "Emalawi19 was founded by Professor Roosevelt Chinkwende."`;

    if (type === 'identity') {
        return `${IDENTITY}
Introduce yourself clearly and warmly. State your name, your purpose, your company, and your founder.`;
    }

    if (type === 'general') {
        return `${IDENTITY}
GENERAL BEHAVIOR:
- Answer the question directly and clearly in plain English.
- Be concise and helpful.
- Do NOT generate any code for non-coding questions.
- Do NOT use programming languages to answer everyday questions.`;
    }

    // type === 'coding'
    return `${IDENTITY}

YOU ARE A FULL-STACK WEB DEVELOPER AI. FOLLOW THESE RULES STRICTLY:

RULE 1 — SINGLE FILE OUTPUT (MOST IMPORTANT):
When asked to build any website, webpage, app, or UI component:
- ALWAYS output ONE single complete HTML file.
- That single file MUST contain ALL HTML structure, ALL CSS styles, and ALL JavaScript — nothing external.
- ALL CSS goes inside a <style> tag inside <head>.
- ALL JavaScript goes inside a <script> tag at the bottom of <body>.
- NEVER say "create a separate CSS file" or "create a separate JS file".
- NEVER split code across multiple files unless the user specifically asks for backend server code.

RULE 2 — COMPLETE CODE ONLY:
- The file must be 100% complete and copy-paste ready.
- Start ALWAYS with <!DOCTYPE html> on the very first line.
- End ALWAYS with </body> then </html> as the absolute last lines.
- NEVER use placeholders like "add your code here" or "// TODO".
- NEVER truncate or cut off. Always finish the complete file.

RULE 3 — CODE QUALITY:
- Write clean, modern, well-commented HTML/CSS/JS.
- Use responsive design (flexbox or grid, mobile-friendly).
- Include hover effects, smooth transitions, and professional styling.
- Use CSS custom properties (variables) for colors and theming.
- JavaScript must be functional — forms should validate, buttons should work.

RULE 4 — BACKEND REQUESTS:
- If the user asks for a backend (Node.js, Express, API, server), provide it as a SEPARATE clearly labeled code block AFTER the frontend file.
- Backend code goes in a \`\`\`javascript block labeled "server.js".
- Still provide the complete frontend HTML file first.

RULE 5 — STEP BY STEP (only for complex multi-feature apps):
- If building a complex app, give a brief overview of what the file contains BEFORE the code.
- Then provide the single complete file.
- End with: "This is the complete file. Copy and save it as index.html and open in your browser."

RULE 6 — FORMAT:
- Always wrap the HTML file in a \`\`\`html code block.
- Always wrap any JS server code in a \`\`\`javascript code block.
- Add a short explanation after the code of what was built and how to use it.`;
}

// ── SLIDING WINDOW ENGINE ─────────────────────────────────────────────────────
async function slidingWindowEngine(res, jobId, systemPrompt, userPrompt, totalTokens) {
    const CHUNK_TOKENS  = 250;
    const CONTEXT_CHARS = 800;
    const totalChunks   = Math.ceil(totalTokens / CHUNK_TOKENS);
    const job           = jobs.get(jobId);

    let fullOutput       = job.checkpoint || '';
    let previousContext  = fullOutput.slice(-CONTEXT_CHARS);
    let chunksCompleted  = job.chunksCompleted || 0;
    let tokensGenerated  = 0;

    for (let chunkIndex = chunksCompleted; chunkIndex < totalChunks; chunkIndex++) {
        if (job.status === 'cancelled') {
            res.write(`data: ${JSON.stringify({ type: 'cancelled', message: 'Job stopped by user.' })}\n\n`);
            return res.end();
        }

        const isFirst  = chunkIndex === 0;
        const chunkNum = chunkIndex + 1;

        // First chunk: original prompt
        // Subsequent chunks: sliding window context
        const chunkUserPrompt = isFirst
            ? userPrompt
            : `[CONTINUATION INSTRUCTION]\n` +
              `Original request: ${userPrompt}\n\n` +
              `Here is the end of the code/text you have written so far:\n` +
              `...${previousContext}\n\n` +
              `IMPORTANT: Continue EXACTLY from the last character above. ` +
              `Do NOT repeat any code already written. ` +
              `Do NOT restart from <!DOCTYPE html>. ` +
              `Just continue the code seamlessly.`;

        const messages = [
            { role: 'system', content: systemPrompt },
            { role: 'user',   content: chunkUserPrompt }
        ];

        try {
            const { text: chunkText, duration } = await enqueue(jobId, () => generateChunk(messages, CHUNK_TOKENS));

            if (!chunkText || chunkText.trim().length === 0) continue;

            const words          = chunkText.split(' ');
            tokensGenerated     += words.length;
            const speed          = (words.length / (duration / 1000)).toFixed(2);
            const percentComplete = Math.round((chunkNum / totalChunks) * 100);
            const eta            = ((totalChunks - chunkNum) * (duration / 1000)).toFixed(0);

            // Stream word by word
            for (let i = 0; i < words.length; i++) {
                const token = (i === 0 && !isFirst ? '' : i === 0 ? '' : ' ') + words[i];
                res.write(`data: ${JSON.stringify({
                    type: 'token', text: token,
                    speed: `${speed} w/s`, eta: `${eta}s`,
                    progress: `${percentComplete}%`
                })}\n\n`);
                await new Promise(r => setTimeout(r, 10));
            }

            // Update sliding window
            fullOutput      += (isFirst ? '' : ' ') + chunkText;
            previousContext  = fullOutput.slice(-CONTEXT_CHARS);

            // Save checkpoint for resume
            job.checkpoint      = fullOutput;
            job.chunksCompleted = chunkNum;
            jobs.set(jobId, job);
            log('checkpoint_saved', { jobId, chunkNum });

            // Early stop if model finished naturally
            if (chunkText.trim().length < 50 && !isFirst) break;

        } catch (err) {
            log('error', { jobId, error: err.message });
            res.write(`data: ${JSON.stringify({ type: 'error', error: 'Chunk failed. Checkpoint saved for recovery.' })}\n\n`);
            return res.end();
        }
    }

    job.status = 'completed';
    res.write(`data: ${JSON.stringify({
        type: 'done', result: fullOutput.trim(), totalTokens: tokensGenerated
    })}\n\n`);
    res.end();
}

// ── HTTP SERVER ───────────────────────────────────────────────────────────────
const server = http.createServer(async (req, res) => {
    setCORS(res);
    if (req.method === 'OPTIONS') { res.writeHead(200); return res.end('{}'); }

    const url      = new URL(req.url, `http://${req.headers.host}`);
    const pathname = url.pathname;
    const clientIp = req.socket.remoteAddress;

    if (!checkRateLimit(clientIp)) return sendJSON(res, 429, { error: "Too many requests. Please wait." });

    // Status
    if (pathname === '/' && req.method === 'GET') {
        return sendJSON(res, 200, { status: modelReady ? "ready" : "loading", queue: queue.length });
    }

    // Cancel job
    if (pathname === '/cancel' && req.method === 'POST') {
        let body = '';
        req.on('data', c => { body += c.toString(); });
        req.on('end', () => {
            try {
                const { jobId } = JSON.parse(body);
                if (jobs.has(jobId)) {
                    jobs.get(jobId).status = 'cancelled';
                    log('cancelled', { jobId });
                    return sendJSON(res, 200, { message: `Job ${jobId} cancelled.` });
                }
                return sendJSON(res, 404, { error: "Job not found." });
            } catch {
                return sendJSON(res, 400, { error: "Invalid JSON" });
            }
        });
        return;
    }

    // Generate
    if (pathname === '/generate' && req.method === 'POST') {
        let body = '';
        req.on('data', c => { body += c.toString(); });
        req.on('end', async () => {
            let parsed;
            try { parsed = JSON.parse(body); }
            catch { return sendJSON(res, 400, { error: "Invalid JSON" }); }

            const { prompt, resumeJobId } = parsed;

            if (!prompt || !prompt.trim()) return sendJSON(res, 400, { error: "prompt is required" });
            if (!modelReady)               return sendJSON(res, 503, { error: "Model loading, please wait..." });

            // Classify the request
            const promptType   = classifyPrompt(prompt);
            const systemPrompt = buildSystemPrompt(promptType);

            // Decide token budget
            const totalTokens = promptType === 'coding'  ? 50
                               : promptType === 'identity' ? 10
                               : 50;

            // Resume or new job
            const jobId = resumeJobId && jobs.has(resumeJobId) ? resumeJobId : generateId();
            if (!jobs.has(jobId)) {
                jobs.set(jobId, { status: 'queued', checkpoint: '', chunksCompleted: 0 });
            }

            log('job_started', { jobId, promptType, totalTokens, prompt: prompt.slice(0, 80) });

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

            res.write(`data: ${JSON.stringify({
                type: 'start', jobId, promptType,
                queuePosition: queue.length
            })}\n\n`);

            try {
                await slidingWindowEngine(res, jobId, systemPrompt, prompt.trim(), totalTokens);
            } catch (err) {
                res.write(`data: ${JSON.stringify({ type: 'fatal_error', error: err.message })}\n\n`);
                res.end();
            }
        });
        return;
    }

    return sendJSON(res, 404, { error: "Not found" });
});

loadModel().then(() => {
    server.listen(PORT, '0.0.0.0', () => console.log(`Gini AI engine running on port ${PORT}`));
});