const { generateObject } = require('ai');
const { createGoogleGenerativeAI } = require('@ai-sdk/google');
const { z } = require('zod');

console.log('Content script loaded!');

// Initialize Google AI with stored API key
browser.storage.local.get('geminiApiKey').then(result => {
    const googleAI = createGoogleGenerativeAI({
        apiKey: result.geminiApiKey || 'default-key-if-none'
    });
}).catch(error => {
    console.error('Error retrieving API key:', error);
});

// Reset AI invocation count on page load
browser.storage.local.set({ aiInvocationCount: '0' }).then(() => {
    console.log('AI Invocation Count reset to 0');
}).catch(error => {
    console.error('Error resetting AI Invocation Count:', error);
});

// Function to extract label text for an input element
function extractLabelText(input) {
    // If input has an id and there's a label with matching 'for' attribute
    if (input.id) {
        const label = document.querySelector(`label[for="${input.id}"]`);
        if (label) return label.textContent.trim();
    }

    // Check for wrapping label
    const parentLabel = input.closest('label');
    if (parentLabel) {
        const labelText = parentLabel.textContent.trim();
        // Remove the input's value from the label text if present
        return labelText.replace(input.value, '').trim();
    }

    // Check for aria-label
    if (input.getAttribute('aria-label')) {
        return input.getAttribute('aria-label');
    }

    return '';
}

// Function to get form history from our API
function getFormHistory() {
    return new Promise((resolve, reject) => {
        const xhr = new XMLHttpRequest();
        xhr.open('GET', 'http://localhost:8080/get', true);
        xhr.onload = () => resolve(xhr.responseText);
        xhr.onerror = () => reject(xhr.statusText);
        xhr.send();
    }).then(response => {
        console.log('Form History:', JSON.parse(response));
        return response;
    });
}

// Function to save form data
function saveFormData(key, value) {
    return new Promise((resolve, reject) => {
        const xhr = new XMLHttpRequest();
        xhr.open('POST', 'http://localhost:8080/set', true);
        xhr.setRequestHeader('Content-Type', 'application/json');
        xhr.onload = () => resolve(xhr.responseText);
        xhr.onerror = () => reject(xhr.statusText);
        xhr.send(JSON.stringify({ key, value }));
    });
}

// Function to track AI invocations
async function trackAIInvocation(fn, ...args) {
    // Get current count from storage
    const currentCount = await browser.storage.local.get('aiInvocationCount');
    
    // Increment count
    const newCount = (currentCount.aiInvocationCount || 0) + 1;
    await browser.storage.local.set({ aiInvocationCount: newCount.toString() });
    
    // Log the count
    console.log(`AI Invocation Count: ${newCount}`);
    
    // Call the original function
    try {
        const result = await fn(...args);
        return result;
    } catch (error) {
        console.error('Error in AI invocation:', error);
        throw error;
    }
}

// Wrapper for AI suggestions
async function getAISuggestionsWithTracking(formData, history) {
    return trackAIInvocation(getAISuggestionsImpl, formData, history);
}

// Define the schema for field suggestions
const SuggestionSchema = z.object({
    suggestions: z.array(z.object({
        fieldIdentifier: z.object({
            id: z.string().optional(),
            name: z.string().optional(),
            type: z.string(),
            label: z.string().optional()
        }),
        value: z.string(),
        confidence: z.number().min(0).max(1),
        fillLogic: z.string().describe('JavaScript code that when executed will fill this field with the suggested value')
    }))
});

// Original AI suggestions implementation
async function getAISuggestionsImpl(formData, history) {
    try {
        const model = googleAI('gemini-2.0-flash');
        
        const result = await generateObject({
            model,
            schema: SuggestionSchema,
            schemaName: 'FormFieldSuggestions',
            schemaDescription: 'Suggestions for form field values with JavaScript logic to fill them',
            prompt: `You are a helpful AI assistant that suggests form field values and provides JavaScript logic to fill them. Based on the following form fields and historical data, suggest appropriate values:

Form Fields:
${JSON.stringify(formData, null, 2)}

Historical Data:
${JSON.stringify(history, null, 2)}

For each form field that you can suggest a value for, provide:
1. The field identifier (id, name, type, and label if available)
2. The suggested value
3. A confidence score (0-1) based on how well the suggestion matches the context
4. JavaScript code that when executed will fill this field

The JavaScript code should:
- Use querySelector with the most specific selector possible (id, name, or other unique attributes)
- Handle both direct input fields and contenteditable elements
- Include error handling
- Return true if successful, false if not

Example fillLogic:
try {
    const field = document.querySelector('#email');
    if (field) {
        field.value = 'example@email.com';
        field.dispatchEvent(new Event('input', { bubbles: true }));
        return true;
    }
    return false;
} catch (e) {
    console.error('Error filling field:', e);
    return false;
}

Return suggestions only for fields where you have high confidence in the suggestions.`
        });

        console.log('AI Suggestions:', result.object);
        return result.object;
    } catch (error) {
        if (error.name === 'NoObjectGeneratedError') {
            console.error('Failed to generate valid suggestions:', {
                cause: error.cause,
                text: error.text,
                response: error.response,
                usage: error.usage
            });
        } else {
            console.error('Error getting AI suggestions:', error);
        }
        return { suggestions: [] };
    }
}

// Function to execute fill logic for a suggestion
function executeFillLogic(fillLogic) {
    try {
        // Create a new function from the string and execute it
        return new Function(fillLogic)();
    } catch (error) {
        console.error('Error executing fill logic:', error);
        return false;
    }
}

// Function to get all input elements
async function getAllInputElements() {
    const inputs = document.querySelectorAll('input, textarea, select');
    
    // Convert NodeList to Array and map to get relevant properties
    const inputsArray = Array.from(inputs)
        // Filter out hidden inputs and submit buttons
        .filter(input => input.type !== 'hidden' && input.type !== 'submit')
        .map(input => ({
            type: input.type || 'textarea',
            id: input.id,
            name: input.name,
            className: input.className,
            value: input.value,
            placeholder: input.placeholder,
            tagName: input.tagName.toLowerCase(),
            html: input.outerHTML,
            possibleLabel: extractLabelText(input)
        }));

    return inputsArray;
}

// Listen for messages from the popup
browser.runtime.onMessage.addListener((message, sender, sendResponse) => {
    console.log('Received message:', message);
    
    switch (message.action) {
        case 'detectFields':
            console.log('Detecting fields...');
            getAllInputElements().then(inputs => {
                console.log('Detected form fields:', inputs);
                sendResponse({ success: true, message: 'Fields detected', data: inputs });
            }).catch(error => {
                console.error('Error detecting fields:', error);
                sendResponse({ success: false, message: error.message });
            });
            break;
        
        case 'generateSuggestions':
            console.log('Generating suggestions...');
            // First check if we have cached suggestions
            browser.storage.local.get('cachedSuggestions').then(result => {
                const cachedSuggestions = result.cachedSuggestions;
                if (cachedSuggestions) {
                    console.log('Using cached suggestions');
                    sendResponse({ 
                        success: true, 
                        message: 'Using cached suggestions', 
                        data: JSON.parse(cachedSuggestions) 
                    });
                    return true;
                }
            }).catch(error => {
                console.error('Error retrieving cached suggestions:', error);
            });

            Promise.all([getAllInputElements(), getFormHistory()])
                .then(([formFields, history]) => {
                    console.log('Got form fields and history:', { formFields, history });
                    return getAISuggestionsWithTracking(formFields, history);
                })
                .then(suggestions => {
                    console.log('Generated suggestions:', suggestions);
                    // Cache the suggestions
                    browser.storage.local.set({ cachedSuggestions: JSON.stringify(suggestions) });
                    sendResponse({ success: true, message: 'Suggestions generated', data: suggestions });
                })
                .catch(error => {
                    console.error('Error:', error);
                    sendResponse({ success: false, message: error.message });
                });
            break;

        case 'clearSuggestions':
            console.log('Clearing cached suggestions');
            browser.storage.local.remove('cachedSuggestions');
            sendResponse({ success: true, message: 'Suggestions cleared' });
            break;

        case 'executeFillLogic':
            console.log('Executing fill logic:', message.fillLogic);
            try {
                const success = executeFillLogic(message.fillLogic);
                
                if (success) {
                    // Update cached suggestions to mark this one as applied
                    browser.storage.local.get('cachedSuggestions').then(result => {
                        const cached = result.cachedSuggestions;
                        if (cached) {
                            const updatedSuggestions = {
                                ...cached,
                                suggestions: cached.suggestions.map(s => 
                                    s.fillLogic === message.fillLogic 
                                        ? { ...s, applied: true }
                                        : s
                                )
                            };
                            browser.storage.local.set({ cachedSuggestions: JSON.stringify(updatedSuggestions) });
                        }
                    }).catch(error => {
                        console.error('Error updating cached suggestions:', error);
                    });
                }
                
                sendResponse({ 
                    success: true, 
                    message: success ? 'Field filled successfully' : 'Failed to fill field' 
                });
            } catch (error) {
                console.error('Error executing fill logic:', error);
                sendResponse({ success: false, message: error.message });
            }
            break;
        
        default:
            console.log('Unknown action:', message.action);
            sendResponse({ success: false, message: 'Unknown action' });
    }
    
    // Required for async response
    return true;
});