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 & MEMORY ─────────────────────────────────────────────────────
const queue = [];
const jobs = new Map();         // Stores job state, checkpoints, and cancellation flags
const rateLimits = new Map();   // IP-based rate limiting
const MAX_PARALLEL = 1;         // Keep at 1 for transformers.js to prevent thread blocking
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); // Log to file
}

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

// ── QUEUE & 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; // Max 10 requests per minute
    limit.count++;
    rateLimits.set(ip, limit);
    return true;
}

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();
    
    // Check if job was cancelled while in queue
    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 INITIALIZATION ──────────────────────────────────────────────────────
async function loadModel() {
    log('system', { message: "Loading model..." });
    generator = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
    modelReady = true;
    log('system', { message: "Model ready!" });
}

// ── GENERATION & VALIDATION ───────────────────────────────────────────────────
async function generateChunk(messages, maxTokens, attempt = 0) {
    const startTime = Date.now();
    try {
        const output = await generator(messages, {
            max_new_tokens: maxTokens,
            temperature: 0.3,
            repetition_penalty: 1.15,
            do_sample: false
        });
        const generated = output[0].generated_text;
        let text = Array.isArray(generated) ? generated.at(-1)?.content || '' : String(generated || '');
        
        // Output Validation: Detect truncated code blocks
        const openBlocks = (text.match(/```/g) || []).length;
        if (openBlocks % 2 !== 0) text += '\n```'; // Repair formatting
        
        return { text, duration: Date.now() - startTime };
    } catch (err) {
        if (attempt < 3) { // Failure Recovery: Exponential backoff
            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;
    }
}

// ── SLIDING WINDOW ENGINE (AUTO-RESUME & CHUNKING) ────────────────────────────
async function slidingWindowEngine(res, jobId, systemPrompt, userPrompt, totalTokens) {
    const CHUNK_TOKENS = 250;
    const CONTEXT_CHARS = 800; // Expanded for better Context Awareness
    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;
    const globalStartTime = Date.now();

    for (let chunkIndex = chunksCompleted; chunkIndex < totalChunks; chunkIndex++) {
        // Cancellation Check
        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;

        let chunkUserPrompt = isFirst ? userPrompt :
            `[SYSTEM: Intelligent Context Compression]\n` +
            `Original Request: ${userPrompt}\n` +
            `Recent Output Context (Maintain consistency, variables, and formatting):\n...${previousContext}\n\n` +
            `INSTRUCTION: Continue generating seamlessly from the exact last character above. Do not repeat the context.`;

        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; // Rough estimation
            
            // Progress Tracking Metrics
            const speed = (words.length / (duration / 1000)).toFixed(2); // words per second
            const percentComplete = Math.round((chunkNum / totalChunks) * 100);
            const eta = ((totalChunks - chunkNum) * (duration / 1000)).toFixed(0);

            // Stream Output
            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)); // Streaming effect
            }

            // Auto Resume & Checkpoint System Update
            fullOutput += (isFirst ? '' : ' ') + chunkText;
            previousContext = fullOutput.slice(-CONTEXT_CHARS);
            
            job.checkpoint = fullOutput;
            job.chunksCompleted = chunkNum;
            jobs.set(jobId, job);
            log('checkpoint_saved', { jobId, chunkNum });

            // Stop condition
            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;

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

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

    // Cancellation Endpoint
    if (pathname === '/cancel' && req.method === 'POST') {
        let body = '';
        req.on('data', c => { body += c.toString(); });
        req.on('end', () => {
            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." });
        });
        return;
    }

    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 (!modelReady) return sendJSON(res, 503, { error: "Model loading..." });

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

            const totalTokens = 2000; // Simulated dynamic estimate
            log('job_started', { jobId, promptLength: prompt?.length });

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

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

            try {
                await slidingWindowEngine(res, jobId, "You are a helpful coding assistant.", prompt || "Continue", 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(`Engine running on port ${PORT}`));
});