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

# 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

# Extract partial command from cursor position
def get_current_command(text, cursor_pos):
    if cursor_pos <= 0:
        return ""
    
    # Look for the last backslash before cursor
    start_pos = text[:cursor_pos].rfind("\\")
    if start_pos == -1:
        return ""
    
    # Extract the partial command (from \ to cursor)
    partial_cmd = text[start_pos:cursor_pos]
    return partial_cmd

# Find matching commands for autocomplete
def find_matching_commands(partial_cmd, all_commands):
    if not partial_cmd.startswith("\\"):
        return []
    
    partial_without_backslash = partial_cmd[1:].lower()
    matches = []
    
    for cmd in all_commands:
        cmd_without_backslash = cmd[1:].lower()
        # Check if command contains the partial input (without \)
        if cmd_without_backslash.startswith(partial_without_backslash):
            matches.append(cmd)
    
    return matches

# 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"
    }
}

# Create a flat list of all LaTeX commands for autocomplete
all_latex_commands = []
for category, commands in latex_commands.items():
    all_latex_commands.extend(list(commands.keys()))

# 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}
"""

# Add custom CSS with improved sidebar styling
st.markdown("""
<style>
    /* Editor styling */
    .editor-container {
        border: 1px solid #ccc;
        border-radius: 5px;
        padding: 10px;
        background-color: #f8f9fa;
    }
    .stTextArea textarea {
        font-family: 'Courier New', Courier, monospace !important;
        font-size: 14px !important;
        line-height: 1.5 !important;
    }
    
    /* 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;
    }
    
    /* Autocomplete suggestions */
    .autocomplete-container {
        background-color: #f8f9fa;
        border: 1px solid #dee2e6;
        border-radius: 4px;
        padding: 10px;
        margin-top: 10px;
        max-height: 200px;
        overflow-y: auto;
    }
    
    .suggestion-item {
        padding: 5px;
        cursor: pointer;
        border-radius: 3px;
    }
    
    .suggestion-item:hover {
        background-color: #e9ecef;
    }
    
    /* Preview toggle button */
    .preview-button {
        margin-top: 15px;
        width: 100%;
    }
</style>
""", unsafe_allow_html=True)

# Main application
def main():
    st.title("LaTeX Editor & PDF Compiler")
    
    # Initialize session state variables
    if 'latex_code' not in st.session_state:
        st.session_state.latex_code = default_template
    if 'cursor_pos' not in st.session_state:
        st.session_state.cursor_pos = 0
    if 'show_preview' not in st.session_state:
        st.session_state.show_preview = False
    
    # 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.")
        
        # Show packages.txt content suggestion
        with st.expander("Required packages.txt content"):
            st.code("""texlive
texlive-latex-base
texlive-latex-extra
texlive-fonts-recommended
texlive-science
python3-dev
python3-pip
poppler-utils""", language="text")
    
    # Create layout
    col1, col2 = st.columns([3, 2])
    
    with col1:
        st.subheader("LaTeX Editor")
        
        # Detect current command for autocomplete
        if 'last_input' in st.session_state:
            current_command = get_current_command(st.session_state.latex_code, st.session_state.cursor_pos)
            matching_commands = find_matching_commands(current_command, all_latex_commands)
            
            # Display autocomplete suggestions when typing a command
            if current_command and matching_commands:
                st.markdown("### Command Suggestions")
                st.markdown('<div class="autocomplete-container">', unsafe_allow_html=True)
                for cmd in matching_commands[:10]:  # Limit to 10 suggestions
                    if st.button(cmd, key=f"suggest_{cmd}"):
                        # Replace partial command with the full command
                        text_before = st.session_state.latex_code[:st.session_state.cursor_pos - len(current_command)]
                        text_after = st.session_state.latex_code[st.session_state.cursor_pos:]
                        st.session_state.latex_code = text_before + cmd + text_after
                        st.rerun()
                st.markdown('</div>', unsafe_allow_html=True)
        
        # LaTeX editor
        latex_code = st.text_area(
            "Edit your LaTeX document:",
            value=st.session_state.latex_code,
            height=500,
            key="latex_editor",
            on_change=lambda: setattr(st.session_state, 'last_input', True)
        )
        st.session_state.latex_code = latex_code
        
        # Get cursor position for autocomplete (using JS - note this is limited in Streamlit)
        # This is a workaround since Streamlit doesn't directly expose cursor position
        st.markdown("""
        <script>
        document.addEventListener('DOMContentLoaded', function() {
            const textArea = document.querySelector('textarea');
            if (textArea) {
                textArea.addEventListener('keyup', function(e) {
                    const cursorPos = this.selectionStart;
                    // Would need a way to communicate this to Python
                    console.log("Cursor position:", cursorPos);
                });
            }
        });
        </script>
        """, unsafe_allow_html=True)
        
        # 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 quick search
    quick_search = st.sidebar.text_input("Type to find commands (e.g. 'fr' for '\\frac')")
    
    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}"):
                        st.session_state.latex_code += f"\n{cmd}"
                        st.rerun()
        else:
            st.sidebar.info("No matching commands found")
    
    # Regular categories
    if not quick_search:
        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}"):
                                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}"):
                                st.session_state.latex_code += f"\n{pkg}"
                                st.rerun()
                    st.markdown('</div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()