tinkering.ts

/// <reference types="@webgpu/types" />
import {
    AutoTokenizer,
    AutoModelForCausalLM,
    TextStreamer,
    pipeline,
    env,
    InterruptableStoppingCriteria
} from "@huggingface/transformers"

import { detectExtension, hasExtensionMarker } from './src/extension-bridge';
import { runPipeline, type OrchestratorState, type Models } from './src/orchestrator';
import { StreamingMarkdownRenderer } from './src/markdown-renderer';

// Configure environment to be more resilient
env.allowRemoteModels = true;

// ============================================
// GPU DETECTION AND SELECTION
// ============================================

async function checkAndRequestGPU(): Promise<string> {
    if (!navigator.gpu) {
        console.warn("WebGPU is not available");
        return "WebGPU not available";
    }

    try {
        // Request high-performance GPU adapter
        const adapter = await navigator.gpu.requestAdapter({
            powerPreference: "high-performance"
        });

        if (!adapter) {
            console.warn("No WebGPU adapter found");
            return "No adapter found";
        }

        // Get adapter info
        const info = adapter.info;
        console.log(info);
        const gpuName = info?.device || info?.description || "Unknown GPU";
        const gpuVendor = info?.vendor || "Unknown vendor";

        console.log("=".repeat(50));
        console.log("🎮 WebGPU Adapter Info:");
        console.log(`   Vendor: ${gpuVendor}`);
        console.log(`   Device: ${gpuName}`);
        console.log(`   Architecture: ${info?.architecture || "N/A"}`);
        console.log("=".repeat(50));

        // Check if it's likely an integrated GPU
        const gpuNameLower = gpuName.toLowerCase();
        const vendorLower = gpuVendor.toLowerCase();
        const isIntegrated =
            gpuNameLower.includes("intel") ||
            gpuNameLower.includes("radeon graphics") ||
            gpuNameLower.includes("vega") ||
            (vendorLower.includes("amd") && !gpuNameLower.includes("rx "));

        if (isIntegrated) {
            console.warn("⚠️ Using integrated GPU! For better performance:");
            console.warn("   1. Go to Windows Settings → Display → Graphics");
            console.warn("   2. Add your browser and set to 'High performance'");
            console.warn("   3. Or use NVIDIA Control Panel → Manage 3D settings");
        }

        return `${gpuVendor} ${gpuName}`;
    } catch (e) {
        console.error("GPU check failed:", e);
        return "GPU check failed";
    }
}

// ============================================
// UI ELEMENTS
// ============================================

// Loading
const loadingOverlay = document.getElementById('loading-overlay') as HTMLElement;
const loadingMessage = document.getElementById('loading-message') as HTMLElement;
const progressBar = document.getElementById('progress-bar') as HTMLElement;

// Header
const statusBadge = document.getElementById('status-badge') as HTMLElement;

// Chat
const welcomeScreen = document.getElementById('welcome-screen') as HTMLElement;
const chatContainer = document.getElementById('chat-container') as HTMLElement;
const userInput = document.getElementById('user-input') as HTMLTextAreaElement;
const sendBtn = document.getElementById('send-btn') as HTMLButtonElement;
const sendIcon = document.getElementById('send-icon') as HTMLImageElement;

// Monitoring Panels
const processContent = document.getElementById('process-content') as HTMLElement;
const processEmpty = document.getElementById('process-empty') as HTMLElement;
const reasoningContent = document.getElementById('reasoning-content') as HTMLElement;
const reasoningEmpty = document.getElementById('reasoning-empty') as HTMLElement;

// Modal
const onboardingModal = document.getElementById('onboarding-modal') as HTMLElement;
const enableSearchBtn = document.getElementById('enable-search-btn') as HTMLButtonElement;
const skipSearchBtn = document.getElementById('skip-search-btn') as HTMLButtonElement;

// Stopping criteria for interrupting generation
let stoppingCriteria: InterruptableStoppingCriteria | null = null;

// ============================================
// UI HELPER FUNCTIONS
// ============================================

function updateStatus(text: string, ready: boolean = false) {
    if (statusBadge) {
        statusBadge.innerHTML = `<span class="${ready ? 'text-accent' : 'text-muted'}">${text}</span>`;
    }
}

function updateProgress(percent: number, message: string) {
    if (progressBar) progressBar.style.width = `${percent}%`;
    if (loadingMessage) loadingMessage.textContent = message;
}

function showChatUI() {
    loadingOverlay?.classList.add('hidden');
    welcomeScreen?.classList.remove('hidden');
    chatContainer?.classList.add('hidden');
    if (userInput) userInput.disabled = false;
    if (sendBtn) sendBtn.disabled = false;
    userInput?.focus();
}

function showOnboardingModal() {
    onboardingModal?.classList.remove('hidden');
}

function hideOnboardingModal() {
    onboardingModal?.classList.add('hidden');
}

// Toggle between send and stop button
function setSendButtonState(isGenerating: boolean) {
    if (!sendBtn || !sendIcon) return;

    if (isGenerating) {
        sendIcon.src = 'Assets/Stop-Response.svg';
        sendIcon.alt = 'Stop';
        sendBtn.classList.add('stop-btn');
        sendBtn.title = 'Stop generation';
    } else {
        sendIcon.src = 'Assets/Send-Button.svg';
        sendIcon.alt = 'Send';
        sendBtn.classList.remove('stop-btn');
        sendBtn.title = 'Send message';
    }
}

// ============================================
// PROCESS PANEL FUNCTIONS
// ============================================

function clearProcessPanel() {
    if (!processContent || !processEmpty) return;
    // Remove all process items but keep empty state
    const items = processContent.querySelectorAll('.process-item');
    items.forEach(item => item.remove());
    processEmpty.classList.remove('hidden');
}

function addProcessItem(icon: string, text: string) {
    if (!processContent || !processEmpty) return;

    // Hide empty state
    processEmpty.classList.add('hidden');

    const item = document.createElement('div');
    item.className = 'process-item';
    item.innerHTML = `
        <img src="Assets/${icon}" alt="" class="process-item-icon">
        <span class="process-item-text">${text}</span>
    `;
    processContent.appendChild(item);
    processContent.scrollTop = processContent.scrollHeight;
}

function updateStepIndicator(step: OrchestratorState['currentStep']) {
    // Map steps to process panel items
    const stepConfig: Record<string, { icon: string; text: string }> = {
        'classifying': { icon: 'GuardIcon.svg', text: 'Sentinel is analyzing query...' },
        'searching': { icon: 'SearchIcon.svg', text: 'Searching the web...' },
        'reasoning': { icon: 'Brain-Icon.svg', text: 'Generating response...' }
    };

    const config = stepConfig[step];
    if (config) {
        addProcessItem(config.icon, config.text);
    }
}

// ============================================
// REASONING PANEL FUNCTIONS
// ============================================

function clearReasoningPanel() {
    if (!reasoningContent || !reasoningEmpty) return;
    // Remove all content but keep empty state
    const items = reasoningContent.querySelectorAll('.reasoning-text');
    items.forEach(item => item.remove());
    reasoningEmpty.classList.remove('hidden');
}

function updateReasoningPanel(text: string) {
    if (!reasoningContent || !reasoningEmpty) return;

    // Hide empty state
    reasoningEmpty.classList.add('hidden');

    // Find or create reasoning text element
    let reasoningText = reasoningContent.querySelector('.reasoning-text') as HTMLElement;
    if (!reasoningText) {
        reasoningText = document.createElement('div');
        reasoningText.className = 'reasoning-text';
        reasoningContent.appendChild(reasoningText);
    }

    reasoningText.textContent = text;
    reasoningContent.scrollTop = reasoningContent.scrollHeight;
}

function addMessage(content: string, isUser: boolean): HTMLElement {
    // Show chat container, hide welcome on first message
    if (welcomeScreen && !welcomeScreen.classList.contains('hidden')) {
        welcomeScreen.classList.add('hidden');
        chatContainer.classList.remove('hidden');
    }

    const messageDiv = document.createElement('div');
    messageDiv.className = `message ${isUser ? 'user' : 'bot'}`;

    const messageBubble = document.createElement('div');
    messageBubble.className = 'message-bubble';
    messageBubble.textContent = content;

    messageDiv.appendChild(messageBubble);
    chatContainer.appendChild(messageDiv);

    // Scroll to bottom
    chatContainer.scrollTop = chatContainer.scrollHeight;

    return messageBubble;
}

/**
 * Add a streaming message with markdown support and optional thinking panel
 * Returns the StreamingMarkdownRenderer for pushing content
 */
function addStreamingMessage(): { renderer: StreamingMarkdownRenderer; thinkingPanel: HTMLElement } {
    // Show chat container, hide welcome on first message
    if (welcomeScreen && !welcomeScreen.classList.contains('hidden')) {
        welcomeScreen.classList.add('hidden');
        chatContainer.classList.remove('hidden');
    }

    const messageDiv = document.createElement('div');
    messageDiv.className = 'message bot';

    // Mobile: Inline reasoning toggle (collapsible)
    const mobileReasoning = document.createElement('div');
    mobileReasoning.className = 'message-reasoning mobile-only hidden';

    const reasoningToggle = document.createElement('div');
    reasoningToggle.className = 'reasoning-toggle collapsed';
    reasoningToggle.innerHTML = `
        <img src="Assets/DownArrow.svg" alt="" class="reasoning-toggle-icon">
        <span>Reasoning</span>
    `;
    reasoningToggle.onclick = () => {
        reasoningToggle.classList.toggle('collapsed');
    };

    const mobileReasoningContent = document.createElement('div');
    mobileReasoningContent.className = 'reasoning-content';

    mobileReasoning.appendChild(reasoningToggle);
    mobileReasoning.appendChild(mobileReasoningContent);

    // Main message bubble with markdown content
    const messageBubble = document.createElement('div');
    messageBubble.className = 'message-bubble markdown-body';

    messageDiv.appendChild(mobileReasoning);
    messageDiv.appendChild(messageBubble);
    chatContainer.appendChild(messageDiv);

    // Create the streaming renderer that updates both mobile inline and desktop panel
    const renderer = new StreamingMarkdownRenderer({
        contentContainer: messageBubble,
        thinkingContainer: mobileReasoningContent,
        onThinkingStateChange: (isThinking) => {
            if (isThinking) {
                // Thinking started - show and expand reasoning
                mobileReasoning.classList.remove('hidden');
                reasoningToggle.classList.remove('collapsed');
            } else {
                // Thinking ended - collapse into dropdown
                reasoningToggle.classList.add('collapsed');
            }
        },
        onThinkingContent: (fullThinkingText) => {
            // Update desktop reasoning panel on every chunk
            updateReasoningPanel(fullThinkingText);
        },
        autoScroll: true,
        scrollContainer: chatContainer
    });

    // Scroll to bottom
    chatContainer.scrollTop = chatContainer.scrollHeight;

    return { renderer, thinkingPanel: mobileReasoning };
}

function addTypingIndicator(): HTMLElement {
    const messageDiv = document.createElement('div');
    messageDiv.className = 'message bot';
    messageDiv.id = 'typing-indicator';

    const messageBubble = document.createElement('div');
    messageBubble.className = 'message-bubble';

    const typingDots = document.createElement('div');
    typingDots.className = 'typing-indicator';
    typingDots.innerHTML = '<span></span><span></span><span></span>';

    messageBubble.appendChild(typingDots);
    messageDiv.appendChild(messageBubble);
    chatContainer.appendChild(messageDiv);

    chatContainer.scrollTop = chatContainer.scrollHeight;

    return messageDiv;
}

function removeTypingIndicator() {
    const indicator = document.getElementById('typing-indicator');
    if (indicator) {
        indicator.remove();
    }
}

// ============================================
// CUSTOM STREAMER FOR UI
// ============================================

class UITextStreamer extends TextStreamer {
    private targetElement: HTMLElement;
    private currentText: string = '';

    constructor(tokenizer: any, targetElement: HTMLElement, options: any = {}) {
        super(tokenizer, options);
        this.targetElement = targetElement;
    }

    put(tokens: any): void {
        super.put(tokens);
    }

    on_finalized_text(text: string, stream_end: boolean): void {
        this.currentText += text;
        this.targetElement.textContent = this.currentText;

        // Auto-scroll as text streams in
        chatContainer.scrollTop = chatContainer.scrollHeight;
    }
}

/**
 * Markdown-enabled streamer that uses StreamingMarkdownRenderer
 * Handles <think> tags and renders markdown in real-time
 */
class MarkdownUIStreamer extends TextStreamer {
    private renderer: StreamingMarkdownRenderer;

    constructor(tokenizer: any, renderer: StreamingMarkdownRenderer, options: any = {}) {
        super(tokenizer, options);
        this.renderer = renderer;
    }

    put(tokens: any): void {
        super.put(tokens);
    }

    on_finalized_text(text: string, stream_end: boolean): void {
        // Push text to the markdown renderer
        this.renderer.push(text);

        // If stream ended, flush the renderer
        if (stream_end) {
            this.renderer.flush();
        }
    }
}

// ============================================
// STATE & MODELS
// ============================================

const MODEL_CONFIGS = {
    qwen: { id: "onnx-community/Qwen3-0.6B-ONNX", name: "Qwen3" },
    classifier: { id: "vmanvs/halugate-sentinel-onnx", name: "Classifier" }
};

const state = {
    models: {
        qwen: null as any,
        classifier: null as any
    },
    tokenizers: {
        qwen: null as any
    },
    isGenerating: false,
    hasSearchCapability: false,
    mode: 'basic' as 'basic' | 'full' | 'orchestrated'
};

// ============================================
// MODEL LOADING
// ============================================

async function loadCausalModel(
    key: 'qwen',
    modelId: string,
    startPercent: number,
    endPercent: number
) {
    const range = endPercent - startPercent;
    console.log(`Loading ${key} model: ${modelId}...`);

    try {
        // 1. Load Tokenizer
        updateProgress(startPercent + (range * 0.1), `Loading ${key} tokenizer...`);
        state.tokenizers[key] = await AutoTokenizer.from_pretrained(modelId);
        console.log(`✓ ${key} Tokenizer loaded`);

        // 2. Load Model
        updateProgress(startPercent + (range * 0.2), `Downloading ${key} model...`);

        // Timer for shader compilation message
        let shaderTimer: number | null = null;
        const startShaderTimer = () => {
            if (shaderTimer) clearTimeout(shaderTimer);
            shaderTimer = window.setTimeout(() => {
                updateProgress(endPercent - 5, `Compiling Shaders for ${key}...`);
            }, 5000);
        };

        state.models[key] = await AutoModelForCausalLM.from_pretrained(modelId, {
            device: "webgpu",
            dtype: "q4f16",
            progress_callback: (progress: any) => {
                const subProgress = progress.progress || 0; // 0-100 for this step
                // Map download progress to range [0.2 -> 0.8] of this model's total allocation
                const effectivePercent = startPercent + (range * 0.2) + (subProgress / 100) * (range * 0.6);

                if (progress.status === "download") {
                    // updateProgress(effectivePercent, `Downloading ${key}...`);
                } else if (progress.status === "done") {
                    startShaderTimer();
                }
            }
        });

        if (shaderTimer) clearTimeout(shaderTimer);
        console.log(`✓ ${key} Model loaded`);

        // 3. Warmup
        updateProgress(endPercent - 2, `Warming up ${key}...`);
        console.log(`Warming up ${key}...`);

        // Simple warmup
        const warmupInputs = state.tokenizers[key]("a");
        await state.models[key].generate({ ...warmupInputs, max_new_tokens: 1 });
        console.log(`✓ ${key} Warmup complete`);

        updateProgress(endPercent, `${key} Ready`);

    } catch (e) {
        console.error(`Failed to load ${key}:`, e);
        throw e;
    }
}

async function loadClassifierModel(startPercent: number, endPercent: number) {
    console.log("Loading classifier...");
    updateProgress(startPercent, "Loading classifier...");

    try {
        state.models.classifier = await pipeline("text-classification", MODEL_CONFIGS.classifier.id, {
            device: "webgpu",
            dtype: "q4",
            progress_callback: (progress: any) => {
                // Simplified progress for classifier
            }
        });
        console.log("✓ Classifier loaded");
        updateProgress(endPercent, "Classifier Ready");
    } catch (e) {
        console.error("Failed to load classifier:", e);
        throw e;
    }
}

async function initializeBasicMode() {
    try {
        updateStatus("Initializing basic mode...");
        state.mode = 'basic';

        // Load only Qwen (0-100%)
        await loadCausalModel('qwen', MODEL_CONFIGS.qwen.id, 0, 100);

        updateProgress(100, "Ready (Basic Mode)");
        updateStatus("Ready (Basic)", true);

        // No model selector in new UI - uses orchestrated/basic mode automatically

        showChatUI();

    } catch (e: any) {
        console.error("Initialization failed:", e);
        updateStatus("Error");
        loadingMessage.textContent = `Error: ${e.message}`;
    }
}

async function initializeFullMode() {
    try {
        updateStatus("Initializing full mode...");
        state.mode = 'full';
        state.hasSearchCapability = true;

        // 1. Load Qwen (0-70%)
        await loadCausalModel('qwen', MODEL_CONFIGS.qwen.id, 0, 70);

        // 2. Load Classifier (70-100%)
        await loadClassifierModel(70, 100);

        updateProgress(100, "All systems go!");
        updateStatus("Ready (Full)", true);

        // No model selector in new UI - uses orchestrated mode automatically

        showChatUI();

    } catch (e: any) {
        console.error("Initialization failed:", e);
        updateStatus("Error");
        loadingMessage.textContent = `Error: ${e.message}`;
    }
}

// ============================================
// CHAT LOGIC
// ============================================

async function chat(userMessage: string) {
    // Handle orchestrated mode
    if (state.hasSearchCapability) {
        await runOrchestrated(userMessage);
        return;
    }

    // Basic mode - just use Qwen
    const model = state.models.qwen;
    const tokenizer = state.tokenizers.qwen;

    if (!model || !tokenizer) {
        addMessage("Error: Model not loaded.", false);
        return;
    }

    state.isGenerating = true;
    updateStatus('Generating...');

    // Initialize stopping criteria for this generation
    stoppingCriteria = new InterruptableStoppingCriteria();

    try {
        // Format messages
        const messages = [{ role: "user", content: userMessage }];

        const inputs = tokenizer.apply_chat_template(messages, {
            add_generation_prompt: true,
            return_dict: true
        });

        console.log('Generating with Qwen...');

        addTypingIndicator();
        await new Promise(r => setTimeout(r, 300));
        removeTypingIndicator();

        // Use new streaming markdown renderer
        const { renderer } = addStreamingMessage();
        const streamer = new MarkdownUIStreamer(tokenizer, renderer, {
            skip_prompt: true,
            skip_special_tokens: true
        });

        await model.generate({
            ...inputs,
            max_new_tokens: 4096,
            temperature: 0.7,
            top_p: 0.9,
            streamer: streamer,
            stopping_criteria: stoppingCriteria
        });

    } catch (e: any) {
        console.error("Generation error:", e);
        addMessage(`Error: ${e.message}`, false);
    } finally {
        state.isGenerating = false;
        updateStatus('Ready', true);
    }
}

async function runOrchestrated(userMessage: string) {
    state.isGenerating = true;
    updateStatus('Processing...');

    // Initialize stopping criteria for this generation
    stoppingCriteria = new InterruptableStoppingCriteria();

    try {
        addTypingIndicator();

        const models: Models = {
            classifier: state.models.classifier,
            qwen: {
                model: state.models.qwen,
                tokenizer: state.tokenizers.qwen
            }
        };

        // Prepare streamer with markdown support
        removeTypingIndicator();
        const { renderer } = addStreamingMessage();
        const streamer = new MarkdownUIStreamer(state.tokenizers.qwen, renderer, {
            skip_prompt: true,
            skip_special_tokens: true
        });

        // Run pipeline with stopping criteria
        const finalState = await runPipeline(
            userMessage,
            models,
            streamer,
            (pipelineState) => {
                updateStepIndicator(pipelineState.currentStep);
            },
            state.hasSearchCapability,
            stoppingCriteria
        );

        if (finalState.currentStep === 'error') {
            addMessage(`Error: ${finalState.error}`, false);
        }

    } catch (e: any) {
        console.error("Orchestration error:", e);
        removeTypingIndicator();
        addMessage(`Error: ${e.message}`, false);
    } finally {
        state.isGenerating = false;
        updateStepIndicator('idle');
        updateStatus('Ready', true);
    }
}

async function runClassifier(userMessage: string) {
    state.isGenerating = true;
    updateStatus('Classifying...');

    try {
        addTypingIndicator();

        const result = await state.models.classifier(userMessage);

        removeTypingIndicator();

        console.log("Classification result:", result);
        const label = result[0]?.label;
        const score = result[0]?.score;

        addMessage(`🛡️ Classification Result:\nLabel: ${label}\nScore: ${score?.toFixed(4)}`, false);

    } catch (e: any) {
        console.error("Classification error:", e);
        addMessage(`Error: ${e.message}`, false);
    } finally {
        state.isGenerating = false;
        updateStatus('Ready', true);
    }
}


// ============================================
// EVENT HANDLERS
// ============================================

async function handleSendMessage() {
    // If generating, stop the generation
    if (state.isGenerating) {
        if (stoppingCriteria) {
            stoppingCriteria.interrupt();
            stoppingCriteria = null;
        }
        state.isGenerating = false;
        setSendButtonState(false);
        updateStatus('Ready', true);
        return;
    }

    const message = userInput?.value.trim();
    if (!message) return;

    userInput.value = '';
    addMessage(message, true);

    // Clear panels for new conversation turn
    clearProcessPanel();
    clearReasoningPanel();

    state.isGenerating = true;
    setSendButtonState(true);

    try {
        await chat(message);
    } finally {
        state.isGenerating = false;
        setSendButtonState(false);
        userInput?.focus();
    }
}

sendBtn?.addEventListener('click', handleSendMessage);
userInput?.addEventListener('keydown', (e: KeyboardEvent) => {
    if (e.key === 'Enter' && !e.shiftKey) {
        e.preventDefault();
        handleSendMessage();
    }
});

// Auto-resize textarea
userInput?.addEventListener('input', () => {
    if (userInput) {
        userInput.style.height = 'auto';
        userInput.style.height = Math.min(userInput.scrollHeight, 150) + 'px';
    }
});

// Onboarding handlers
enableSearchBtn?.addEventListener('click', () => {
    // Open Chrome Web Store or extension installation page
    window.open('https://chromewebstore.google.com/detail/bognlgebgmapkcjncomjflcppdobeelh?utm_source=item-share-cb', '_blank');
    hideOnboardingModal();
    // Still load basic mode for now
    initializeBasicMode();
});

skipSearchBtn?.addEventListener('click', () => {
    hideOnboardingModal();
    initializeBasicMode();
});

// ============================================
// MAIN ENTRY
// ============================================

(async () => {
    await checkAndRequestGPU();

    // Check if extension is installed
    // Give content script time to inject marker
    await new Promise(r => setTimeout(r, 100));

    const extensionInstalled = hasExtensionMarker() || await detectExtension();

    if (extensionInstalled) {
        console.log('Extension detected! Loading full mode...');
        state.hasSearchCapability = true;
        hideOnboardingModal();
        await initializeFullMode();
    } else {
        console.log('Extension not detected. Showing onboarding...');
        showOnboardingModal();
    }
})();