import streamlit as st
import sys
from io import StringIO
from contextlib import redirect_stdout
import json
from datetime import datetime
from typing import Dict, List, Optional
import base64
from PIL import Image
import io
import os
import zipfile
from pathlib import Path
import shutil

class CodePlayground:
    def __init__(self):
        self.initialize_session_state()

    @staticmethod
    def initialize_session_state():
        """Initialize session state for artifacts and history"""
        if 'artifacts' not in st.session_state:
            st.session_state.artifacts = []
        if 'code_history' not in st.session_state:
            st.session_state.code_history = []
        if 'current_artifact' not in st.session_state:
            st.session_state.current_artifact = None

    @staticmethod
    def display():
        """Display the enhanced code playground interface"""
        st.header("Interactive Learning Playground")
        
        # Create tabs for different views
        code_tab, artifacts_tab = st.tabs(["Code Editor", "Learning Artifacts"])
        
        playground = CodePlayground()
        
        with code_tab:
            playground.display_code_editor()
        
        with artifacts_tab:
            playground.display_artifacts()

    def display_code_editor(self):
        """Display the code editor interface"""
        # Code editor
        code = st.text_area(
            "Write your Python code here",
            height=200,
            key="code_editor",
            help="Write Python code to execute. You can use print() to see outputs."
        )
        
        # Button columns
        col1, col2, col3 = st.columns(3)
        
        with col1:
            if st.button("Run Code", key="run_code"):
                output, artifact_data = self.execute_code(code)
                if artifact_data:
                    self.store_artifact(artifact_data)
        
        with col2:
            if st.button("Save as Artifact", key="save_artifact"):
                self.create_manual_artifact(code)
        
        with col3:
            if st.button("Clear Output", key="clear_output"):
                if 'code_output' in st.session_state:
                    del st.session_state.code_output

        # Display output
        if 'code_output' in st.session_state:
            st.markdown("### Output:")
            st.code(st.session_state.code_output)

    def display_artifacts(self):
        """Display stored learning artifacts"""
        st.subheader("Learning Artifacts")
        
        if not st.session_state.artifacts:
            st.info("No artifacts created yet. Run some code or create manual artifacts to see them here!")
            return
        
        # Filter options
        filter_type = st.selectbox(
            "Filter by type",
            options=["All", "Code", "Image", "Text", "Mixed Media"],
            key="artifact_filter"
        )
        
        # Export buttons
        col1, col2 = st.columns(2)
        with col1:
            if st.button("Export All Artifacts"):
                try:
                    # Use temporary directory for export
                    temp_dir = Path("temp_export")
                    temp_dir.mkdir(exist_ok=True)
                    
                    zip_path = self.export_all_artifacts(str(temp_dir))
                    
                    # Read the zip file
                    with open(zip_path, "rb") as fp:
                        zip_data = fp.read()
                        
                    # Create download button
                    st.download_button(
                        label="Download Artifacts",
                        data=zip_data,
                        file_name="learning_artifacts.zip",
                        mime="application/zip"
                    )
                    st.success("Artifacts exported successfully!")
                except Exception as e:
                    st.error(f"Error exporting artifacts: {str(e)}")
                finally:
                    # Cleanup
                    if temp_dir.exists():
                        shutil.rmtree(temp_dir)

        # Display artifacts in reverse chronological order
        for artifact in reversed(st.session_state.artifacts):
            if filter_type == "All" or artifact['type'] == filter_type:
                with st.expander(f"{artifact['title']} - {artifact['date']}", expanded=False):
                    st.write(f"**Type:** {artifact['type']}")
                    st.write(f"**Created:** {artifact['date']}")
                    
                    # Display content based on type
                    if artifact['type'] == "Code":
                        st.code(artifact['content'], language="python")
                    elif artifact['type'] == "Image":
                        st.image(base64.b64decode(artifact['content']))
                    elif artifact['type'] == "Text":
                        st.write(artifact['content'])
                    elif artifact['type'] == "Mixed Media":
                        for item in artifact['content']:
                            if item['type'] == "code":
                                st.code(item['data'], language="python")
                            elif item['type'] == "image":
                                st.image(base64.b64decode(item['data']))
                            elif item['type'] == "text":
                                st.write(item['data'])
                    
                    # Artifact actions
                    col1, col2 = st.columns(2)
                    with col1:
                        if st.button("Load in Editor", key=f"load_{artifact['id']}"):
                            if artifact['type'] == "Code":
                                st.session_state.code_editor = artifact['content']
                                st.rerun()
                    
                    with col2:
                        if st.button("Delete", key=f"delete_{artifact['id']}"):
                            self.delete_artifact(artifact['id'])
                            st.rerun()

    def execute_code(self, code: str) -> tuple:
        """Execute code and capture outputs"""
        output = StringIO()
        artifact_data = None
        
        try:
            # Create a restricted environment
            globals_dict = {
                '__builtins__': __builtins__,
                'print': print,
                'Image': Image,
                'io': io,
                'base64': base64,
                'create_artifact': create_artifact
            }
            
            # Capture stdout
            with redirect_stdout(output):
                # Execute the code
                exec(code, globals_dict, {})
            
            # Store output
            output_text = output.getvalue()
            st.session_state.code_output = output_text
            
            # Check for generated artifacts in globals
            if '_artifact_output' in globals_dict:
                artifact_data = globals_dict['_artifact_output']
            
            # Store in history
            self.store_in_history(code, output_text)
            
            return output_text, artifact_data
            
        except Exception as e:
            error_msg = f"Error: {str(e)}"
            st.session_state.code_output = error_msg
            return error_msg, None

    def export_artifact(self, artifact: Dict, export_dir: str):
        """Export a single artifact to the specified directory"""
        # Create directory if it doesn't exist
        artifact_dir = Path(export_dir) / str(artifact['id'])
        artifact_dir.mkdir(parents=True, exist_ok=True)

        # Export metadata
        metadata = {
            'id': artifact['id'],
            'title': artifact['title'],
            'type': artifact['type'],
            'date': artifact['date']
        }
        
        with open(artifact_dir / 'metadata.json', 'w') as f:
            json.dump(metadata, f, indent=2)

        # Export content based on type
        if artifact['type'] == "Code":
            with open(artifact_dir / 'content.py', 'w') as f:
                f.write(artifact['content'])
        
        elif artifact['type'] == "Image":
            img_data = base64.b64decode(artifact['content'])
            with open(artifact_dir / 'content.png', 'wb') as f:
                f.write(img_data)
        
        elif artifact['type'] == "Text":
            with open(artifact_dir / 'content.txt', 'w') as f:
                f.write(artifact['content'])
        
        elif artifact['type'] == "Mixed Media":
            media_dir = artifact_dir / 'content'
            media_dir.mkdir(exist_ok=True)
            
            for idx, item in enumerate(artifact['content']):
                if item['type'] == 'code':
                    with open(media_dir / f'code_{idx}.py', 'w') as f:
                        f.write(item['data'])
                elif item['type'] == 'image':
                    img_data = base64.b64decode(item['data'])
                    with open(media_dir / f'image_{idx}.png', 'wb') as f:
                        f.write(img_data)
                elif item['type'] == 'text':
                    with open(media_dir / f'text_{idx}.txt', 'w') as f:
                        f.write(item['data'])

    def export_all_artifacts(self, export_dir: str):
        """Export all artifacts to a directory and create a zip file"""
        # Create temporary directory for exports
        temp_dir = Path(export_dir) / 'temp_artifacts'
        temp_dir.mkdir(parents=True, exist_ok=True)

        try:
            # Export each artifact
            for artifact in st.session_state.artifacts:
                self.export_artifact(artifact, str(temp_dir))

            # Create zip file
            zip_path = Path(export_dir) / 'learning_artifacts.zip'
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for root, _, files in os.walk(temp_dir):
                    for file in files:
                        file_path = Path(root) / file
                        arcname = file_path.relative_to(temp_dir)
                        zipf.write(file_path, arcname)

            return str(zip_path)
        finally:
            # Clean up temporary directory
            shutil.rmtree(temp_dir)

    def store_artifact(self, artifact_data: Dict):
        """Store a new artifact"""
        artifact = {
            'id': len(st.session_state.artifacts),
            'date': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            **artifact_data
        }
        st.session_state.artifacts.append(artifact)
        st.success(f"Created new artifact: {artifact['title']}")

    def create_manual_artifact(self, code: str):
        """Create a manual artifact from current code"""
        title = st.text_input("Artifact Title")
        if title:
            artifact = {
                'id': len(st.session_state.artifacts),
                'title': title,
                'type': "Code",
                'content': code,
                'date': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
            st.session_state.artifacts.append(artifact)
            st.success(f"Created new artifact: {title}")

    def delete_artifact(self, artifact_id: int):
        """Delete an artifact"""
        st.session_state.artifacts = [
            a for a in st.session_state.artifacts if a['id'] != artifact_id
        ]

    def store_in_history(self, code: str, output: str):
        """Store code execution in history"""
        history_entry = {
            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            'code': code,
            'output': output
        }
        st.session_state.code_history.append(history_entry)

def create_artifact(title: str, content: any, artifact_type: str = "Mixed Media"):
    """
    Create an artifact from code execution.
    Usage example in code:
    
    # Create a text artifact
    create_artifact("My Text", "Hello World", "Text")
    
    # Create an image artifact
    img = Image.new('RGB', (100, 100), color='red')
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    create_artifact("My Image", img_str, "Image")
    """
    globals()['_artifact_output'] = {
        'title': title,
        'type': artifact_type,
        'content': content
    }