import streamlit as st
import pandas as pd
import re
import io
import tempfile
import os

# Configure Streamlit to prevent permission issues
os.environ['STREAMLIT_SERVER_HEADLESS'] = 'true'
os.environ['STREAMLIT_SERVER_ENABLE_CORS'] = 'false'
os.environ['STREAMLIT_SERVER_ENABLE_XSRF_PROTECTION'] = 'false'

st.set_page_config(
    page_title="🧪 LaTeX MCQs", 
    layout="wide",
    initial_sidebar_state="collapsed"
)
st.title("🧪 Chemistry MCQs with LaTeX Rendering")

# Helper to safely convert mixed text + LaTeX into proper LaTeX
def format_for_latex(text):
    """Convert mixed text and LaTeX expressions into properly formatted LaTeX."""
    if pd.isna(text) or text == "":
        return ""
    
    text = str(text).strip()
    
    # Remove existing \( \) or $ $ wrappers if any
    text = re.sub(r'\\\(|\\\)', '', text)
    text = re.sub(r'\$(.*?)\$', r'\1', text)
    
    # Handle common chemistry and math patterns
    latex_commands = r'(\\(?:ce|text|frac|sqrt|sum|int|lim|mathrm|mathbf|mathit|alpha|beta|gamma|delta|epsilon|theta|lambda|mu|pi|sigma|omega|rightarrow|leftarrow|leftrightarrow|cdot|times|div|pm|neq|leq|geq|approx|infty|partial|nabla|Delta|Omega|therefore|because)\b[^{]*(?:\{[^}]*\})*)'
    
    parts = re.split(f'({latex_commands})', text)
    
    formatted = ""
    for part in parts:
        if not part:
            continue
            
        if re.match(r'\\[a-zA-Z]+', part):
            formatted += part
        else:
            if part.strip():
                part = re.sub(r'([A-Za-z])(\d+)', r'\1_{\2}', part)
                part = re.sub(r'\^(\d+)', r'^{\1}', part)
                
                if not re.search(r'[{}\\]', part):
                    formatted += f"\\text{{{part}}}"
                else:
                    formatted += part
    
    return formatted

def validate_dataframe(df):
    """Validate the uploaded dataframe structure."""
    required_columns = ["question_latex", "option_1", "option_2", "option_3", "option_4"]
    missing_columns = [col for col in required_columns if col not in df.columns]
    
    if missing_columns:
        return False, f"Missing required columns: {missing_columns}"
    
    empty_questions = df["question_latex"].isna().sum()
    if empty_questions > 0:
        st.warning(f"⚠️ Found {empty_questions} empty questions that will be skipped.")
    
    return True, "Valid"

def render_question(idx, row):
    """Render a single question with its options."""
    if pd.isna(row["question_latex"]) or str(row["question_latex"]).strip() == "":
        return
    
    st.markdown(f"### Question {idx+1}")
    
    # Render question
    try:
        question_latex = format_for_latex(row["question_latex"])
        if question_latex:
            st.latex(question_latex)
        else:
            st.markdown(f"**{row['question_latex']}**")
    except Exception as e:
        st.error(f"Error rendering question {idx+1}: {e}")
        st.markdown(f"**{row['question_latex']}**")
    
    # Render options
    options = {
        "A": row["option_1"],
        "B": row["option_2"],
        "C": row["option_3"],
        "D": row["option_4"]
    }
    
    col1, col2 = st.columns(2)
    
    for i, (key, value) in enumerate(options.items()):
        with col1 if i % 2 == 0 else col2:
            st.markdown(f"**{key}.**")
            try:
                option_latex = format_for_latex(value)
                if option_latex:
                    st.latex(option_latex)
                else:
                    st.markdown(f"{value}")
            except Exception as e:
                st.warning(f"Error rendering option {key}: {e}")
                st.markdown(f"{value}")
    
    st.markdown("---")

def safe_read_excel(uploaded_file):
    """Safely read Excel file with multiple fallback methods."""
    try:
        # Method 1: Direct reading
        df = pd.read_excel(uploaded_file)
        return df, None
    except Exception as e1:
        try:
            # Method 2: BytesIO
            bytes_data = uploaded_file.getvalue()
            df = pd.read_excel(io.BytesIO(bytes_data))
            return df, None
        except Exception as e2:
            try:
                # Method 3: Temporary file in /tmp (writable in containers)
                with tempfile.NamedTemporaryFile(delete=False, suffix='.xlsx', dir='/tmp') as tmp_file:
                    tmp_file.write(uploaded_file.getbuffer())
                    tmp_file.flush()
                    
                    df = pd.read_excel(tmp_file.name)
                    
                    # Cleanup
                    try:
                        os.unlink(tmp_file.name)
                    except:
                        pass
                    
                    return df, None
            except Exception as e3:
                error_msg = f"All methods failed - Method 1: {str(e1)[:100]}, Method 2: {str(e2)[:100]}, Method 3: {str(e3)[:100]}"
                return None, error_msg

# Main app
st.markdown("""
### 📋 Instructions
1. Upload an Excel file (.xlsx or .xls) with your questions
2. Make sure your file has the required columns (see sample format below)
3. Navigate through questions using the controls
""")

uploaded_file = st.file_uploader(
    "📥 Upload Excel file with LaTeX questions", 
    type=["xlsx", "xls"],
    help="Upload an Excel file containing your MCQ questions with LaTeX formatting"
)

if uploaded_file is not None:
    st.info(f"📄 Processing: {uploaded_file.name} ({uploaded_file.size} bytes)")
    
    with st.spinner("🔄 Reading Excel file..."):
        df, error_msg = safe_read_excel(uploaded_file)
    
    if df is None:
        st.error(f"❌ Error: {error_msg}")
        st.markdown("### Troubleshooting:")
        st.markdown("""
        - Ensure file is valid Excel format (.xlsx/.xls)
        - Check file isn't corrupted or password-protected
        - Try re-uploading the file
        - Verify all required columns exist
        """)
    else:
        is_valid, message = validate_dataframe(df)
        
        if not is_valid:
            st.error(f"❌ {message}")
            st.markdown("### Required Columns:")
            st.code("question_latex, option_1, option_2, option_3, option_4")
        else:
            st.success(f"✅ Loaded {len(df)} questions successfully!")
            
            # Initialize session state
            if 'current_question' not in st.session_state:
                st.session_state.current_question = 0
            if 'navigator_page' not in st.session_state:
                st.session_state.navigator_page = 0
            
            # Navigation system
            st.markdown("### 📋 Question Navigator")
            
            total_questions = len(df)
            questions_per_page = 50
            questions_per_row = 10
            
            total_nav_pages = (total_questions + questions_per_page - 1) // questions_per_page
            
            # Navigator pagination (if needed)
            if total_questions > questions_per_page:
                nav_cols = st.columns([1, 1, 2, 1, 1])
                
                with nav_cols[0]:
                    if st.button("⏪ First Set", disabled=st.session_state.navigator_page == 0):
                        st.session_state.navigator_page = 0
                        st.rerun()
                
                with nav_cols[1]:
                    if st.button("◀️ Prev Set", disabled=st.session_state.navigator_page == 0):
                        st.session_state.navigator_page -= 1
                        st.rerun()
                
                with nav_cols[2]:
                    start_q = st.session_state.navigator_page * questions_per_page + 1
                    end_q = min((st.session_state.navigator_page + 1) * questions_per_page, total_questions)
                    st.markdown(f"**Questions {start_q}-{end_q} of {total_questions}**")
                
                with nav_cols[3]:
                    if st.button("Next Set ▶️", disabled=st.session_state.navigator_page >= total_nav_pages - 1):
                        st.session_state.navigator_page += 1
                        st.rerun()
                
                with nav_cols[4]:
                    if st.button("Last Set ⏩", disabled=st.session_state.navigator_page >= total_nav_pages - 1):
                        st.session_state.navigator_page = total_nav_pages - 1
                        st.rerun()
                
                st.markdown("---")
            
            # Question number buttons
            start_idx = st.session_state.navigator_page * questions_per_page
            end_idx = min(start_idx + questions_per_page, total_questions)
            
            questions_to_show = end_idx - start_idx
            rows_needed = (questions_to_show + questions_per_row - 1) // questions_per_row
            
            for row in range(rows_needed):
                row_start_idx = start_idx + (row * questions_per_row)
                row_end_idx = min(row_start_idx + questions_per_row, end_idx)
                
                cols = st.columns(row_end_idx - row_start_idx)
                
                for i, col in enumerate(cols):
                    question_idx = row_start_idx + i
                    question_num = question_idx + 1
                    
                    with col:
                        if question_idx == st.session_state.current_question:
                            button_label = f"**Q{question_num}**"
                            button_type = "primary"
                        else:
                            button_label = f"Q{question_num}"
                            button_type = "secondary"
                        
                        if st.button(button_label, key=f"q_btn_{question_idx}", type=button_type):
                            st.session_state.current_question = question_idx
                            st.rerun()
            
            st.markdown("---")
            
            # Main navigation controls
            nav_cols = st.columns([1, 1, 2, 1, 1])
            
            with nav_cols[0]:
                if st.button("⏮️ First", disabled=st.session_state.current_question == 0):
                    st.session_state.current_question = 0
                    st.session_state.navigator_page = 0
                    st.rerun()
            
            with nav_cols[1]:
                if st.button("⬅️ Previous", disabled=st.session_state.current_question == 0):
                    st.session_state.current_question -= 1
                    new_page = st.session_state.current_question // questions_per_page
                    if new_page != st.session_state.navigator_page:
                        st.session_state.navigator_page = new_page
                    st.rerun()
            
            with nav_cols[2]:
                question_options = [f"Question {i+1}" for i in range(total_questions)]
                selected_q = st.selectbox(
                    "Jump to:",
                    question_options,
                    index=st.session_state.current_question
                )
                
                new_index = question_options.index(selected_q)
                if new_index != st.session_state.current_question:
                    st.session_state.current_question = new_index
                    new_page = new_index // questions_per_page
                    if new_page != st.session_state.navigator_page:
                        st.session_state.navigator_page = new_page
                    st.rerun()
            
            with nav_cols[3]:
                if st.button("Next ➡️", disabled=st.session_state.current_question >= total_questions - 1):
                    st.session_state.current_question += 1
                    new_page = st.session_state.current_question // questions_per_page
                    if new_page != st.session_state.navigator_page:
                        st.session_state.navigator_page = new_page
                    st.rerun()
            
            with nav_cols[4]:
                if st.button("Last ⏭️", disabled=st.session_state.current_question >= total_questions - 1):
                    st.session_state.current_question = total_questions - 1
                    st.session_state.navigator_page = total_nav_pages - 1
                    st.rerun()
            
            # Progress bar
            progress = (st.session_state.current_question + 1) / total_questions
            st.progress(progress)
            st.caption(f"Progress: {st.session_state.current_question + 1}/{total_questions} ({progress:.1%})")
            
            st.markdown("---")
            
            # Render current question
            if 0 <= st.session_state.current_question < len(df):
                current_row = df.iloc[st.session_state.current_question]
                render_question(st.session_state.current_question, current_row)

# Sample data section
with st.expander("📋 Sample Excel Format"):
    sample_data = pd.DataFrame({
        "question_latex": [
            "What is the molecular formula for water?",
            "Calculate the pH of a 0.1 M HCl solution. Use \\text{pH} = -\\log[\\text{H}^+]",
            "Balance the equation: \\ce{C2H6 + O2 -> CO2 + H2O}"
        ],
        "option_1": ["H2O", "1", "\\ce{2C2H6 + 7O2 -> 4CO2 + 6H2O}"],
        "option_2": ["HO2", "0.1", "\\ce{C2H6 + O2 -> CO2 + H2O}"],
        "option_3": ["H2O2", "-1", "\\ce{C2H6 + 3O2 -> 2CO2 + 3H2O}"],
        "option_4": ["OH2", "14", "\\ce{2C2H6 + 5O2 -> 4CO2 + 2H2O}"],
        "correct_answer": ["A", "A", "A"]
    })
    st.dataframe(sample_data)
    
    csv = sample_data.to_csv(index=False)
    st.download_button(
        label="📥 Download Sample CSV",
        data=csv,
        file_name="sample_mcq_format.csv",
        mime="text/csv"
    )