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

# 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

# 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 VS Code-like styling
st.markdown("""
<style>
    /* VS Code-like styling */
    .vscode-editor textarea {
        font-family: 'Consolas', 'Monaco', 'Courier New', monospace !important;
        font-size: 14px !important;
        line-height: 1.5 !important;
        background-color: #1e1e1e !important;
        color: #d4d4d4 !important;
        padding: 10px !important;
        border-radius: 4px !important;
        border: 1px solid #252526 !important;
    }
    
    /* Editor container styling */
    .vscode-container {
        background-color: #1e1e1e;
        border-radius: 6px;
        padding: 8px;
        border: 1px solid #333;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    
    /* Make scrollbars VS Code style */
    .vscode-editor textarea::-webkit-scrollbar {
        width: 14px;
        height: 14px;
    }
    
    .vscode-editor textarea::-webkit-scrollbar-thumb {
        background-color: #424242;
        border-radius: 7px;
        border: 3px solid #1e1e1e;
    }
    
    .vscode-editor textarea::-webkit-scrollbar-track {
        background-color: #1e1e1e;
    }
    
    /* Button styling */
    .vscode-button {
        background-color: #0e639c;
        color: white;
        border: none;
        padding: 8px 12px;
        border-radius: 2px;
        cursor: pointer;
        font-size: 13px;
        margin-right: 10px;
        margin-top: 10px;
        transition: background-color 0.2s;
    }
    
    .vscode-button:hover {
        background-color: #1177bb;
    }
    
    /* Download button styling */
    .download-button {
        display: inline-block;
        padding: 0.7em 1.4em;
        background-color: #3d995e;
        color: white !important;
        text-align: center;
        text-decoration: none;
        font-size: 14px;
        border-radius: 2px;
        transition: background-color 0.3s;
        margin-top: 10px;
        font-weight: normal;
    }
    .download-button:hover {
        background-color: #4eb772;
    }
    
    /* Status bar styling */
    .status-bar {
        background-color: #007acc;
        color: white;
        padding: 2px 8px;
        font-size: 12px;
        border-radius: 2px 2px 0 0;
        display: flex;
        justify-content: space-between;
    }
    
    /* Terminal/output styling */
    .terminal-output {
        background-color: #1e1e1e;
        color: #cccccc;
        padding: 10px;
        font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
        font-size: 12px;
        border-radius: 0 0 4px 4px;
        border-top: 1px solid #333;
        max-height: 200px;
        overflow-y: auto;
    }
    
    /* PDF preview container */
    .pdf-preview-container {
        border: 1px solid #333;
        border-radius: 4px;
        padding: 15px;
        background-color: #252526;
    }
    
    /* Info and error messages */
    .stInfo {
        background-color: #063b49;
        color: #bbbbbb;
        border: 1px solid #145b6c;
    }
    
    .stError {
        background-color: #5a1d1d;
        color: #bbbbbb;
        border: 1px solid #6c2b2b;
    }
    
    .stSuccess {
        background-color: #143d27;
        color: #bbbbbb;
        border: 1px solid #1e5a3a;
    }
    
    /* Hide Streamlit elements */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    
    /* Customize the rest of Streamlit UI */
    .stApp {
        background-color: #252526;
    }
    
    h1, h2, h3, h4, h5, h6, p, div {
        color: #cccccc;
    }
    
    .stTabs [data-baseweb="tab-list"] {
        background-color: #2d2d2d;
    }
    
    .stTabs [data-baseweb="tab"] {
        color: #cccccc;
    }
</style>
""", unsafe_allow_html=True)

# Function to create a VS Code-like editor
def vs_code_editor(key, height=500):
    editor_html = f"""
    <div class="vscode-container">
        <div class="status-bar">
            <span>document.tex - LaTeX</span>
            <span>UTF-8</span>
        </div>
    </div>
    """
    st.markdown(editor_html, unsafe_allow_html=True)
    
    # Create the actual editor with VS Code styling
    return st.text_area("", value=st.session_state.get(key, ""), 
                        height=height, key=key, label_visibility="collapsed",
                        help="Type your LaTeX code here")

# Main application
def main():
    # Set up a clean, dark theme
    st.markdown("<h1 style='color: #cccccc; margin-bottom: 20px;'>LaTeX Editor</h1>", unsafe_allow_html=True)
    
    # 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
    
    # Display installation status
    if not is_pdflatex_installed():
        st.warning("⚠️ LaTeX is not installed correctly. The PDF compilation feature will not work.")
    
    # Create layout - full width editor
    col1, col2 = st.columns([3, 2])
    
    with col1:
        # VS Code-like editor with custom class for styling
        st.markdown('<div class="vscode-editor">', unsafe_allow_html=True)
        latex_code = vs_code_editor("latex_editor", height=500)
        st.markdown('</div>', unsafe_allow_html=True)
        st.session_state.latex_code = latex_code
        
        # Control buttons with VS Code styling
        st.markdown("""
        <div style="display: flex; gap: 10px;">
            <button class="vscode-button" onclick="document.querySelector('[data-testid=\\\"stFormSubmitButton\\\"]').click()">
                Compile PDF
            </button>
            <button class="vscode-button" onclick="document.querySelector('[key=\\\"load_template\\\"]').click()">
                Load Template
            </button>
            <button class="vscode-button" onclick="document.querySelector('[key=\\\"clear_editor\\\"]').click()">
                Clear Editor
            </button>
        </div>
        """, unsafe_allow_html=True)
        
        # Hidden buttons to handle the clicks
        if st.button("Compile PDF", key="compile", help="Compile LaTeX to PDF"):
            st.session_state.compile_clicked = True
        
        if st.button("Load Template", key="load_template", help="Load default template"):
            st.session_state.latex_code = default_template
            st.rerun()
        
        if st.button("Clear Editor", key="clear_editor", help="Clear editor content"):
            st.session_state.latex_code = ""
            st.rerun()
    
    with col2:
        st.markdown("<h3 style='color: #cccccc; margin-bottom: 10px;'>Output</h3>", unsafe_allow_html=True)
        
        # PDF compilation
        if 'compile_clicked' in st.session_state and st.session_state.compile_clicked:
            with st.spinner("Compiling..."):
                pdf_data, stdout, stderr = latex_to_pdf(latex_code)
                
                if pdf_data:
                    st.session_state.pdf_data = pdf_data
                    st.success("Compilation successful")
                    
                    # Toggle button for preview
                    if st.button("Toggle Preview", help="Show or hide the PDF preview"):
                        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:
                        st.markdown('<div class="pdf-preview-container">', unsafe_allow_html=True)
                        preview_images = render_pdf_preview(pdf_data)
                        if preview_images:
                            for i, img in enumerate(preview_images):
                                st.image(img, caption=f"Page {i+1}", use_container_width=True, 
                                        output_format="PNG")
                        st.markdown('</div>', unsafe_allow_html=True)
                    
                    # Terminal output in collapsible section
                    with st.expander("Terminal Output"):
                        st.markdown('<div class="terminal-output">', unsafe_allow_html=True)
                        st.text(stdout)
                        st.markdown('</div>', unsafe_allow_html=True)
                    
                    st.session_state.compile_clicked = False
                else:
                    st.error("Compilation failed")
                    st.markdown('<div class="terminal-output">', unsafe_allow_html=True)
                    st.text(stderr)
                    st.markdown('</div>', unsafe_allow_html=True)
                    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("Toggle Preview", help="Show or hide the PDF preview"):
                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:
                st.markdown('<div class="pdf-preview-container">', unsafe_allow_html=True)
                preview_images = render_pdf_preview(st.session_state.pdf_data)
                if preview_images:
                    for i, img in enumerate(preview_images):
                        st.image(img, caption=f"Page {i+1}", use_container_width=True, 
                                output_format="PNG")
                st.markdown('</div>', unsafe_allow_html=True)
        else:
            st.info("Click 'Compile PDF' to generate output")

if __name__ == "__main__":
    main()