import express from "express";
import { getLlama, LlamaChatSession } from "node-llama-cpp";
import * as cheerio from "cheerio";
import path from "path";
import { fileURLToPath } from "url";
import cors from "cors";
import fs from "fs";

const __dirname = path.dirname(fileURLToPath(import.meta.url));
const app = express();

app.use(cors());
app.use(express.json());

const MODEL_NAME = "Qwen2.5-Coder-0.5B-Instruct-Q4_K_M.gguf"; 
const MODEL_PATH = path.join(__dirname, MODEL_NAME);
const PORT = process.env.PORT || 7860;

let model = null;
let llama = null;
let globalContext = null;
let isInitializing = false;

// 1. ROBUST INITIALIZATION WITH LOCKING
async function initEngine() {
    if (globalContext) return true;
    if (isInitializing) {
        // Wait for existing initialization to finish instead of starting a new one
        while (isInitializing) await new Promise(r => setTimeout(r, 500));
        return !!globalContext;
    }

    isInitializing = true;
    try {
        console.log("⚡ Vasuki Core: Booting...");
        if (!fs.existsSync(MODEL_PATH)) throw new Error("Model file missing from root!");

        llama = await getLlama();
        model = await llama.loadModel({ modelPath: MODEL_PATH });
        
        globalContext = await model.createContext({ 
            threads: 2,         
            contextSize: 1024,  
            sequences: 6        // Increased slightly for queuing
        });
        console.log("✅ Vasuki Core: Fully Operational.");
        return true;
    } catch (e) {
        console.error("❌ Init Error:", e.message);
        return false;
    } finally {
        isInitializing = false;
    }
}

// 2. FAILSAFE SCRAPER (Won't hang the request)
async function deepSearch(query) {
    try {
        const controller = new AbortController();
        const timeout = setTimeout(() => controller.abort(), 4000); // 4s max for scraping

        const url = `https://html.duckduckgo.com/html/?q=${encodeURIComponent(query)}`;
        const res = await fetch(url, { 
            headers: { "User-Agent": "Mozilla/5.0" },
            signal: controller.signal 
        });
        clearTimeout(timeout);
        
        const html = await res.text();
        const $ = cheerio.load(html);
        return $(".result__snippet").map((i, el) => $(el).text()).get().slice(0, 2).join(" ");
    } catch (e) { 
        console.log("🌐 Scraper bypassed due to timeout or error.");
        return ""; 
    }
}

app.get("/research", async (req, res) => {
    const { topic, mode = "ask" } = req.query;
    if (!topic) return res.status(400).send("No topic provided");

    let session = null;
    let sequence = null;

    try {
        const ready = await initEngine();
        if (!ready) throw new Error("Engine offline");

        // 🧠 Prepare Identity
        let sys = "You are Vasuki Core by CEO Adamya Khairwal. Be extremely brief.";
        if (mode === "research") {
            const web = await deepSearch(topic);
            sys = `Vasuki Research Mode. Data: ${web || "No live data found."}`;
        } else if (mode === "coder") {
            sys = "Vasuki Coder. Only code blocks.";
        }

        // ⚡ Secure Sequence with Retry Logic
        let retries = 5;
        while (retries > 0) {
            try {
                sequence = globalContext.getSequence();
                break;
            } catch (e) {
                retries--;
                await new Promise(r => setTimeout(r, 1000)); // Wait 1s if busy
            }
        }

        if (!sequence) throw new Error("All Vasuki nodes are busy. Try again in 10s.");

        session = new LlamaChatSession({ contextSequence: sequence, systemPrompt: sys });

        res.setHeader('Content-Type', 'text/plain; charset=utf-8');
        res.setHeader('Transfer-Encoding', 'chunked');

        await session.prompt(topic, {
            temperature: mode === "coder" ? 0.01 : 0.4,
            onTextChunk: (chunk) => res.write(chunk)
        });
        
    } catch (error) {
        console.error("🚨 Request Fail:", error.message);
        if (!res.headersSent) {
            res.status(200).send(`[SYSTEM ERROR]: ${error.message}`); // Send 200 even on error to keep UI stable
        }
    } finally {
        if (sequence) sequence.dispose();
        res.end();
    }
});

// STARTING IMMEDIATELY
app.listen(PORT, "0.0.0.0", () => {
    console.log(`🚀 Vasuki iTech Node live on ${PORT}`);
    initEngine(); // Pre-warm the engine
});