import React, { useState, useEffect } from 'react';
import { createRoot } from 'react-dom/client';
import html2canvas from 'html2canvas';

const transformerImageDataUri = '/0_WGCl24jT_rWfgZWL.png';

const App = () => {
    // Component definitions for the Transformer architecture
    const components = [
        { id: 'inputs', label: 'Inputs', side: 'encoder' },
        { id: 'input-embedding', label: 'Input Embedding', side: 'encoder' },
        { id: 'positional-encoding', label: 'Positional Encoding', side: 'encoder' },
        { id: 'multi-head-attention', label: 'Multi-Head Attention', side: 'encoder' },
        { id: 'add-norm', label: 'Add & Norm', side: 'encoder' },
        { id: 'feedforward', label: 'Feedforward', side: 'encoder' },
        { id: 'outputs', label: 'Outputs (shifted right)', side: 'decoder' },
        { id: 'output-embedding', label: 'Output Embedding', side: 'decoder' },
        { id: 'decoder-positional-encoding', label: 'Positional Encoding', side: 'decoder' },
        { id: 'masked-multi-head-attention', label: 'Masked Multi-Head Attention', side: 'decoder' },
        { id: 'decoder-add-norm', label: 'Add & Norm', side: 'decoder' },
        { id: 'cross-attention', label: 'Cross Attention', side: 'decoder' },
        { id: 'decoder-feedforward', label: 'Feedforward', side: 'decoder' },
        { id: 'linear', label: 'Linear', side: 'decoder' },
        { id: 'softmax', label: 'Softmax', side: 'decoder' },
        { id: 'output-probabilities', label: 'Output Probabilities', side: 'decoder' },
    ];

    const [notes, setNotes] = useState<{ [key: string]: string }>({});
    const [showToast, setShowToast] = useState(false);
    const [name, setName] = useState('');

    // Load notes from localStorage on initial render
    useEffect(() => {
        try {
            const savedNotes = localStorage.getItem('transformer-notes');
            if (savedNotes) {
                setNotes(JSON.parse(savedNotes));
            }
        } catch (error) {
            console.error("Failed to load notes from localStorage", error);
        }
    }, []);

    const handleNoteChange = (id: string, value: string) => {
        setNotes((prev: { [key: string]: string }) => ({ ...prev, [id]: value }));
    };

    const handleSaveNotes = async () => {
        const appContainer = document.querySelector('.app-container') as HTMLElement;
        if (!appContainer) return;
        try {
            const canvas = await html2canvas(appContainer, { backgroundColor: '#121212' });
            const image = canvas.toDataURL('image/png');
            const link = document.createElement('a');
            link.href = image;
            link.download = 'transformer-notes.png';
            link.click();
        } catch (error) {
            console.error('Failed to capture screenshot', error);
        }
    };

    const encoderComponents = components.filter(c => c.side === 'encoder');
    const decoderComponents = components.filter(c => c.side === 'decoder');

    const renderNoteColumn = (title: string, componentsToRender: { id: string; label: string; side: string }[]) => (
        <div className="notes-column">
            <h2 className="column-title">{title}</h2>
            {componentsToRender.map(comp => (
                <div className="note-block" key={comp.id}>
                    <label className="note-label" htmlFor={comp.id}>{comp.label}</label>
                    <textarea
                        id={comp.id}
                        aria-label={`Notes for ${comp.label}`}
                        value={notes[comp.id] || ''}
                        onChange={(e: React.ChangeEvent<HTMLTextAreaElement>) => handleNoteChange(comp.id, e.target.value)}
                        placeholder={`Explain ${comp.label}...`}
                    />
                </div>
            ))}
        </div>
    );

    return (
        <div className="app-container">
            <header className="header">
                <h1>Interactive Transformer Architecture</h1>
                <p>An interactive diagram of the Transformer model from "Attention Is All You Need". Add your notes in the text boxes to explain each part. Your notes will be saved locally in your browser.</p>
            </header>
            
            <div className="save-button-container" style={{ flexDirection: 'column', alignItems: 'center', gap: '1rem' }}>
                <input
                    type="text"
                    className="name-input"
                    placeholder="Enter your name..."
                    value={name}
                    onChange={e => setName(e.target.value)}
                    style={{ padding: '10px', borderRadius: '6px', border: '1px solid #333', width: '250px', fontSize: '1rem', marginBottom: '0.5rem' }}
                />
                <button className="save-button" onClick={handleSaveNotes}>
                    Save Notes
                </button>
            </div>
            
            <div className="architecture-container">
                {renderNoteColumn('Encoder', encoderComponents)}
                
                <div className="image-column">
                    <img 
                        src={transformerImageDataUri}
                        alt="Transformer Architecture Diagram" 
                        className="architecture-image"
                    />
                </div>

                {renderNoteColumn('Decoder', decoderComponents)}
            </div>

            <div className={`toast ${showToast ? 'show' : ''}`}>
                Notes saved successfully!
            </div>
        </div>
    );
};

const container = document.getElementById('root');
if (container) {
    const root = createRoot(container);
    root.render(<App />);
}