document.addEventListener('DOMContentLoaded', () => {
    const state = {
        conversationFiles: [],
        currentIndex: 0,
        identities: {},
        totalConversations: 74 // Based on the file list
    };

    const EPSTEIN_ID = 'e:jeeitunes@gmail.com';
    const EPSTEIN_AVATAR_URL = 'https://m1s5.c20.e2-5.dev/files/jeeitunes/pfp/jeffrey-epstein-pfp.png';

    const prevBtn = document.getElementById('prev-btn');
    const nextBtn = document.getElementById('next-btn');
    const exportBtn = document.getElementById('export-btn');
    const conversationNumberEl = document.getElementById('conversation-number');
    const messageContainer = document.getElementById('message-container');
    const participantsHeader = document.querySelector('.imessage-header .participant-names');
    const avatarsHeader = document.querySelector('.imessage-header .avatars');

    // Modal elements
    const modal = document.getElementById('identity-modal');
    const modalIdentifierEl = document.getElementById('modal-identifier');
    const modalNameInput = document.getElementById('modal-name-input');
    const modalAvatarInput = document.getElementById('modal-avatar-input');
    const modalSaveBtn = document.getElementById('modal-save-btn');
    const closeModalBtn = document.querySelector('.close-btn');

    function init() {
        for (let i = 1; i <= state.totalConversations; i++) {
            state.conversationFiles.push(`data/output_${i}.json`);
        }
        const savedIdentities = localStorage.getItem('speakerIdentities');
        if (savedIdentities) {
            state.identities = JSON.parse(savedIdentities);
        }

        // Hardcode Epstein's identity
        state.identities[EPSTEIN_ID] = { name: 'Jeffrey Epstein', avatar: EPSTEIN_AVATAR_URL, initials: 'JE' };

        prevBtn.addEventListener('click', () => loadConversation(state.currentIndex - 1));
        nextBtn.addEventListener('click', () => loadConversation(state.currentIndex + 1));
        exportBtn.addEventListener('click', exportConversation);
        messageContainer.addEventListener('click', handleAvatarClick);
        closeModalBtn.addEventListener('click', () => modal.style.display = 'none');
        modalSaveBtn.addEventListener('click', saveIdentity);

        loadConversation(0);
    }

    async function loadConversation(index) {
        if (index < 0 || index >= state.totalConversations) {
            return;
        }
        state.currentIndex = index;

        try {
            const res = await fetch(`data/epstein_imessage_${index + 1}.json`);
            if (!res.ok) {
                const fallbackRes = await fetch(`data/output_${index + 1}.json`);
                if (!fallbackRes.ok) throw new Error(`Conversation ${index + 1} not found`);
                const data = await fallbackRes.json();
                renderConversation(data, index + 1);
                return;
            }
            const data = await res.json();
            renderConversation(data, index + 1);
        } catch (error) {
            console.error('Error loading conversation:', error);
            messageContainer.innerHTML = `<p style="text-align:center;">Could not load conversation ${index + 1}.</p>`;
        }
    }

    function parseMessages(rawMessages) {
        const messages = [];
        let tempGroup = [];

        for (const item of rawMessages) {
            if (item.Message && tempGroup.length > 0) {
                const messageObject = tempGroup.reduce((acc, curr) => ({ ...acc, ...curr }), {});
                messages.push(messageObject);
                tempGroup = [];
            }
            tempGroup.push(item);
        }

        if (tempGroup.length > 0) {
            const messageObject = tempGroup.reduce((acc, curr) => ({ ...acc, ...curr }), {});
            messages.push(messageObject);
        }

        return messages.map(msg => ({
            text: msg.Message || '',
            sender: msg.Sender || 'me',
            time: msg.Time || ''
        })).filter(msg => msg.text);
    }

    function getIdentity(senderId) {
        // Explicitly handle Epstein
        if (senderId === EPSTEIN_ID) {
            return state.identities[EPSTEIN_ID];
        }
        // Check for user-defined identities
        if (state.identities[senderId]) {
            return state.identities[senderId];
        }
        // Default for everyone else
        return {
            name: '?',
            avatar: null,
            initials: '?'
        };
    }
    
    function generateColor(str) {
        // Keep a consistent color for the generic "?" avatar
        if (str === '?') return '#8e8e93';

        let hash = 0;
        for (let i = 0; i < str.length; i++) {
            hash = str.charCodeAt(i) + ((hash << 5) - hash);
        }
        let color = '#';
        for (let i = 0; i < 3; i++) {
            const value = (hash >> (i * 8)) & 0xFF;
            color += ('00' + value.toString(16)).substr(-2);
        }
        return color;
    }


    function renderConversation(data, convNumber) {
        const messages = parseMessages(data.Messages);
        messageContainer.innerHTML = '';

        const participants = [...new Set(messages.map(m => m.sender))];
        
        participantsHeader.innerHTML = participants.map(p => getIdentity(p).name).join(', ');

        avatarsHeader.innerHTML = '';
        participants.forEach(p => {
             const identity = getIdentity(p);
             const avatarEl = document.createElement('div');
             avatarEl.className = 'avatar';
             if (identity.avatar) {
                 avatarEl.style.backgroundImage = `url(${identity.avatar})`;
             } else {
                 avatarEl.textContent = identity.initials;
                 avatarEl.style.backgroundColor = generateColor(identity.name); // Use name for color generation
             }
             avatarsHeader.appendChild(avatarEl);
        });

        messages.forEach(msg => {
            const bubble = document.createElement('div');
            bubble.className = 'message-bubble';
            const identity = getIdentity(msg.sender);

            const avatar = document.createElement('div');
            avatar.className = 'avatar';
            avatar.dataset.senderId = msg.sender;
            if (identity.avatar) {
                avatar.style.backgroundImage = `url('${identity.avatar}')`;
            } else {
                avatar.textContent = identity.initials;
                avatar.style.backgroundColor = generateColor(identity.name);
            }

            const content = document.createElement('div');
            content.className = 'message-content';
            
            const urlRegex = /(\b(https?|ftp|file):\/\/[-A-Z0-9+&@#\/%?=~_|!:,.;]*[-A-Z0-9+&@#\/%=~_|])/ig;
            content.innerHTML = msg.text.replace(urlRegex, url => `<a href="${url}" target="_blank">${url}</a>`);
            
            // Set sent/received based on Epstein's ID
            bubble.classList.add(msg.sender === EPSTEIN_ID ? 'sent' : 'received');

            bubble.appendChild(avatar);
            bubble.appendChild(content);
            messageContainer.appendChild(bubble);
        });

        conversationNumberEl.textContent = `Conversation #${convNumber}`;
        prevBtn.disabled = state.currentIndex === 0;
        nextBtn.disabled = state.currentIndex >= state.totalConversations - 1;
    }

    function exportConversation() {
        const conversationId = state.currentIndex + 1;
        html2canvas(messageContainer, {
            scrollY: -window.scrollY,
            scale: 2
        }).then(canvas => {
            const link = document.createElement('a');
            link.download = `epstein_${conversationId}.png`;
            link.href = canvas.toDataURL('image/png');
            link.click();
        });
    }

    function handleAvatarClick(event) {
        if (event.target.classList.contains('avatar')) {
            const senderId = event.target.dataset.senderId;
            // Prevent editing Epstein's identity
            if (senderId && senderId !== EPSTEIN_ID) {
                openIdentityModal(senderId);
            }
        }
    }

    function openIdentityModal(senderId) {
        const identity = getIdentity(senderId);
        modal.style.display = 'block';
        modalIdentifierEl.textContent = senderId;
        // If the current name is '?', show an empty input
        modalNameInput.value = identity.name !== '?' ? identity.name : '';
        modalAvatarInput.value = identity.avatar || '';
        modalSaveBtn.dataset.senderId = senderId;
    }

    function saveIdentity() {
        const senderId = modalSaveBtn.dataset.senderId;
        const name = modalNameInput.value.trim();
        const avatar = modalAvatarInput.value.trim();

        if (senderId) {
            state.identities[senderId] = {
                name: name || '?',
                avatar: avatar || null,
                initials: (name || '?').substring(0, 2).toUpperCase()
            };
            localStorage.setItem('speakerIdentities', JSON.stringify(state.identities));
            modal.style.display = 'none';
            loadConversation(state.currentIndex); // Re-render
        }
    }

    init();
});