import streamlit as st
import subprocess
import tempfile
import base64
from pathlib import Path
import os
import shutil
import io
from PIL import Image
import fitz  # PyMuPDF
import re
import json
import time
import components.html as html

# Set page configuration
st.set_page_config(page_title="LaTeX Editor & Compiler", page_icon="📝", layout="wide")

# Check if pdflatex is available
def is_pdflatex_installed():
    return shutil.which("pdflatex") is not None

# Function to convert LaTeX to PDF
def latex_to_pdf(latex_code):
    # Check if pdflatex is installed
    if not is_pdflatex_installed():
        st.error("pdflatex not found. Debug info:")
        st.code(f"PATH: {os.environ.get('PATH')}")
        result = subprocess.run(["which", "pdflatex"], capture_output=True, text=True)
        st.code(f"which pdflatex: {result.stdout} {result.stderr}")
        return None, "", "Error: pdflatex is not installed or not in PATH."
    
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)
        tex_file = temp_path / "document.tex"
        pdf_file = temp_path / "document.pdf"
        
        # Write LaTeX code to file
        with open(tex_file, "w") as f:
            f.write(latex_code)
        
        try:
            # Run pdflatex to compile the LaTeX file
            process = subprocess.run(
                ["pdflatex", "-interaction=nonstopmode", "-output-directory", temp_dir, str(tex_file)],
                capture_output=True,
                text=True
            )
            
            # Check if PDF was created
            if pdf_file.exists():
                with open(pdf_file, "rb") as file:
                    pdf_data = file.read()
                return pdf_data, process.stdout, process.stderr
            else:
                return None, process.stdout, process.stderr
        except Exception as e:
            return None, "", str(e)

# Function to create download link for PDF
def get_download_link(pdf_data, filename="document.pdf"):
    b64_pdf = base64.b64encode(pdf_data).decode()
    return f'<a href="data:application/pdf;base64,{b64_pdf}" download="{filename}" class="download-button">Download PDF</a>'

# Convert PDF to image for preview
def render_pdf_preview(pdf_data):
    if not pdf_data:
        return None
    
    try:
        # Create a file-like object from the PDF data
        pdf_stream = io.BytesIO(pdf_data)
        
        # Open PDF with PyMuPDF (fitz)
        pdf_document = fitz.open(stream=pdf_stream, filetype="pdf")
        
        # Render pages as images
        images = []
        for page_num in range(min(3, len(pdf_document))):  # Preview first 3 pages max
            page = pdf_document.load_page(page_num)
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # Zoom factor 2 for better resolution
            img_data = pix.tobytes("png")
            img = Image.open(io.BytesIO(img_data))
            images.append(img)
        
        pdf_document.close()
        return images
    except Exception as e:
        st.error(f"Error rendering PDF preview: {str(e)}")
        return None

# LaTeX package reference
latex_packages = {
    "Document": {
        "\\usepackage{geometry}": "Page layout customization",
        "\\usepackage{fancyhdr}": "Custom headers and footers",
        "\\usepackage{titlesec}": "Title formatting",
        "\\usepackage{hyperref}": "Hyperlinks and PDF metadata"
    },
    "Math": {
        "\\usepackage{amsmath}": "Enhanced math formatting",
        "\\usepackage{amssymb}": "Mathematical symbols",
        "\\usepackage{mathtools}": "Extensions to amsmath",
        "\\usepackage{physics}": "Physics notation"
    },
    "Graphics": {
        "\\usepackage{graphicx}": "Include images",
        "\\usepackage{tikz}": "Create vector graphics",
        "\\usepackage{pgfplots}": "Create plots",
        "\\usepackage{float}": "Better figure placement"
    },
    "Tables": {
        "\\usepackage{tabularx}": "Flexible tables",
        "\\usepackage{booktabs}": "Professional tables",
        "\\usepackage{colortbl}": "Colored tables",
        "\\usepackage{multirow}": "Multi-row cells"
    },
    "Content": {
        "\\usepackage{listings}": "Code syntax highlighting",
        "\\usepackage{minted}": "Advanced code highlighting",
        "\\usepackage{biblatex}": "Bibliography management",
        "\\usepackage{xcolor}": "Color support"
    }
}

# LaTeX commands reference
latex_commands = {
    "Document Structure": {
        "\\documentclass{article}": "Specifies the type of document",
        "\\begin{document}": "Starts the document content",
        "\\end{document}": "Ends the document content",
        "\\title{...}": "Sets the document title",
        "\\author{...}": "Sets the document author",
        "\\date{...}": "Sets the document date",
        "\\maketitle": "Prints the title, author, and date"
    },
    "Sections": {
        "\\section{...}": "Creates a section",
        "\\subsection{...}": "Creates a subsection",
        "\\subsubsection{...}": "Creates a subsubsection",
        "\\paragraph{...}": "Creates a paragraph heading",
        "\\tableofcontents": "Generates a table of contents"
    },
    "Text Formatting": {
        "\\textbf{...}": "Bold text",
        "\\textit{...}": "Italic text",
        "\\underline{...}": "Underlined text",
        "\\emph{...}": "Emphasized text",
        "\\texttt{...}": "Typewriter text",
        "\\textsc{...}": "Small caps text",
        "\\textsf{...}": "Sans-serif text",
        "\\color{red}{...}": "Colored text (requires xcolor)"
    },
    "Math": {
        "$...$": "Inline math mode",
        "$$...$$": "Display math mode",
        "\\begin{equation}...\\end{equation}": "Numbered equation",
        "\\begin{align}...\\end{align}": "Aligned equations",
        "\\frac{num}{denom}": "Fraction",
        "\\dfrac{num}{denom}": "Display fraction",
        "\\sqrt{...}": "Square root",
        "\\sqrt[n]{...}": "nth root",
        "\\sum_{lower}^{upper}": "Summation",
        "\\prod_{lower}^{upper}": "Product",
        "\\int_{lower}^{upper}": "Integral",
        "\\lim_{x \\to value}": "Limit",
        "\\vec{...}": "Vector",
        "\\overline{...}": "Overline",
        "\\hat{...}": "Hat accent",
        "\\partial": "Partial derivative"
    },
    "Lists": {
        "\\begin{itemize}...\\end{itemize}": "Bulleted list",
        "\\begin{enumerate}...\\end{enumerate}": "Numbered list",
        "\\begin{description}...\\end{description}": "Description list",
        "\\item": "List item",
        "\\item[custom]": "Custom label item"
    },
    "Tables": {
        "\\begin{table}...\\end{table}": "Table environment",
        "\\begin{tabular}{cols}...\\end{tabular}": "Create a table",
        "\\hline": "Horizontal line in table",
        "\\cline{i-j}": "Partial horizontal line",
        "cell1 & cell2 & cell3 \\\\": "Table row",
        "\\multicolumn{n}{align}{content}": "Span multiple columns"
    },
    "Figures": {
        "\\begin{figure}...\\end{figure}": "Figure environment",
        "\\includegraphics[width=0.8\\textwidth]{filename}": "Include an image",
        "\\caption{...}": "Add a caption to a figure or table",
        "\\label{...}": "Add a label for cross-referencing",
        "\\ref{...}": "Reference a labeled item"
    },
    "Citations": {
        "\\cite{key}": "Citation",
        "\\bibliography{file}": "Bibliography source",
        "\\bibliographystyle{style}": "Bibliography style"
    }
}

# Flatten all commands for autocomplete
all_commands = []
for category, commands in latex_commands.items():
    for cmd in commands:
        all_commands.append(cmd)

# Default LaTeX template
default_template = r"""\documentclass{article}
\usepackage[utf8]{inputenc}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{graphicx}
\usepackage{hyperref}

\title{LaTeX Document}
\author{Your Name}
\date{\today}

\begin{document}

\maketitle

\section{Introduction}
Your introduction here. Insert some text to demonstrate LaTeX.

\section{Mathematical Expressions}
\subsection{Equations}
The famous Einstein's equation:
\begin{equation}
    E = mc^2
\end{equation}

The quadratic formula:
\begin{equation}
    x = \frac{-b \pm \sqrt{b^2 - 4ac}}{2a}
\end{equation}

\subsection{Calculus}
An integral example:
\begin{equation}
    \int_{0}^{\pi} \sin(x) \, dx = 2
\end{equation}

\section{Lists and Items}
\subsection{Bullet Points}
\begin{itemize}
    \item First item
    \item Second item
    \item Third item
\end{itemize}

\subsection{Numbered List}
\begin{enumerate}
    \item First step
    \item Second step
    \item Third step
\end{enumerate}

\section{Tables}
\begin{table}[h]
    \centering
    \begin{tabular}{|c|c|c|}
        \hline
        Cell 1 & Cell 2 & Cell 3 \\
        \hline
        Data 1 & Data 2 & Data 3 \\
        \hline
    \end{tabular}
    \caption{A simple table}
    \label{tab:simple}
\end{table}

\section{Conclusion}
Your conclusion here.

\end{document}
"""

# Create the Monaco Editor component with autocomplete
def create_monaco_editor():
    # Convert the LaTeX commands to a format usable by Monaco
    monaco_suggestions = []
    for category, commands in latex_commands.items():
        for cmd, desc in commands.items():
            # Format command for Monaco
            cmd_text = cmd.replace("{...}", "").replace("...", "")
            monaco_suggestions.append({
                "label": cmd_text,
                "kind": 14,  # Snippet
                "insertText": cmd_text,
                "detail": desc,
                "documentation": f"Category: {category}"
            })
    
    # JSON stringify the suggestions for passing to JavaScript
    suggestions_json = json.dumps(monaco_suggestions)
    
    # Create the Monaco editor component with LaTeX syntax highlighting and autocomplete
    monaco_editor = f"""
    <div id="container" style="width:100%;height:500px;border:1px solid #ccc;"></div>
    <script src="https://cdnjs.cloudflare.com/ajax/libs/monaco-editor/0.36.1/min/vs/loader.min.js"></script>
    <script>
        require.config({{ paths: {{ 'vs': 'https://cdnjs.cloudflare.com/ajax/libs/monaco-editor/0.36.1/min/vs' }}});
        
        // Before loading vs/editor/editor.main, define a global MonacoEnvironment that overwrites
        // the default worker url location (used when creating WebWorkers). The problem here is that
        // HTML5 does not allow cross-domain web workers, so we need to proxy the instantiation of
        // a web worker through a same-domain script
        window.MonacoEnvironment = {{
            getWorkerUrl: function(workerId, label) {{
                return `data:text/javascript;charset=utf-8,
                    self.MonacoEnvironment = {{
                        baseUrl: 'https://cdnjs.cloudflare.com/ajax/libs/monaco-editor/0.36.1/min/'
                    }};
                    importScripts('https://cdnjs.cloudflare.com/ajax/libs/monaco-editor/0.36.1/min/vs/base/worker/workerMain.js');`;
            }}
        }};
        
        require(['vs/editor/editor.main'], function() {{
            // Define LaTeX language
            monaco.languages.register({{ id: 'latex' }});
            
            // Define LaTeX syntax highlighting
            monaco.languages.setMonarchTokensProvider('latex', {{
                tokenizer: {{
                    root: [
                        [/\\\\[a-zA-Z]+/, 'keyword'],
                        [/\\\\begin\\{{[^}}]*\\}}/, 'keyword'],
                        [/\\\\end\\{{[^}}]*\\}}/, 'keyword'],
                        [/\\$\\$.*?\\$\\$/, 'string'],
                        [/\\$.*?\\$/, 'string'],
                        [/%.*$/, 'comment'],
                        [/\\{{/, 'delimiter.curly'],
                        [/\\}}/, 'delimiter.curly'],
                        [/\\[/, 'delimiter.square'],
                        [/\\]/, 'delimiter.square']
                    ]
                }}
            }});
            
            // Setup autocomplete provider for LaTeX
            monaco.languages.registerCompletionItemProvider('latex', {{
                provideCompletionItems: function(model, position) {{
                    const textUntilPosition = model.getValueInRange({{
                        startLineNumber: position.lineNumber,
                        startColumn: 1,
                        endLineNumber: position.lineNumber,
                        endColumn: position.column
                    }});
                    
                    // Check if we're typing a LaTeX command
                    const match = textUntilPosition.match(/\\\\([a-zA-Z]*)$/);
                    if (!match) return {{ suggestions: [] }};
                    
                    const word = match[1];
                    const suggestions = {suggestions_json};
                    
                    // Filter suggestions based on what has been typed
                    return {{
                        suggestions: suggestions.filter(function(s) {{
                            return s.label.substring(1).toLowerCase().indexOf(word.toLowerCase()) === 0;
                        }})
                    }};
                }},
                triggerCharacters: ['\\\\']
            }});
            
            // Get the initial value from localStorage or use default
            let initialValue = localStorage.getItem('latexEditorContent') || 
                               {json.dumps(default_template)};
            
            // Create the editor
            const editor = monaco.editor.create(document.getElementById('container'), {{
                value: initialValue,
                language: 'latex',
                theme: 'vs',
                automaticLayout: true,
                minimap: {{ enabled: false }},
                fontSize: 14,
                fontFamily: "'Courier New', monospace",
                lineNumbers: "on",
                scrollBeyondLastLine: false,
                tabSize: 2,
                wordWrap: "on"
            }});
            
            // Set up change handler to save content to localStorage
            editor.onDidChangeModelContent(function() {{
                const value = editor.getValue();
                localStorage.setItem('latexEditorContent', value);
                
                // Send value back to Streamlit
                if (window.Streamlit) {{
                    window.Streamlit.setComponentValue(value);
                }}
            }});
            
            // Handle initialization from Streamlit
            window.addEventListener('message', function(event) {{
                if (event.data.type === 'streamlit:render') {{
                    const args = event.data.args;
                    if (args.value) {{
                        editor.setValue(args.value);
                    }}
                }}
            }});
        }});
    </script>
    """
    
    return monaco_editor

# Add custom CSS
st.markdown("""
<style>
    /* Editor styling */
    .editor-container {
        border: 1px solid #ccc;
        border-radius: 5px;
        padding: 10px;
        background-color: #f8f9fa;
    }
    
    /* Download button styling */
    .download-button {
        display: inline-block;
        padding: 0.7em 1.4em;
        background-color: #4CAF50;
        color: white !important;
        text-align: center;
        text-decoration: none;
        font-size: 18px;
        border-radius: 4px;
        transition: background-color 0.3s;
        margin-top: 10px;
        font-weight: bold;
        box-shadow: 0 2px 5px rgba(0,0,0,0.2);
    }
    .download-button:hover {
        background-color: #45a049;
        box-shadow: 0 4px 8px rgba(0,0,0,0.3);
    }
    
    /* Sidebar styling */
    .sidebar .sidebar-content {
        background-color: #f0f2f6;
    }
    
    /* LaTeX command styling */
    .latex-command {
        background-color: #e9ecef;
        padding: 4px 8px;
        border-radius: 4px;
        font-family: 'Courier New', Courier, monospace;
        color: #1e1e1e;
        cursor: pointer;
        display: inline-block;
        margin-bottom: 4px;
        border: 1px solid #ced4da;
    }
    .latex-command:hover {
        background-color: #d0d7de;
        border-color: #adb5bd;
    }
    
    /* Command description styling */
    .command-description {
        color: #495057;
        padding-left: 8px;
        display: inline-block;
    }
    
    /* Category title styling */
    .category-title {
        font-weight: bold;
        color: #212529;
        margin-top: 15px;
        margin-bottom: 8px;
    }
    
    /* Expander styling */
    .streamlit-expanderHeader {
        font-weight: bold;
        color: #212529;
        background-color: #e9ecef;
        border-radius: 4px;
    }
    
    /* Command list container */
    .command-list {
        background-color: #f8f9fa;
        border-radius: 4px;
        padding: 8px;
        border: 1px solid #dee2e6;
    }
    
    /* Insert button styling */
    .insert-button {
        background-color: #007bff;
        color: white;
        border: none;
        border-radius: 4px;
        padding: 2px 8px;
        margin-left: 8px;
        cursor: pointer;
        font-size: 12px;
    }
    .insert-button:hover {
        background-color: #0069d9;
    }
    
    /* PDF preview container */
    .pdf-preview-container {
        border: 1px solid #dee2e6;
        border-radius: 5px;
        padding: 15px;
        background-color: #f8f9fa;
    }
    
    /* Hide Streamlit footer */
    footer {display: none !important;}
    #MainMenu {visibility: hidden;}
</style>
""", unsafe_allow_html=True)

# Create a component for the Monaco editor
def monaco_editor_component(key, default_value=""):
    if key not in st.session_state:
        st.session_state[key] = default_value
    
    # Create a placeholder for the editor
    placeholder = st.empty()
    
    # Display the Monaco editor
    editor_html = create_monaco_editor()
    component_value = placeholder.html(editor_html, height=520)
    
    # Return the current value
    if component_value is not None:
        st.session_state[key] = component_value
    
    return st.session_state[key]

# Main application
def main():
    st.title("LaTeX Editor & PDF Compiler")
    
    # Display installation status
    if not is_pdflatex_installed():
        st.warning("⚠️ LaTeX is not installed correctly. The PDF compilation feature will not work.")
        st.info("For Hugging Face Spaces, make sure you have a packages.txt file with the necessary LaTeX packages.")
    
    # Initialize session state
    if 'latex_code' not in st.session_state:
        st.session_state.latex_code = default_template
    if 'show_preview' not in st.session_state:
        st.session_state.show_preview = False
    
    # Create layout
    col1, col2 = st.columns([3, 2])
    
    with col1:
        st.subheader("LaTeX Editor")
        
        # Use the Monaco editor component
        latex_code = monaco_editor_component("monaco_editor", st.session_state.latex_code)
        st.session_state.latex_code = latex_code
        
        # Create a small spacer
        st.write("")
        
        # Control buttons
        col1_1, col1_2, col1_3 = st.columns(3)
        
        with col1_1:
            if st.button("Compile PDF", use_container_width=True):
                st.session_state.compile_clicked = True
        
        with col1_2:
            if st.button("Load Template", use_container_width=True):
                st.session_state.latex_code = default_template
                st.rerun()
        
        with col1_3:
            if st.button("Clear Editor", use_container_width=True):
                st.session_state.latex_code = ""
                st.rerun()
    
    with col2:
        st.subheader("PDF Output")
        
        # PDF compilation
        if 'compile_clicked' in st.session_state and st.session_state.compile_clicked:
            with st.spinner("Compiling LaTeX to PDF..."):
                pdf_data, stdout, stderr = latex_to_pdf(latex_code)
                
                if pdf_data:
                    st.session_state.pdf_data = pdf_data
                    st.success("PDF compiled successfully!")
                    
                    # Toggle button for preview
                    if st.button("Show/Hide Preview", use_container_width=True):
                        st.session_state.show_preview = not st.session_state.show_preview
                    
                    # Download button always available
                    st.markdown(get_download_link(pdf_data), unsafe_allow_html=True)
                    
                    # Optional preview
                    if st.session_state.show_preview:
                        preview_images = render_pdf_preview(pdf_data)
                        if preview_images:
                            st.write("PDF Preview (First Pages):")
                            for i, img in enumerate(preview_images):
                                st.image(img, caption=f"Page {i+1}", use_container_width=True, 
                                        output_format="PNG")
                    
                    st.session_state.compile_clicked = False
                else:
                    st.error("Compilation Error")
                    with st.expander("Error Details"):
                        st.text(stdout)
                        st.text(stderr)
                    st.session_state.compile_clicked = False
        
        # Display previous PDF if available
        elif 'pdf_data' in st.session_state and st.session_state.pdf_data:
            # Toggle button for preview
            if st.button("Show/Hide Preview", use_container_width=True):
                st.session_state.show_preview = not st.session_state.show_preview
            
            # Download button always available
            st.markdown(get_download_link(st.session_state.pdf_data), unsafe_allow_html=True)
            
            # Optional preview
            if st.session_state.show_preview:
                preview_images = render_pdf_preview(st.session_state.pdf_data)
                if preview_images:
                    st.write("PDF Preview (First Pages):")
                    for i, img in enumerate(preview_images):
                        st.image(img, caption=f"Page {i+1}", use_container_width=True, 
                                output_format="PNG")
        else:
            st.info("Compile your LaTeX document to generate a PDF for download")
    
    # LaTeX Reference Sidebar
    st.sidebar.title("LaTeX Reference")
    
    # Command search
    quick_search = st.sidebar.text_input("Find LaTeX Commands", "")
    
    if quick_search:
        # Find and display matching commands
        matching_cmds = []
        for category, commands in latex_commands.items():
            for cmd, desc in commands.items():
                cmd_without_backslash = cmd.replace("\\", "").lower()
                if quick_search.lower() in cmd_without_backslash:
                    matching_cmds.append((cmd, desc, category))
        
        if matching_cmds:
            st.sidebar.markdown("### Matching Commands")
            for cmd, desc, category in matching_cmds[:15]:  # Limit to 15 results
                col1, col2 = st.sidebar.columns([4, 1])
                with col1:
                    st.markdown(f"<div><span class='latex-command'>{cmd}</span> <small>{category}</small></div>", unsafe_allow_html=True)
                with col2:
                    if st.button("Insert", key=f"quick_{cmd}"):
                        # Update LaTeX code with the inserted command
                        st.session_state.latex_code += f"\n{cmd}"
                        st.rerun()
        else:
            st.sidebar.info("No matching commands found")
    
    # Regular categories
    tab1, tab2 = st.sidebar.tabs(["Commands", "Packages"])
    
    with tab1:
        for category, commands in latex_commands.items():
            with st.expander(category, expanded=category=="Math"):
                st.markdown('<div class="command-list">', unsafe_allow_html=True)
                for cmd, desc in commands.items():
                    col1, col2 = st.sidebar.columns([4, 1])
                    with col1:
                        st.markdown(f"<div><span class='latex-command'>{cmd}</span></div>", unsafe_allow_html=True)
                    with col2:
                        if st.button("Insert", key=f"btn_{cmd}"):
                            # Update LaTeX code with the inserted command
                            st.session_state.latex_code += f"\n{cmd}"
                            st.rerun()
                st.markdown('</div>', unsafe_allow_html=True)
    
    with tab2:
        for category, packages in latex_packages.items():
            with st.expander(category):
                st.markdown('<div class="command-list">', unsafe_allow_html=True)
                for pkg, desc in packages.items():
                    col1, col2 = st.sidebar.columns([4, 1])
                    with col1:
                        st.markdown(f"<div><span class='latex-command'>{pkg}</span></div>", unsafe_allow_html=True)
                    with col2:
                        if st.button("Insert", key=f"btn_{pkg}"):
                            # Update LaTeX code with the inserted command
                            st.session_state.latex_code += f"\n{pkg}"
                            st.rerun()
                st.markdown('</div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()