rentgen/memory.ts

import ExtendedRequest from './extended-request';
import { getshorthost } from './util';
import { RequestCluster } from './request-cluster';
import { SaferEmitter } from './safer-emitter';
import browserAPI from './lib/browser-api';

// Deklaracja Chrome API dla TypeScript
declare const chrome: any;

function setDomainsCount(counter: number, tabId: number) {
    // Ochrona przed próbą ustawienia badge dla zamkniętej zakładki
    try {
        browserAPI.badge.setBadgeText({ text: counter < 0 ? '0' : counter.toString(), tabId });
        browserAPI.badge.setTitle({
            title: 'Rentgen',
            tabId,
        });
    } catch (e) {
        // Zakładka została zamknięta - ignorujemy błąd
        console.debug(`Zakładka ${tabId} już nie istnieje, pomijanie aktualizacji badge`);
    }
}

// Cachowany RequestCluster dla popupu Chrome (ma metody, ale używa danych z cache)
class CachedRequestCluster extends RequestCluster {
    private _hasCookies: boolean = false;
    private _exposesOrigin: boolean = false;
    private _hasMarks: boolean = false;
    public lastFullUrl: string | null = null;
    public lastModified: number = 0;
    
    constructor(id: string, cached: any) {
        super(id);
        this._hasCookies = cached.hasCookies || false;
        this._exposesOrigin = cached.exposesOrigin || false;
        this._hasMarks = cached.hasMarks || false;
        this.lastFullUrl = cached.lastFullUrl || null;
        this.lastModified = cached.lastModified || 0;
        this.requests = []; // Pusta tablica zapobiegająca błędom
    }
    
    hasCookies(): boolean {
        return this._hasCookies;
    }
    
    exposesOrigin(): boolean {
        return this._exposesOrigin;
    }
    
    hasMarks(): boolean {
        return this._hasMarks;
    }
    
    // Automatyczne zaznaczenie dla cachowanego clustra
    autoMark(): void {
        this._hasMarks = true;
        // ✅ Trigger storage sync
        if (process.env.TARGET === 'chrome') {
            const memory = getMemory();
            if (memory instanceof Memory) {
                memory.scheduleSyncToStorage?.();
            }
        }
    }
    
    // Cofnięcie zaznaczenia dla cachowanego clustra
    undoMark(): void {
        this._hasMarks = false;
        // ✅ Trigger storage sync
        if (process.env.TARGET === 'chrome') {
            const memory = getMemory();
            if (memory instanceof Memory) {
                memory.scheduleSyncToStorage?.();
            }
        }
    }
    
    // Nadpisanie metody - zwraca pustą tablicę
    calculateRepresentativeStolenData(): any[] {
        return [];
    }
    
    // Nadpisanie metody - zwraca tylko główną domenę
    getFullHosts(): string[] {
        return [this.id];
    }
    
    // Metody wymagane przez report-window
    hasMarkedCookies(): boolean {
        return this._hasCookies && this._hasMarks;
    }
    
    getMarkedRequests(): any[] {
        return [];
    }
    
    getMarkedEntries(): any[] {
        return [];
    }
    
    exposesOriginWhere(): any[] {
        return this._exposesOrigin ? [{ path: '', source: 'cached', key: '' }] : [];
    }
    
    getDataTypeDescription(noun = 'Twojej'): string {
        let types_of_data: string[] = [];
        if (this.exposesOrigin()) {
            types_of_data.push(`część ${noun} historii przeglądania`);
        }
        if (this.hasMarkedCookies()) {
            types_of_data.push('unikalne ID z cookies');
        }
        if (types_of_data.length > 1) {
            types_of_data[types_of_data.length - 1] = 'oraz ' + types_of_data[types_of_data.length - 1];
        }
        return types_of_data.join(', ');
    }
}

export default class Memory extends SaferEmitter {
    origin_to_history = {} as Record<string, Record<string, RequestCluster>>;
    isReady: boolean = true; // Firefox jest zawsze gotowy
    private readyPromise: Promise<void> | null = null;
    private syncScheduled: boolean = false;
    private lastSyncTime: number = 0;
    
    // Chrome: Śledzenie pełnych URLi dla zakładek (do generowania screenshotów)
    private tabUrls: Map<number, string> = new Map();
    
    async register(request: ExtendedRequest) {
        await request.init();
        
        if (!request.isThirdParty()) {
            return;
        }
        if (!this.origin_to_history[request.origin]) {
            this.origin_to_history[request.origin] = {};
        }
        const shorthost = getshorthost(new URL(request.url).host);
        
        let isNewCluster = false;
        if (!this.origin_to_history[request.origin][shorthost]) {
            const cluster = new RequestCluster(shorthost);
            this.origin_to_history[request.origin][shorthost] = cluster;
            isNewCluster = true;
        }
        
        this.origin_to_history[request.origin][shorthost].add(request);
        
        // Chrome: Automatyczne zaznaczanie podejrzanych domen w service workerze
        if (process.env.TARGET === 'chrome' && isNewCluster) {
            // Zaznacz od razu po utworzeniu clustra
            this.origin_to_history[request.origin][shorthost].autoMark();
        }
        
        this.emit('change', shorthost);

        // Owinięcie operacji badge w try-catch
        try {
            Object.values(this.getClustersForOrigin(request.origin)).some((cluster) =>
                cluster.hasCookies()
            )
                ? browserAPI.badge.setBadgeBackgroundColor({ color: '#ff726b' })
                : browserAPI.badge.setBadgeBackgroundColor({ color: '#ffb900' });

            if (request.tabId >= 0) {
                setDomainsCount(
                    Object.values(this.getClustersForOrigin(request.origin)).length,
                    request.tabId
                );
            }
        } catch (e) {
            // Zakładka zamknięta - ignorujemy błędy badge
            console.debug('Aktualizacja badge nie powiodła się - zakładka prawdopodobnie zamknięta');
        }

        // Chrome: Throttlowana synchronizacja do storage (max co 500ms)
        if (process.env.TARGET === 'chrome') {
            this.scheduleSyncToStorage();
        }
    }

    constructor() {
        super();

        browserAPI.webRequest.onBeforeRequest.addListener(
            async (request) => {
                // Chrome: Śledzenie nawigacji main_frame dla pełnego URL
                if (process.env.TARGET === 'chrome' && request.type === 'main_frame' && request.tabId >= 0) {
                    this.tabUrls.set(request.tabId, request.url);
                    console.log(`📍 Zapamiętano URL zakładki ${request.tabId}:`, request.url);
                }
                
                const extReq = new ExtendedRequest(request);
                
                // Chrome: Wstrzyknięcie pełnego URL ze śledzonych zakładek
                if (process.env.TARGET === 'chrome' && request.tabId >= 0) {
                    const fullUrl = this.tabUrls.get(request.tabId);
                    if (fullUrl && !extReq.originalURL) {
                        extReq.originalURL = fullUrl;
                    }
                }
            },
            { urls: ['<all_urls>'] },
            ['requestBody']
        );
        
        browserAPI.webRequest.onBeforeSendHeaders.addListener(
            async (request) => {
                const extendedRequest = ExtendedRequest.by_id[request.requestId].addHeaders(
                    request.requestHeaders || []
                );
                this.register(extendedRequest);
            },
            { urls: ['<all_urls>'] },
            ['requestHeaders']
        );

        // Chrome: Czyszczenie URLi zakładek po zamknięciu
        if (process.env.TARGET === 'chrome') {
            browserAPI.tabs.onRemoved?.addListener?.((tabId: number) => {
                this.tabUrls.delete(tabId);
            });
        }

        // Chrome: Wczytywanie z storage przy starcie (dla popupu)
        if (process.env.TARGET === 'chrome') {
            this.isReady = false;
            this.readyPromise = this.loadFromStorage().then(() => {
                this.isReady = true;
                this.emit('ready');
                console.log('🔵 Memory gotowa do użycia przez popup');
            });
        }
    }

    emit(eventName: string, data = 'any'): boolean {
        setTimeout(() => {
            super.emit(eventName, data);

            // ✅ Sync to storage when marks change (Chrome only)
            if (process.env.TARGET === 'chrome' && eventName === 'change') {
                this.scheduleSyncToStorage();
        }
        }, 0);
        return true;
    }

    getClustersForOrigin(origin: string): Record<string, RequestCluster> {
        return this.origin_to_history[origin] || {};
    }

    // Chrome: Oczekiwanie na zakończenie wczytywania ze storage
    async waitUntilReady(): Promise<void> {
        if (this.isReady) return;
        if (this.readyPromise) {
            await this.readyPromise;
        }
    }

    async removeCookiesFor(origin: string, shorthost?: string): Promise<void> {
        if (shorthost) {
            const cookies = await browserAPI.cookies.getAll({ domain: shorthost });
            for (const cookie of cookies) {
                await browserAPI.cookies.remove({
                    name: cookie.name,
                    url: `https://${cookie.domain}`,
                });
            }
        } else {
            const clusters = this.getClustersForOrigin(origin);

            await Promise.all(
                Object.values(clusters)
                    .filter((cluster) => !shorthost || cluster.id === shorthost)
                    .map((cluster) => this.removeCookiesFor(origin, cluster.id))
            );
        }

        // Chrome: Throttlowana synchronizacja do storage
        if (process.env.TARGET === 'chrome') {
            this.scheduleSyncToStorage();
        }
    }

    async removeRequestsFor(origin: string) {
        this.origin_to_history[origin] = {};

        // Chrome: Throttlowana synchronizacja do storage
        if (process.env.TARGET === 'chrome') {
            this.scheduleSyncToStorage();
        }
    }

    // ========================================
    // CHROME: Metody synchronizacji ze storage (THROTTLED)
    // ========================================

    public scheduleSyncToStorage(): void {
        if (this.syncScheduled) return;
        
        this.syncScheduled = true;
        setTimeout(() => {
            this.syncToStorage();
            this.syncScheduled = false;
        }, 500); // Synchronizacja max co 500ms
    }

    private syncToStorage(): void {
        if (typeof chrome !== 'undefined' && chrome.storage?.session) {
            const now = Date.now();
            // Pomijamy jeśli synchronizowano niedawno (w ciągu 300ms)
            if (now - this.lastSyncTime < 300) {
                return;
            }
            this.lastSyncTime = now;
            
            const serializable: Record<string, any> = {};
            
            for (const [origin, clusters] of Object.entries(this.origin_to_history)) {
                serializable[origin] = {};
                for (const [shorthost, cluster] of Object.entries(clusters)) {
                    // Zapisujemy tylko niezbędne dane dla UI
                    serializable[origin][shorthost] = {
                        id: cluster.id,
                        hasCookies: cluster.hasCookies(),
                        exposesOrigin: cluster.exposesOrigin(),
                        hasMarks: cluster.hasMarks(),
                        requestCount: cluster.requests?.length || 0,
                        lastFullUrl: cluster.lastFullUrl || null,
                        lastModified: cluster.lastModified || 0,
                    };
                }
            }
            
            chrome.storage.session.set({ rentgen_memory: serializable }).catch((err: any) => {
                console.error('Nie udało się zsynchronizować pamięci do storage:', err);
            });
        }
    }

    private async loadFromStorage(): Promise<void> {
        if (typeof chrome !== 'undefined' && chrome.storage?.session) {
            try {
                const result = await chrome.storage.session.get('rentgen_memory');
                if (result.rentgen_memory) {
                    const serialized = result.rentgen_memory;
                    
                    for (const [origin, clusters] of Object.entries(serialized)) {
                        this.origin_to_history[origin] = {};
                        for (const [shorthost, cached] of Object.entries(clusters as Record<string, any>)) {
                            // Tworzenie CachedRequestCluster ze wszystkimi potrzebnymi danymi
                            const clusterData = cached as { 
                                id: string; 
                                hasCookies: boolean; 
                                exposesOrigin: boolean;
                                hasMarks: boolean;
                                requestCount: number;
                                lastFullUrl: string | null;
                                lastModified: number;
                            };
                            this.origin_to_history[origin][shorthost] = new CachedRequestCluster(
                                clusterData.id,
                                clusterData
                            );
                        }
                    }
                    
                    console.log('🔵 Wczytano pamięć z chrome.storage.session:', Object.keys(this.origin_to_history).length, 'origins');
                }
            } catch (err) {
                console.error('Nie udało się wczytać pamięci ze storage:', err);
            }
        }
    }
}

// ========================================
// INICJALIZACJA: Firefox vs Chrome
// ========================================

export function init() {
    const memory = new Memory();

    if (process.env.TARGET === 'chrome') {
        // Chrome MV3: Service worker używa 'self' zamiast 'window'
        (self as any).memory = memory;
        console.log('🔵 Memory zainicjalizowana w service workerze Chrome (self.memory + chrome.storage.session)');
    } else {
        // Firefox: Standardowa strona tła z 'window'
        (window as any).memory = memory;
        console.log('🦊 Memory zainicjalizowana w stronie tła Firefox (window.memory)');
    }
}

// ========================================
// DOSTĘP DO MEMORY: Firefox vs Chrome
// ========================================

// Cachowana instancja memory dla popupu (tylko Chrome)
let popupMemoryInstance: Memory | null = null;

export function getMemory(): Memory {
    if (process.env.TARGET === 'chrome') {
        // Chrome: Najpierw próba pobrania z service workera
        if (typeof self !== 'undefined' && (self as any).memory) {
            // Jesteśmy W service workerze - bezpośredni dostęp
            return (self as any).memory as Memory;
        } else {
            // Jesteśmy w popupie/contencie - tworzymy RAZ i cachujemy
            if (!popupMemoryInstance) {
                console.log('🔵 Tworzenie instancji Chrome memory dla popupu (czytanie z chrome.storage.session)');
                popupMemoryInstance = new Memory();
            }
            return popupMemoryInstance;
        }
    } else {
        // Firefox: Używamy tradycyjnego getBackgroundPage()
        const backgroundPage = browserAPI.extension.getBackgroundPage();
        if (!backgroundPage) {
            throw new Error('Strona tła nie jest dostępna');
        }
        return (backgroundPage.window as any).memory as Memory;
    }
}