import streamlit as st
import io
import sys
from contextlib import redirect_stdout
import re
import mistune

def extract_python_code(markdown_text):
    """Extract Python code blocks from markdown text."""
    pattern = r"```python\s*(.*?)\s*```"
    matches = re.findall(pattern, markdown_text, re.DOTALL)
    return matches

def execute_code(code):
    """Execute Python code and return the output."""
    # Create string buffer to capture output
    buffer = io.StringIO()
    
    # Create a dictionary for local variables
    local_vars = {}
    
    try:
        # Redirect stdout to our buffer
        with redirect_stdout(buffer):
            # Execute the code
            exec(code, {}, local_vars)
        
        # Get the output
        output = buffer.getvalue()
        return output, None
    except Exception as e:
        return None, str(e)
    finally:
        buffer.close()

def main():
    st.set_page_config(page_title="Python Code Executor", layout="wide")
    
    st.title("📝 Python Code Executor")
    st.markdown("Enter Python code directly or upload a .py/.md file")
    
    # File uploader
    uploaded_file = st.file_uploader("Upload a Python (.py) or Markdown (.md) file", 
                                   type=['py', 'md'])
    
    # Code input area
    if uploaded_file is None:
        code = st.text_area("Or enter your Python code here:", 
                           height=200,
                           placeholder="Enter your Python code...")
    else:
        content = uploaded_file.getvalue().decode()
        if uploaded_file.type == "text/markdown":
            # Extract Python code from markdown
            code_blocks = extract_python_code(content)
            if code_blocks:
                code = code_blocks[0]  # Use the first Python code block found
            else:
                st.error("No Python code blocks found in the markdown file!")
                return
        else:
            # Assume it's a Python file
            code = content
        
        st.code(code, language='python')
    
    # Execute button
    if st.button("▶️ Run Code"):
        if code:
            st.markdown("### Output:")
            output, error = execute_code(code)
            
            if error:
                st.error(f"Error:\n{error}")
            elif output:
                st.code(output)
            else:
                st.info("Code executed successfully with no output.")
        else:
            st.warning("Please enter some code to execute!")
    
    # Add some usage information
    with st.expander("ℹ️ How to use"):
        st.markdown("""
        1. Either upload a Python (.py) or Markdown (.md) file containing Python code
        2. Or enter Python code directly in the text area
        3. Click the 'Run Code' button to execute
        
        **Note:** For markdown files, the code should be in Python code blocks like:
        ````
        ```python
        print("Hello World!")
        ```
        ````
        """)
    
    # Add safety information
    with st.expander("⚠️ Safety Notice"):
        st.markdown("""
        - The code executor runs in a restricted environment
        - Some operations may be limited for security reasons
        - Be careful when executing code from unknown sources
        """)

if __name__ == "__main__":
    main()