offscreen.js

/**
 * Offscreen Document for AI API Access
 * Required to access window.LanguageModel, window.Summarizer, window.Writer, window.Rewriter, window.Translator, window.LanguageDetector, window.Proofreader APIs from service worker
 * 
 * Proofreader API: https://developer.chrome.com/docs/ai/proofreader-api
 * The Proofreader API provides grammar and spelling correction functionality
 */

const LANGUAGE_MODEL_SYSTEM_PROMPT = 'You are a helpful AI assistant integrated into a browser extension. Provide concise, accurate, and helpful responses.'
const SMART_SUGGESTION_SYSTEM_PROMPT = 'You are an intelligent browser assistant that generates helpful, context-aware action suggestions.'

// Supported languages by Chrome's Prompt API
const SUPPORTED_LANGUAGES = ['en', 'es', 'ja']

// Cache user preferences to avoid repeated message passing
let cachedUserPreferences = null;

let defaultLanguageModelSession = null
let defaultLanguageModelInitPromise = null
let smartSuggestionSession = null
let smartSuggestionInitPromise = null
let summarizerInstance = null
let summarizerInitPromise = null
let writerInstance = null
let writerInitPromise = null
let languageDetectorInstance = null
let languageDetectorInitPromise = null
let proofreaderInstance = null
let proofreaderInitPromise = null

const translatorSessions = new Map()
const rewriterSessions = new Map()

// Handle messages from service worker
chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
    if (message.target !== 'offscreen') {
        return false;
    }

    handleAIRequest(message).then(result => {
        sendResponse(result);
    }).catch(error => {
        console.error('Error handling AI request:', error);
        sendResponse({ success: false, error: error.message });
    });

    return true; // Keep channel open for async response
});

/**
 * Route AI requests to appropriate handlers
 */
async function handleAIRequest(message) {
    switch (message.type) {
        case 'PING':
            return { pong: true };
        
        case 'CHECK_AI_AVAILABILITY':
            return await checkAIAvailability();
        
        case 'UPDATE_USER_PREFERENCES':
            // Cache preferences sent from background script
            if (message.data && message.data.preferences) {
                cachedUserPreferences = message.data.preferences;
                // Reset sessions so they can be recreated with new language
                resetDefaultLanguageModelSession();
                resetSmartSuggestionSession();
            }
            return { success: true };
        
        case 'AI_SUMMARIZE':
            return await summarizeContent(message.data.text);
        
        case 'AI_TRANSLATE':
            return await translateText(message.data.text, message.data.targetLanguage, message.data.sourceLanguage);
        
        case 'AI_REWRITE':
            return await rewriteText(message.data.text, message.data.tone, message.data.writingStyle);
        
        case 'AI_PROOFREAD':
            return await proofreadText(message.data.text, message.data.writingStyle);
        
        case 'AI_COMPLETE':
            return await completeText(message.data.text, message.data.prefix, message.data.writingStyle);
        
        case 'AI_EXTRACT':
            return await extractKeyPoints(message.data.text, message.data.writingStyle);
        
        case 'AI_EXPLAIN':
            return await explainContent(message.data.text, message.data.writingStyle);
        
        case 'AI_SYNTHESIZE':
            return await synthesizeTabs(message.data.prompt, message.data.tabContexts, message.data.draftText, message.data.writingStyle);
        
        case 'AI_SMART_SUGGESTIONS':
            return await generateSmartSuggestions(message.data.prompt);
        
        default:
            return { success: false, error: 'Unknown message type' };
    }
}

// Generic instance manager factory
function createInstanceManager(instanceVar, promiseVar) {
    return {
        ensure: async (createFn) => {
            if (instanceVar.value) return instanceVar.value;
            
            if (!promiseVar.value) {
                promiseVar.value = createFn()
                    .then(instance => {
                        instanceVar.value = instance;
                        return instance;
                    })
                    .catch(error => {
                        instanceVar.value = null;
                        throw error;
                    })
                    .finally(() => {
                        promiseVar.value = null;
                    });
            }
            
            return instanceVar.value || promiseVar.value;
        },
        reset: () => {
            instanceVar.value = null;
        }
    };
}

// Wrapper objects to allow pass-by-reference
const defaultLMSession = { value: defaultLanguageModelSession };
const defaultLMPromise = { value: defaultLanguageModelInitPromise };
const smartSugSession = { value: smartSuggestionSession };
const smartSugPromise = { value: smartSuggestionInitPromise };
const summarizerInst = { value: summarizerInstance };
const summarizerProm = { value: summarizerInitPromise };
const writerInst = { value: writerInstance };
const writerProm = { value: writerInitPromise };
const langDetectorInst = { value: languageDetectorInstance };
const langDetectorProm = { value: languageDetectorInitPromise };
const proofreaderInst = { value: proofreaderInstance };
const proofreaderProm = { value: proofreaderInitPromise };

const defaultLMManager = createInstanceManager(defaultLMSession, defaultLMPromise);
const smartSugManager = createInstanceManager(smartSugSession, smartSugPromise);
const summarizerManager = createInstanceManager(summarizerInst, summarizerProm);
const writerManager = createInstanceManager(writerInst, writerProm);
const langDetectorManager = createInstanceManager(langDetectorInst, langDetectorProm);
const proofreaderManager = createInstanceManager(proofreaderInst, proofreaderProm);

/**
 * Get user's preferred language from storage
 * @returns {Promise<string>} Preferred language code (defaults to 'en')
 */
async function getUserPreferredLanguage() {
    try {
        // Use cached preferences if available
        if (cachedUserPreferences && cachedUserPreferences.preferredLanguage) {
            const preferredLang = cachedUserPreferences.preferredLanguage;
            return SUPPORTED_LANGUAGES.includes(preferredLang) ? preferredLang : 'en';
        }
        
        // Request preferences from background script if not cached
        const response = await chrome.runtime.sendMessage({
            type: 'GET_USER_PREFERENCES'
        });
        
        if (response && response.preferences) {
            cachedUserPreferences = response.preferences;
            const preferredLang = response.preferences.preferredLanguage || 'en';
            return SUPPORTED_LANGUAGES.includes(preferredLang) ? preferredLang : 'en';
        }
        
        return 'en';
    } catch (error) {
        console.error('Error getting user preferred language:', error);
        return 'en';
    }
}

async function ensureDefaultLanguageModelSession() {
    const preferredLang = await getUserPreferredLanguage();
    
    const session = await defaultLMManager.ensure(() => window.LanguageModel.create({
        systemPrompt: LANGUAGE_MODEL_SYSTEM_PROMPT,
        language: preferredLang,
        expectedInputs: [
            { type: "text", languages: SUPPORTED_LANGUAGES }
        ],
        expectedOutputs: [
            { type: "text", languages: [preferredLang] }
        ]
    }));
    defaultLanguageModelSession = defaultLMSession.value;
    return session;
}

function resetDefaultLanguageModelSession() {
    defaultLMManager.reset();
    defaultLanguageModelSession = null;
}

async function ensureSmartSuggestionSession() {
    const preferredLang = await getUserPreferredLanguage();
    
    const session = await smartSugManager.ensure(() => window.LanguageModel.create({
        systemPrompt: SMART_SUGGESTION_SYSTEM_PROMPT,
        temperature: 0.7,
        topK: 8,
        language: preferredLang,
        expectedInputs: [
            { type: "text", languages: SUPPORTED_LANGUAGES }
        ],
        expectedOutputs: [
            { type: "text", languages: [preferredLang] }
        ]
    }));
    smartSuggestionSession = smartSugSession.value;
    return session;
}

function resetSmartSuggestionSession() {
    smartSugManager.reset();
    smartSuggestionSession = null;
}

async function runWithLanguageModelSession(ensureFn, resetFn, executor) {
    try {
        const session = await ensureFn();
        return await executor(session);
    } catch (error) {
        resetFn();
        try {
            const session = await ensureFn();
            return await executor(session);
        } catch (retryError) {
            resetFn();
            throw retryError;
        }
    }
}

async function ensureSummarizerInstance() {
    const instance = await summarizerManager.ensure(() => window.Summarizer.create({
        type: 'key-points',
        format: 'markdown',
        length: 'medium'
    }));
    summarizerInstance = summarizerInst.value;
    return instance;
}

function resetSummarizerInstance() {
    summarizerManager.reset();
    summarizerInstance = null;
}

async function ensureWriterInstance() {
    const instance = await writerManager.ensure(() => window.Writer.create());
    writerInstance = writerInst.value;
    return instance;
}

function resetWriterInstance() {
    writerManager.reset();
    writerInstance = null;
}

async function ensureLanguageDetectorInstance() {
    const instance = await langDetectorManager.ensure(() => window.LanguageDetector.create());
    languageDetectorInstance = langDetectorInst.value;
    return instance;
}

function resetLanguageDetectorInstance() {
    langDetectorManager.reset();
    languageDetectorInstance = null;
}

async function ensureProofreaderInstance() {
    const instance = await proofreaderManager.ensure(() => window.Proofreader.create());
    proofreaderInstance = proofreaderInst.value;
    return instance;
}

function resetProofreaderInstance() {
    proofreaderManager.reset();
    proofreaderInstance = null;
}

function getRewriterKey(tone) {
    return tone || 'professional';
}

async function ensureRewriterInstance(tone) {
    const resolvedTone = tone || 'professional';
    const key = getRewriterKey(resolvedTone);
    const cached = rewriterSessions.get(key);

    if (cached) {
        return cached;
    }

    const creationPromise = window.Rewriter.create({ tone: resolvedTone }).then(instance => {
        rewriterSessions.set(key, instance);
        return instance;
    }).catch(error => {
        rewriterSessions.delete(key);
        throw error;
    });

    rewriterSessions.set(key, creationPromise);
    return creationPromise;
}

function resetRewriterInstance(tone) {
    rewriterSessions.delete(getRewriterKey(tone));
}

function getTranslatorKey(sourceLanguage, targetLanguage) {
    return `${sourceLanguage || 'auto'}->${targetLanguage}`;
}

async function ensureTranslatorInstance(sourceLanguage, targetLanguage) {
    const key = getTranslatorKey(sourceLanguage, targetLanguage);
    const cached = translatorSessions.get(key);

    if (cached) {
        return cached;
    }

    const creationPromise = window.Translator.create({
        sourceLanguage,
        targetLanguage
    }).then(instance => {
        translatorSessions.set(key, instance);
        return instance;
    }).catch(error => {
        translatorSessions.delete(key);
        throw error;
    });

    translatorSessions.set(key, creationPromise);
    return creationPromise;
}

function resetTranslatorInstance(sourceLanguage, targetLanguage) {
    translatorSessions.delete(getTranslatorKey(sourceLanguage, targetLanguage));
}

function warmAIResources(status) {
    const warmupTasks = [
        { condition: status.prompt === 'available' && window.LanguageModel, 
          tasks: [ensureDefaultLanguageModelSession, ensureSmartSuggestionSession] },
        { condition: status.summarizer === 'available' && window.Summarizer, 
          tasks: [ensureSummarizerInstance] },
        { condition: status.writer === 'available' && window.Writer, 
          tasks: [ensureWriterInstance] },
        { condition: status.rewriter === 'available' && window.Rewriter, 
          tasks: [() => ensureRewriterInstance('professional')] },
        { condition: status.languageDetector === 'available' && window.LanguageDetector, 
          tasks: [ensureLanguageDetectorInstance] },
        { condition: status.proofreader === 'available' && window.Proofreader, 
          tasks: [ensureProofreaderInstance] }
    ];

    setTimeout(() => {
        warmupTasks.forEach(({ condition, tasks }) => {
            if (condition) tasks.forEach(task => task().catch(() => {}));
        });
    }, 0);
}

/**
 * Check AI API availability
 * Following Chrome's official built-in AI documentation
 * Returns: "available", "downloadable", "downloading", or "unavailable"
 */
async function checkAIAvailability() {
    const status = {
        prompt: 'unavailable',
        writer: 'unavailable',
        summarizer: 'unavailable',
        rewriter: 'unavailable',
        translator: 'unavailable',
        languageDetector: 'unavailable',
        proofreader: 'unavailable'
    };

    const apiChecks = [
        { key: 'prompt', api: window.LanguageModel },
        { key: 'summarizer', api: window.Summarizer },
        { key: 'writer', api: window.Writer },
        { key: 'rewriter', api: window.Rewriter },
        { key: 'translator', api: window.Translator, 
          params: { sourceLanguage: 'en', targetLanguage: 'es' } },
        { key: 'languageDetector', api: window.LanguageDetector },
        { key: 'proofreader', api: window.Proofreader }
    ];

    try {
        for (const { key, api, params } of apiChecks) {
            if (api) {
                status[key] = await api.availability(params);
            }
        }
    } catch (error) {
        console.error('Error checking AI availability:', error);
    }

    warmAIResources(status);

    return status;
}

async function streamResponse(streamFn, action) {
    let fullResult = '';
    const stream = streamFn();
    
    for await (const chunk of stream) {
        fullResult = chunk.length > fullResult.length ? chunk : fullResult + chunk;
        
        chrome.runtime.sendMessage({
            type: 'AI_STREAM_CHUNK',
            action,
            chunk: fullResult,
            done: false
        }).catch(() => {});
    }

    chrome.runtime.sendMessage({
        type: 'AI_STREAM_CHUNK',
        action,
        chunk: fullResult,
        done: true
    }).catch(() => {});

    return fullResult;
}

/**
 * Summarize content using AI with streaming support
 */
async function summarizeContent(text) {
    try {
        if (window.Summarizer) {
            const availability = await window.Summarizer.availability();
            
            if (availability === 'available') {
                try {
                    const summarizer = await ensureSummarizerInstance();

                    if (summarizer.summarizeStreaming) {
                        const result = await streamResponse(
                            () => summarizer.summarizeStreaming(text),
                            'summarize'
                        );
                        return { success: true, result, streamed: true };
                    }

                    const summary = await summarizer.summarize(text);
                    return { success: true, result: summary };
                } catch (summarizerError) {
                    console.error('Summarizer API error:', summarizerError);
                    resetSummarizerInstance();
                }
            }
        }
        
        return await usePromptAPIStreaming(
            `Summarize the following text in bullet points:\n\n${text.slice(0, 4000)}`,
            'summarize'
        );
    } catch (error) {
        return { success: false, error: error.message };
    }
}

async function translateText(text, targetLanguage, sourceLanguage = null) {
    try {
        // If sourceLanguage not provided, try to detect it
        let detectedSourceLang = sourceLanguage;
        
        if (!detectedSourceLang && window.LanguageDetector) {
            try {
                const detectorAvailability = await window.LanguageDetector.availability();
                
                if (detectorAvailability === 'available' || detectorAvailability === 'downloadable') {
                    const detector = await ensureLanguageDetectorInstance();
                    const results = await detector.detect(text);
                    
                    if (results && results.length > 0 && results[0].confidence > 0.5) {
                        detectedSourceLang = results[0].detectedLanguage;
                    }
                }
            } catch (detectionError) {
                resetLanguageDetectorInstance();
            }
        }
        
        // If still no source language detected, use a default based on context
        if (!detectedSourceLang) {
            // Default to common languages - try English first
            detectedSourceLang = 'en';
        }
        
        // Try Chrome's Translator API
        if (window.Translator && detectedSourceLang) {
            try {
                const availability = await window.Translator.availability({
                    sourceLanguage: detectedSourceLang,
                    targetLanguage: targetLanguage
                });
                
                if (availability === 'available' || availability === 'downloadable') {
                    try {
                        const translator = await ensureTranslatorInstance(detectedSourceLang, targetLanguage);
                        const translated = await translator.translate(text);
                        return { success: true, result: translated };
                    } catch (translatorError) {
                        console.error('Translator API error:', translatorError);
                        resetTranslatorInstance(detectedSourceLang, targetLanguage);
                    }
                }
            } catch (translatorAvailabilityError) {
                console.error('Translator availability error:', translatorAvailabilityError);
            }
        }
        
        // Fallback to Prompt API
        const languageNames = {
            'en': 'English',
            'es': 'Spanish',
            'fr': 'French',
            'de': 'German',
            'it': 'Italian',
            'pt': 'Portuguese',
            'ru': 'Russian',
            'ja': 'Japanese',
            'ko': 'Korean',
            'zh': 'Chinese'
        };
        
        const targetLangName = languageNames[targetLanguage] || targetLanguage;

        return await usePromptAPIStreaming(
            `Translate the following text to ${targetLangName}:\n\n${text.slice(0, 4000)}`
        );
    } catch (error) {
        return { success: false, error: error.message };
    }
}

async function rewriteText(text, tone, writingStyle) {
    try {
        if (window.Rewriter) {
            const availability = await window.Rewriter.availability();
            
            if (availability === 'available' || availability === 'downloadable') {
                try {
                    const rewriter = await ensureRewriterInstance(tone);
                    const rewritten = await rewriter.rewrite(text);
                    return { success: true, result: rewritten };
                } catch (rewriterError) {
                    console.error('Rewriter API error:', rewriterError);
                    resetRewriterInstance(tone);
                }
            }
        }

        // Fallback to Prompt API with writing style
        const style = writingStyle || 'professional';
        const styleInstruction = style !== 'auto' ? ` Ensure the text follows a ${style} writing style.` : '';
        
        return await usePromptAPIStreaming(
            `Rewrite the following text in a ${tone} tone.${styleInstruction}\n\n${text}`,
            'rewrite'
        );
    } catch (error) {
        return { success: false, error: error.message };
    }
}

async function proofreadText(text, writingStyle) {
    try {
        // Try Chrome's built-in Proofreader API first
        if (window.Proofreader) {
            const availability = await window.Proofreader.availability();
            
            if (availability === 'available' || availability === 'downloadable') {
                try {
                    const proofreader = await ensureProofreaderInstance();
                    
                    // Call proofread() to get corrections for the input text
                    const proofreadResult = await proofreader.proofread(text);
                    
                    // proofreadResult contains:
                    // - correction: the fully corrected text
                    // - corrections: array of ProofreadCorrection objects with startIndex, endIndex, replacement
                    
                    // Format the result for display with detailed corrections
                    let resultText = proofreadResult.correction;
                    
                    // If there are corrections, add them as metadata
                    if (proofreadResult.corrections && proofreadResult.corrections.length > 0) {
                        const correctionsSummary = proofreadResult.corrections.map((correction, index) => {
                            const originalText = text.substring(correction.startIndex, correction.endIndex);
                            return `${index + 1}. "${originalText}" → "${correction.replacement}"${correction.type ? ` (${correction.type})` : ''}`;
                        }).join('\n');
                        
                        resultText = `${proofreadResult.correction}\n\n---\nCorrections made (${proofreadResult.corrections.length}):\n${correctionsSummary}`;
                    }
                    
                    return { 
                        success: true, 
                        result: resultText,
                        proofreadResult: proofreadResult, // Include raw result for advanced use
                        usedProofreaderAPI: true
                    };
                } catch (proofreaderError) {
                    console.error('Proofreader API error:', proofreaderError);
                    resetProofreaderInstance();
                    // Fall through to Prompt API fallback
                }
            }
        }
        
        // Fallback to Prompt API if Proofreader is not available
        const style = writingStyle || 'professional';
        const styleInstruction = style !== 'auto' ? ` Ensure the text follows a ${style} writing style.` : '';
        
        return await usePromptAPIStreaming(
            `Proofread and correct any grammar, spelling, or punctuation errors in the following text.${styleInstruction}\n\n${text}`,
            'proofread'
        );
    } catch (error) {
        return { success: false, error: error.message };
    }
}

async function completeText(text, prefix, writingStyle) {
    try {
        if (window.Writer) {
            const availability = await window.Writer.availability();
            
            if (availability === 'available' || availability === 'downloadable') {
                try {
                    const writer = await ensureWriterInstance();
                    const completion = await writer.write(prefix + text);
                    return { success: true, result: completion };
                } catch (writerError) {
                    console.error('Writer API error:', writerError);
                    resetWriterInstance();
                }
            }
        }

        // Fallback to Prompt API with writing style
        const style = writingStyle || 'professional';
        const styleInstruction = style !== 'auto' ? ` Ensure the completion follows a ${style} writing style.` : '';
        
        return await usePromptAPIStreaming(
            `Complete the following text naturally.${styleInstruction}\n\n${prefix}${text}`,
            'complete'
        );
    } catch (error) {
        return { success: false, error: error.message };
    }
}

async function extractKeyPoints(text, writingStyle) {
    const style = writingStyle || 'professional';
    const styleInstruction = style !== 'auto' ? ` Present the key points in a ${style} writing style.` : '';
    
    return await usePromptAPIStreaming(
        `Extract the main key points and insights from the following text.${styleInstruction}\n\n${text.slice(0, 4000)}`,
        'extract'
    );
}

async function explainContent(text, writingStyle) {
    const style = writingStyle || 'professional';
    const styleInstruction = style !== 'auto' ? ` Use a ${style} writing style in your explanation.` : '';
    
    return await usePromptAPIStreaming(
        `Explain the following in simple, easy-to-understand terms.${styleInstruction}\n\n${text.slice(0, 4000)}`,
        'explain'
    );
}

async function synthesizeTabs(prompt, tabContexts, draftText, writingStyle) {
    try {
        // Add writing style instruction to the prompt
        const style = writingStyle || 'professional';
        const styleInstruction = style !== 'auto' ? `\n\nImportant: Use a ${style} writing style in your synthesis.` : '';
        
        // Use streaming for better UX
        return await usePromptAPIStreaming(prompt + styleInstruction, 'synthesize');
    } catch (error) {
        return { success: false, error: error.message };
    }
}

async function usePromptAPI(prompt) {
    try {
        if (!window.LanguageModel) {
            return { 
                success: false, 
                error: 'AI language model not available. Please ensure Chrome has Gemini Nano enabled.' 
            };
        }

        const availability = await window.LanguageModel.availability();
        
        if (availability !== 'available') {
            return { 
                success: false, 
                error: `AI language model status: ${availability}. Please check Chrome AI settings.` 
            };
        }

        const result = await runWithLanguageModelSession(
            ensureDefaultLanguageModelSession,
            resetDefaultLanguageModelSession,
            async (session) => session.prompt(prompt)
        );

        return { success: true, result };
    } catch (error) {
        console.error('Prompt API error:', error);
        return { success: false, error: error.message };
    }
}

async function usePromptAPIStreaming(prompt, action) {
    try {
        if (!window.LanguageModel) {
            return { 
                success: false, 
                error: 'AI language model not available. Please ensure Chrome has Gemini Nano enabled.' 
            };
        }

        const availability = await window.LanguageModel.availability();
        
        if (availability !== 'available') {
            return { 
                success: false, 
                error: `AI language model status: ${availability}. Please check Chrome AI settings.` 
            };
        }

        const { result, streamed } = await runWithLanguageModelSession(
            ensureDefaultLanguageModelSession,
            resetDefaultLanguageModelSession,
            async (session) => {
                if (session.promptStreaming) {
                    const result = await streamResponse(
                        () => session.promptStreaming(prompt),
                        action
                    );
                    return { result, streamed: true };
                }

                const response = await session.prompt(prompt);
                return { result: response, streamed: false };
            }
        );

        return { success: true, result, ...(streamed && { streamed: true }) };
    } catch (error) {
        console.error('Prompt API streaming error:', error);
        return { success: false, error: error.message };
    }
}

async function generateSmartSuggestions(prompt) {
    try {
        if (!window.LanguageModel) {
            return { 
                success: false, 
                error: 'AI language model not available. Please ensure Chrome has Gemini Nano enabled.' 
            };
        }

        const availability = await window.LanguageModel.availability();
        
        if (availability !== 'available') {
            return { 
                success: false, 
                error: `AI language model status: ${availability}. Please check Chrome AI settings.` 
            };
        }

        const suggestionSchema = {
            type: 'array',
            items: {
                type: 'object',
                properties: {
                    type: {
                        type: 'string',
                        enum: ['summarize', 'translate', 'rewrite', 'proofread', 'complete', 'extract', 'explain', 'compare', 'synthesize']
                    },
                    title: {
                        type: 'string',
                        maxLength: 40
                    },
                    description: {
                        type: 'string',
                        maxLength: 60
                    },
                    priority: {
                        type: 'string',
                        enum: ['high', 'medium', 'low']
                    }
                },
                required: ['type', 'title', 'description', 'priority']
            },
            minItems: 1,
            maxItems: 3
        };
        
        const result = await runWithLanguageModelSession(
            ensureSmartSuggestionSession,
            resetSmartSuggestionSession,
            async (session) => session.prompt(prompt, {
                responseConstraint: suggestionSchema,
                omitResponseConstraintInput: false
            })
        );
        
        return { success: true, result };
    } catch (error) {
        console.error('Smart suggestions generation error:', error);
        return { success: false, error: error.message };
    }
}