import streamlit as st
import streamlit.components.v1 as components

# Initialize session state
if 'settings' not in st.session_state:
    st.session_state['settings'] = {}

# Function to update widgets based on saved settings
def load_settings():
    settings = st.session_state['settings']
    st.session_state.cube_color = settings.get("cube_color", "#00ff00")
    st.session_state.rotation_speed = settings.get("rotation_speed", 0.01)
    st.session_state.cube_position_x = settings.get("cube_position_x", 0)
    st.session_state.cube_position_y = settings.get("cube_position_y", 0)
    st.session_state.cube_position_z = settings.get("cube_position_z", 0)
    st.session_state.cube_scale = settings.get("cube_scale", 1.0)
    st.session_state.show_bounding_box = settings.get("show_bounding_box", True)
    st.session_state.spotlight_intensity = settings.get("spotlight_intensity", 1.0)
    st.session_state.ambient_light_intensity = settings.get("ambient_light_intensity", 0.5)
    st.session_state.camera_view = settings.get("camera_view", "Default")
    st.write("Loaded settings into session state:", settings)

# Create two columns for layout
col1, col2 = st.columns([1, 3])

with col1:
    # Streamlit widgets to control the scene
    cube_color = st.color_picker("Pick a cube color", st.session_state.get('cube_color', "#00ff00"))
    rotation_speed = st.slider("Rotation Speed", min_value=0.01, max_value=0.1, value=st.session_state.get('rotation_speed', 0.01))
    cube_position_x = st.slider("Cube Position X", min_value=-10, max_value=10, value=st.session_state.get('cube_position_x', 0))
    cube_position_y = st.slider("Cube Position Y", min_value=-10, max_value=10, value=st.session_state.get('cube_position_y', 0))
    cube_position_z = st.slider("Cube Position Z", min_value=-10, max_value=10, value=st.session_state.get('cube_position_z', 0))
    cube_scale = st.slider("Cube Scale", min_value=0.1, max_value=5.0, value=st.session_state.get('cube_scale', 1.0))
    show_bounding_box = st.checkbox("Show Bounding Box", value=st.session_state.get('show_bounding_box', True))
    spotlight_intensity = st.slider("Spotlight Intensity", min_value=0.0, max_value=2.0, value=st.session_state.get('spotlight_intensity', 1.0))
    ambient_light_intensity = st.slider("Ambient Light Intensity", min_value=0.0, max_value=2.0, value=st.session_state.get('ambient_light_intensity', 0.5))
    camera_view = st.selectbox("Camera View", ["Default", "Top", "Side"], index=["Default", "Top", "Side"].index(st.session_state.get('camera_view', "Default")))

    # Save settings button
    if st.button("Save Settings"):
        st.session_state['settings'] = {
            "cube_color": cube_color,
            "rotation_speed": rotation_speed,
            "cube_position_x": cube_position_x,
            "cube_position_y": cube_position_y,
            "cube_position_z": cube_position_z,
            "cube_scale": cube_scale,
            "show_bounding_box": show_bounding_box,
            "spotlight_intensity": spotlight_intensity,
            "ambient_light_intensity": ambient_light_intensity,
            "camera_view": camera_view
        }
        st.write("Settings saved:", st.session_state['settings'])

    # Load settings button
    if st.button("Load Settings"):
        st.write("Loading settings...")
        load_settings()
        st.write("Settings loaded:", st.session_state['settings'])
        st.rerun()

def three_js_component(cube_color, rotation_speed, cube_position_x, cube_position_y, cube_position_z, cube_scale, show_bounding_box, spotlight_intensity, ambient_light_intensity, camera_view):
    component_code = f"""
    <div id="threejs-container" style="width: 100%; height: 600px; background-color: #000;"></div>
    <script>
        (function() {{
            if (document.getElementById('threejs-container-script')) {{
                return;
            }}

            function loadScript(url, callback) {{
                console.log('Loading script:', url);
                var script = document.createElement("script");
                script.type = "text/javascript";
                
                script.onload = function() {{
                    console.log('Loaded script:', url);
                    callback();
                }};
                
                script.onerror = function() {{
                    console.error('Error loading script:', url);
                }};
                
                script.src = url;
                document.getElementsByTagName("head")[0].appendChild(script);
            }}

            function initializeScene() {{
                console.log('Initializing Scene');
                const container = document.getElementById('threejs-container');
                if (!container) {{
                    console.error('Container not found!');
                    return;
                }}

                const scene = new THREE.Scene();
                console.log('Created scene');
                const camera = new THREE.PerspectiveCamera(75, container.clientWidth / container.clientHeight, 0.1, 1000);
                console.log('Created camera');

                const renderer = new THREE.WebGLRenderer();
                renderer.setSize(container.clientWidth, container.clientHeight);
                renderer.shadowMap.enabled = true;
                container.appendChild(renderer.domElement);
                console.log('Created renderer and appended to container');

                const controls = new THREE.OrbitControls(camera, renderer.domElement);
                controls.enableDamping = true;
                controls.dampingFactor = 0.25;
                controls.screenSpacePanning = false;
                controls.minDistance = 1;
                controls.maxDistance = 100;
                controls.maxPolarAngle = Math.PI / 2;
                console.log('Initialized Orbit Controls');

                // Add a spotlight for lighting
                const spotLight = new THREE.SpotLight(0xffffff, {spotlight_intensity});
                spotLight.position.set(10, 10, 10);
                spotLight.castShadow = true;
                scene.add(spotLight);
                console.log('Spotlight added');

                // Add ambient light
                const ambientLight = new THREE.AmbientLight(0x404040, {ambient_light_intensity}); // Soft white light
                scene.add(ambientLight);
                console.log('Ambient light added');

                // Add a plane to receive shadows
                const planeGeometry = new THREE.PlaneGeometry(200, 200);
                const planeMaterial = new THREE.ShadowMaterial({{ opacity: 0.5 }});
                const plane = new THREE.Mesh(planeGeometry, planeMaterial);
                plane.rotation.x = -Math.PI / 2;
                plane.position.y = -5;
                plane.receiveShadow = true;
                scene.add(plane);
                console.log('Plane added');

                // Load texture
                const loader = new THREE.TextureLoader();
                loader.load('https://threejs.org/examples/textures/crate.gif', function(texture) {{
                    // Add a rotating cube with texture
                    const cubeGeometry = new THREE.BoxGeometry();
                    const cubeMaterial = new THREE.MeshStandardMaterial({{ map: texture, color: '{cube_color}' }});
                    const cube = new THREE.Mesh(cubeGeometry, cubeMaterial);
                    cube.castShadow = true;
                    cube.position.set({cube_position_x}, {cube_position_y}, {cube_position_z});
                    cube.scale.set({cube_scale}, {cube_scale}, {cube_scale});
                    scene.add(cube);
                    console.log('Cube added');

                    // Add a sphere
                    const sphereGeometry = new THREE.SphereGeometry(1, 32, 32);
                    const sphereMaterial = new THREE.MeshStandardMaterial({{ color: 0xff0000 }});
                    const sphere = new THREE.Mesh(sphereGeometry, sphereMaterial);
                    sphere.position.set(3, 1, 0);
                    sphere.castShadow = true;
                    scene.add(sphere);
                    console.log('Sphere added');

                    // Load GLTF model
                    const gltfLoader = new THREE.GLTFLoader();
                    gltfLoader.load('https://threejs.org/examples/models/gltf/Flamingo.glb', function(gltf) {{
                        const model = gltf.scene;
                        model.position.set(0, 0, -5);
                        model.scale.set(0.05, 0.05, 0.05);  // Scale down the model
                        scene.add(model);
                        console.log('GLTF model added');

                        // Add a bounding box helper
                        const box = new THREE.Box3().setFromObject(model);
                        const helper = new THREE.Box3Helper(box, 0xffff00);
                        if ({str(show_bounding_box).lower()}) {{
                            scene.add(helper);
                            console.log('Bounding box added');
                        }}

                        // Animate the model
                        const mixer = new THREE.AnimationMixer(model);
                        gltf.animations.forEach((clip) => {{
                            mixer.clipAction(clip).play();
                        }});
                        console.log('Animation added');

                        camera.position.z = 10;

                        if ('{camera_view}' === 'Top') {{
                            camera.position.set(0, 10, 0);
                            camera.lookAt(0, 0, 0);
                            console.log('Switched to top view');
                        }} else if ('{camera_view}' === 'Side') {{
                            camera.position.set(10, 0, 0);
                            camera.lookAt(0, 0, 0);
                            console.log('Switched to side view');
                        }}

                        console.log('Starting animation');
                        function animate() {{
                            requestAnimationFrame(animate);
                            cube.rotation.x += {rotation_speed};
                            cube.rotation.y += {rotation_speed};
                            controls.update();
                            mixer.update(0.01);  // Update the animation
                            renderer.render(scene, camera);
                            console.log('Rendered frame');
                        }}

                        animate();

                        window.addEventListener('resize', () => {{
                            console.log('Resizing window');
                            camera.aspect = container.clientWidth / container.clientHeight;
                            camera.updateProjectionMatrix();
                            renderer.setSize(container.clientWidth, container.clientHeight);
                        }});
                    }});
                }});
            }}

            loadScript("https://unpkg.com/three@0.130.1/build/three.min.js", function() {{
                loadScript("https://unpkg.com/three@0.130.1/examples/js/controls/OrbitControls.js", function() {{
                    loadScript("https://unpkg.com/three@0.130.1/examples/js/loaders/GLTFLoader.js", function() {{
                        console.log('Scripts loaded, initializing scene');
                        initializeScene();
                    }});
                }});
            }});
        }})();
    </script>
    """
    components.html(component_code, height=600)

with col2:
    st.title("3D Streamlit Component with Enhanced UI Controls")
    three_js_component(cube_color, rotation_speed, cube_position_x, cube_position_y, cube_position_z, cube_scale, show_bounding_box, spotlight_intensity, ambient_light_intensity, camera_view)