"""
Code Agent nodes for autonomous code execution.
"""

from typing import Any, Dict, Optional
from ybagent.node import BaseNode, NodeType, NodeInput, NodeOutput
from ybagent.sandbox import (
    DockerExecutor,
    RestrictedExecutor,
    get_monitor,
    SecurityError
)
import json
import httpx
from ybagent.utils.display import DisplayManager


class CodeAgentNode(BaseNode):
    """
    Autonomous code execution agent.
    
    This node:
    1. Receives a task description
    2. Generates Python code to solve the task (via LLM)
    3. Executes code in secure sandbox
    4. Returns results with full execution logs
    
    Security:
    - All code runs in isolated sandbox (Docker or RestrictedPython)
    - No file system or network access by default
    - Resource limits enforced
    - All executions logged for audit
    """
    
    def __init__(
        self,
        id: str,
        sandbox_type: str = "auto",  # "auto", "docker", or "restricted"
        provider: str = "openai",    # "openai", "ollama", "huggingface"
        llm_model: str = "gpt-4",
        llm_api_key: str = "",
        timeout: int = 30,
        memory_limit: str = "256m",
        cpu_limit: float = 1.0,
        enable_network: bool = False,
        additional_imports: list = None,  # List of allowed/installed packages
        max_repairs: int = 3,
        **config
    ):
        """
        Initialize Code Agent node.
        
        Args:
            id: Node ID
            sandbox_type: "auto" (default), "docker", or "restricted"
            provider: LLM provider ("openai", "ollama", "huggingface")
            llm_model: LLM model for code generation
            llm_api_key: API key for LLM (if required)
            timeout: Execution timeout in seconds
            memory_limit: Memory limit (Docker only)
            cpu_limit: CPU cores limit (Docker only)
            enable_network: Allow network access (Docker only)
            additional_imports: List of packages to allow/install (e.g. ['plotly', 'numpy'])
            max_repairs: Maximum number of self-correction attempts
            **config: Additional configuration
        """
        config.update({
            "sandbox_type": sandbox_type,
            "provider": provider,
            "llm_model": llm_model,
            "llm_api_key": llm_api_key,
            "timeout": timeout,
            "memory_limit": memory_limit,
            "cpu_limit": cpu_limit,
            "enable_network": enable_network,
            "additional_imports": additional_imports or [],
            "max_repairs": max_repairs,
        })
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.CODE_AGENT, **config)
        
        # Initialize sandbox executor
        self.executor = self._create_executor()
        
        # Get global monitor
        self.monitor = get_monitor()

    def _create_executor(self):
        """Create sandbox executor based on configuration."""
        sandbox_type = self.metadata.config.get("sandbox_type", "auto")
        raw_imports = self.metadata.config.get("additional_imports", [])
        
        # Ensure io and base64 are always allowed for image handling
        if "io" not in raw_imports: 
            raw_imports.append("io")
        if "base64" not in raw_imports: 
            raw_imports.append("base64")
        
        # Parse imports to get clean package names (e.g. "pandas as pd" -> "pandas")
        packages = self._parse_package_names(raw_imports)
        
        # Auto-detect: try Docker first, fallback to RestrictedPython
        if sandbox_type == "auto":
            try:
                executor = DockerExecutor(
                    timeout=self.metadata.config.get("timeout", 30),
                    memory_limit=self.metadata.config.get("memory_limit", "256m"),
                    cpu_limit=self.metadata.config.get("cpu_limit", 1.0),
                    network_disabled=not self.metadata.config.get("enable_network", False),
                    enable_pip=True if packages else False
                )
                
                # Install packages if needed
                if packages:
                    print(f"📦 Installing packages in Docker sandbox: {', '.join(packages)}...")
                    if executor.install_packages(packages):
                        print("✅ Packages installed successfully")
                    else:
                        print("⚠️  Failed to install packages")
                
                print("✅ Using Docker sandbox (maximum security)")
                self.metadata.config["sandbox_type"] = "docker"
                return executor
            except RuntimeError:
                print("ℹ️  Docker not available, using RestrictedPython sandbox")
                
                # Auto-install packages locally for RestrictedExecutor
                if packages:
                    self._install_local_packages(packages)
                
                self.metadata.config["sandbox_type"] = "restricted"
                return RestrictedExecutor(
                    timeout=self.metadata.config.get("timeout", 10),
                    allowed_imports=packages
                )
        
        # Explicit Docker
        elif sandbox_type == "docker":
            try:
                executor = DockerExecutor(
                    timeout=self.metadata.config.get("timeout", 30),
                    memory_limit=self.metadata.config.get("memory_limit", "256m"),
                    cpu_limit=self.metadata.config.get("cpu_limit", 1.0),
                    network_disabled=not self.metadata.config.get("enable_network", False),
                    enable_pip=True if packages else False
                )
                
                if packages:
                    print(f"📦 Installing packages in Docker sandbox: {', '.join(packages)}...")
                    executor.install_packages(packages)
                    
                return executor
            except RuntimeError as e:
                print(f"⚠️  Docker not available, falling back to RestrictedPython: {e}")
                
                if packages:
                    self._install_local_packages(packages)
                    
                self.metadata.config["sandbox_type"] = "restricted"
                return RestrictedExecutor(
                    timeout=self.metadata.config.get("timeout", 10),
                    allowed_imports=packages
                )
        
        # Explicit RestrictedPython
        elif sandbox_type == "restricted":
            if packages:
                self._install_local_packages(packages)
                
            return RestrictedExecutor(
                timeout=self.metadata.config.get("timeout", 10),
                allowed_imports=packages
            )
        
        else:
            raise ValueError(f"Unknown sandbox type: {sandbox_type}. Use 'auto', 'docker', or 'restricted'")

    def _parse_package_names(self, imports: list) -> list:
        """Extract clean package names from import strings."""
        clean_packages = set()
        for imp in imports:
            # Handle "pandas as pd" -> "pandas"
            name = imp.split(' as ')[0].strip()
            # Handle "plotly.express" -> "plotly"
            root_pkg = name.split('.')[0]
            clean_packages.add(root_pkg)
        return list(clean_packages)

    def _install_local_packages(self, packages: list):
        """Install packages locally using pip, skipping standard library modules."""
        import subprocess
        import sys
        
        # Filter out modules that are part of the Python standard library
        stdlib_modules = {"io", "base64", "sys", "json", "os", "time", "re", "math"}
        packages_to_install = [p for p in packages if p not in stdlib_modules]
        
        if not packages_to_install:
            print("✅ No external packages need to be installed (all are stdlib).")
            return
        
        print(f"📦 Installing packages locally for RestrictedPython: {', '.join(packages_to_install)}...")
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install"] + packages_to_install)
            print("✅ Packages installed successfully")
        except subprocess.CalledProcessError as e:
            print(f"⚠️  Failed to install packages locally: {e}")
            print("Please install them manually: pip install " + " ".join(packages_to_install))
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(
                name="task",
                type="str",
                required=True,
                description="Task description for the agent"
            ),
            NodeInput(
                name="context",
                type="dict",
                required=False,
                description="Optional context variables"
            ),
            NodeInput(
                name="code",
                type="str",
                required=False,
                description="Pre-written code (skips LLM generation)"
            ),
            NodeInput(
                name="llm_model",
                type="str",
                required=False,
                description="Override LLM model (e.g. 'gpt-4', 'llama2')"
            ),
        ]
        self.metadata.outputs = [
            NodeOutput(
                name="result",
                type="any",
                description="Execution result"
            ),
            NodeOutput(
                name="code",
                type="str",
                description="Generated/executed code"
            ),
            NodeOutput(
                name="logs",
                type="dict",
                description="Execution logs and metrics"
            ),
        ]
    
    def _process_visualization_output(self, code: str, result) -> Any:
        """
        Process the execution result to extract visualizations.
        
        Args:
            code: The executed code
            result: Execution result object
            
        Returns:
            Processed output (visualization object or text)
        """
        if 'plotly' in code.lower() and 'fig' in code:
            try:
                if hasattr(result, 'return_value') and result.return_value:
                    try:
                        fig_obj = result.return_value
                        if hasattr(fig_obj, 'to_html'):
                            return fig_obj.to_html()
                    except:
                        pass
                return result.output
            except:
                return result.output
        
        return result.output
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute code agent workflow.
        
        Args:
            inputs: Task description and optional context
            context: Execution context
        
        Returns:
            Execution result, code, and logs
        """
        task = inputs.get("task")
        exec_context = inputs.get("context", {})
        pre_code = inputs.get("code")
        model_override = inputs.get("llm_model")
        state = context.get("state") if isinstance(context, dict) else None
        if state:
            state.set_variable("suppress_llm_response", True)
            state.set_variable("suppress_vector_display", True)
        
        if not task and not pre_code:
            raise ValueError("Either 'task' or 'code' must be provided")
        
        
        DisplayManager.header(f"🤖 Code Agent: {self.metadata.name}", level=3)
        if task:
            # Extract clean task text if it's a dictionary
            if isinstance(task, dict):
                task_text = task.get('task', str(task))
            else:
                task_text = str(task)
            DisplayManager.info(f"Task: {task_text}")


        # Generate code if not provided
        if pre_code:
            code = pre_code
        else:
            code = self._generate_code(task, exec_context, model_override)
        
        # Modify code to capture figure object if it's a visualization
        if 'plotly' in code.lower() and 'fig' in code:
            display_calls = ['fig.show(', 'plt.show(', 'print(fig.to_html', 'display(']
            if not code.strip().endswith('fig') and not any(dc in code for dc in display_calls):
                code = code.rstrip() + '\nfig'
        
        # Display generated code
        DisplayManager.code(code, title="Generated Code")
        
        # Self-correction loop
        max_repairs = self.metadata.config.get("max_repairs", 3)
        attempts = 0
        last_error = None
        result = None
        
        while attempts <= max_repairs:
            if attempts > 0:
                DisplayManager.markdown(f"🔄 **Attempt {attempts}/{max_repairs}:** Self-correcting code based on error...")
                # Regenerate code with error context
                code = self._generate_code(task, exec_context, model_override, error_context=last_error, previous_code=code)
                DisplayManager.code(code, title="Corrected Code")
            
            # Execute
            result = self.executor.execute(code, exec_context)
            
            if result.success:
                break
                
            # If failed, prepare for retry
            last_error = result.error
            DisplayManager.error(f"Execution failed: {last_error}")
            attempts += 1
            
            if attempts > max_repairs:
                DisplayManager.error(f"❌ Failed after {max_repairs} repair attempts.")
        
        # Process output
        processed_output = self._process_visualization_output(code, result)

        # Avoid double visualization in RestrictedPython when code explicitly displays
        display_calls = ['fig.show(', 'plt.show(', 'print(fig.to_html', 'display(']
        has_explicit_display = any(dc in code for dc in display_calls)
        sandbox_type = self.metadata.config.get("sandbox_type", "restricted")

        should_render = not (has_explicit_display and sandbox_type == "restricted")

        if should_render:
            # Display visualization if it's HTML
            if processed_output and isinstance(processed_output, str) and '<html>' in processed_output.lower():
                DisplayManager.render_html(processed_output, title="Visualization Output")
            elif processed_output and isinstance(processed_output, str) and processed_output.strip():
                # Display text output if not empty
                DisplayManager.collapsible_log("Output", processed_output, level="info")
        
        # Log execution
        DisplayManager.collapsible_log("Execution Logs", f"Success: {result.success}\nExecution Time: {result.execution_time}s\nOutput Type: {type(processed_output).__name__}\nError: {result.error}")
        
        if result.success:
            DisplayManager.success("Code executed successfully.")
        
        # Log execution
        self.monitor.log_execution(
            code=code,
            success=result.success,
            output=str(processed_output)[:1000] if processed_output else None,  # Truncate for logging
            error=result.error,
            execution_time=result.execution_time,
            memory_used=result.memory_used,
            sandbox_type=self.metadata.config.get("sandbox_type", "docker"),
            context=exec_context
        )
        
        # Prepare logs
        logs = {
            "success": result.success,
            "execution_time": result.execution_time,
            "memory_used": result.memory_used,
            "exit_code": result.exit_code,
            "sandbox_type": self.metadata.config.get("sandbox_type"),
            "attempts": attempts + 1,
            "repaired": attempts > 0
        }
        
        if result.success:
            return {
                "result": processed_output,
                "code": code,
                "logs": logs
            }
        else:
            # Return error but don't fail the node
            return {
                "result": None,
                "code": code,
                "logs": {
                    **logs,
                    "error": result.error
                }
            }

    def _generate_code(self, task: str, context: Dict[str, Any], model_override: Optional[str] = None, error_context: Optional[str] = None, previous_code: Optional[str] = None) -> str:
        """
        Generate Python code using configured LLM provider.
        
        Args:
            task: Task description
            context: Available context variables
            model_override: Optional model to use instead of config
            error_context: Error message from previous execution (for self-correction)
            previous_code: Code that caused the error (for self-correction)
        
        Returns:
            Generated Python code
        """
        provider = self.metadata.config.get("provider", "openai").lower()
        # Use override if provided, otherwise config, otherwise default
        model = model_override or self.metadata.config.get("llm_model", "gpt-4")
        api_key = self.metadata.config.get("llm_api_key", "")
        
        # Build prompt
        context_str = json.dumps(context, indent=2) if context else "None"
        
        # Build system prompt
        # Get list of available modules for the prompt
        available_modules = self.metadata.config.get("additional_imports", [])
        modules_info = ""
        if available_modules:
            # Extract module names and aliases
            module_list = []
            for imp in available_modules:
                if ' as ' in imp:
                    module_list.append(f"import {imp}")
                else:
                    module_list.append(f"import {imp.strip()}")
            modules_info = f"\n\nAvailable modules (you can import these):\n" + "\n".join(module_list)
        
        system_prompt = f"""You are a Python code generation assistant. Generate clean, executable Python code based on the task description.

Rules:

Output ONLY valid Python code, no explanations or markdown
Use available context variables if provided
Keep code concise and efficient
Handle errors gracefully
For visualizations, ensure the figure object is returned or displayed just show it once not twice.
Do not use input() or any blocking operations
You can use import statements for the available modules listed below.{modules_info}
MANDATORY: Read data from the original file using the appropriate pandas function:
    -For CSV files: df = pd.read_csv(r'{{original_file_path}}')
    -For Excel files: df = pd.read_excel(r'{{original_file_path}}')
    -For JSON files: df = pd.read_json(r'{{original_file_path}}')
    -For Parquet files: df = pd.read_parquet(r'{{original_file_path}}')
    -For other formats: use corresponding pd.read_* methods
"""

        # Build user prompt
        if error_context and previous_code:
            user_prompt = f"""Previous code failed with error:
{error_context}

Previous code:
```python
{previous_code}
```

Task: {task}
Available context: {context_str}

Generate corrected Python code that fixes the error:"""
        else:
            user_prompt = f"""Task: {task}
Available context: {context_str}

Generate Python code to complete this task:"""
        
        try:
            # 1. OpenAI Provider
            if provider == "openai":
                if not api_key:
                    raise ValueError("OpenAI API key not provided")
                    
                response = httpx.post(
                    "https://api.openai.com/v1/chat/completions",
                    headers={
                        "Authorization": f"Bearer {api_key}",
                        "Content-Type": "application/json"
                    },
                    json={
                        "model": model,
                        "messages": [
                            {"role": "system", "content": system_prompt},
                            {"role": "user", "content": user_prompt}
                        ],
                        "temperature": 0.3,
                        "max_tokens": 1000
                    },
                    timeout=30.0
                )
                response.raise_for_status()
                content = response.json()["choices"][0]["message"]["content"]

            # 2. Ollama Provider (Local)
            elif provider == "ollama":
                # Default to localhost:11434
                base_url = self.metadata.config.get("base_url", "http://localhost:11434").rstrip("/")
                
                try:
                    response = httpx.post(
                        f"{base_url}/api/generate",
                        json={
                            "model": model,
                            "prompt": f"{system_prompt}\n\n{user_prompt}",
                            "stream": False,
                            "options": {
                                "temperature": 0.3
                            }
                        },
                        timeout=60.0
                    )
                    response.raise_for_status()
                    content = response.json()["response"]
                except httpx.HTTPStatusError as e:
                    if e.response.status_code == 404:
                        # Check if it's a model missing error
                        try:
                            error_msg = e.response.json().get("error", "")
                        except:
                            error_msg = e.response.text
                            
                        if "model" in error_msg or "not found" in error_msg:
                            raise ValueError(f"Ollama model '{model}' not found. Please run: ollama pull {model}")
                    raise e

            # 3. HuggingFace Provider
            elif provider == "huggingface":
                if not api_key:
                    raise ValueError("HuggingFace API key not provided")
                
                response = httpx.post(
                    f"https://api-inference.huggingface.co/models/{model}",
                    headers={"Authorization": f"Bearer {api_key}"},
                    json={
                        "inputs": f"{system_prompt}\n\n{user_prompt}",
                        "parameters": {
                            "max_new_tokens": 1000,
                            "temperature": 0.3,
                            "return_full_text": False
                        }
                    },
                    timeout=60.0
                )
                response.raise_for_status()
                # HF API format varies, usually list of dicts
                result = response.json()
                if isinstance(result, list) and len(result) > 0:
                    content = result[0].get("generated_text", "")
                else:
                    content = str(result)

            else:
                raise ValueError(f"Unsupported provider: {provider}")

            # Clean up markdown code blocks
            code = content.strip()
            if code.startswith("```python"):
                code = code[len("```python"):].strip()
            if code.startswith("```"):
                code = code[3:].strip()
            if code.endswith("```"):
                code = code[:-3].strip()
            
            return code
        
        except Exception as e:
            # Re-raise specific ValueErrors (like model missing) directly
            if isinstance(e, ValueError):
                raise e
            raise RuntimeError(f"Code generation failed: {e}")




class DataAnalysisNode(BaseNode):
    """
    Data Analysis Agent that generates analysis code and injects it into Jupyter notebook cells.
    """
    
    def __init__(self, id: str, provider: str = "ollama", model: str = "llama3.2", additional_imports: list = None, temperature: float = 0.7, max_tokens: int = 2000, **config):
        config["provider"] = provider
        config["model"] = model
        config["temperature"] = temperature
        config["max_tokens"] = max_tokens
        config["additional_imports"] = additional_imports or ["pandas", "numpy", "matplotlib", "seaborn"]
        config.pop('node_type', None)
        super().__init__(id=id, node_type=NodeType.DATA_ANALYSIS, **config)
        
        # Auto-install packages
        self._install_packages()
    
    def _install_packages(self):
        """Install required packages automatically."""
        import subprocess
        import sys
        
        packages = self.metadata.config.get("additional_imports", [])
        if not packages:
            return
        
        print(f"📦 Checking required packages: {', '.join(packages)}")
        
        for package in packages:
            # Extract package name (handle cases like "pandas as pd")
            pkg_name = package.split()[0] if ' ' in package else package
            
            try:
                __import__(pkg_name)
            except ImportError:
                print(f"Installing {pkg_name}...")
                try:
                    subprocess.check_call([sys.executable, "-m", "pip", "install", pkg_name, "-q"])
                    print(f"✅ Installed {pkg_name}")
                except subprocess.CalledProcessError:
                    print(f"⚠️  Failed to install {pkg_name}")
    
    def _define_inputs_outputs(self):
        self.metadata.inputs = [
            NodeInput(name="dataset_path", type="str", required=True, description="Path to dataset"),
            NodeInput(name="analysis_goal", type="str", required=True, description="Analysis objective"),
            NodeInput(name="provider", type="str", required=False, description="LLM provider"),
            NodeInput(name="model", type="str", required=False, description="LLM model")
        ]
        self.metadata.outputs = [
            NodeOutput(name="cells_created", type="int", description="Number of cells created"),
            NodeOutput(name="analysis_plan", type="str", description="Analysis plan text")
        ]
    
    def _get_dataset_info(self, dataset_path: str) -> Dict[str, Any]:
        """Extract dataset metadata for LLM context."""
        import pandas as pd
        import os
        
        if not os.path.exists(dataset_path):
            raise FileNotFoundError(f"Dataset not found: {dataset_path}")
        
        # Load dataset
        if dataset_path.endswith('.csv'):
            df = pd.read_csv(dataset_path, nrows=5)
        elif dataset_path.endswith(('.xls', '.xlsx')):
            df = pd.read_excel(dataset_path, nrows=5)
        else:
            raise ValueError(f"Unsupported file type: {dataset_path}")
        
        return {
            "columns": df.columns.tolist(),
            "dtypes": df.dtypes.astype(str).to_dict(),
            "shape": df.shape,
            "sample": df.head(3).to_dict(),
            "path": dataset_path
        }
    
    def _generate_analysis_plan(self, dataset_info: Dict[str, Any], analysis_goal: str, provider: str, model: str) -> Dict[str, Any]:
        """Use LLM to generate analysis plan and code."""
        
        # Get available imports and format them
        available_imports = self.metadata.config.get("additional_imports", [])
        imports_list = []
        for imp in available_imports:
            if ' as ' in imp:
                imports_list.append(f"import {imp}")
            else:
                imports_list.append(f"import {imp}")
        imports_info = "\n".join(imports_list) if imports_list else "No additional imports specified"
        
        prompt = f"""You are a data analysis expert. Generate a complete data analysis workflow.

Dataset Information:
- Path: {dataset_info['path']}
- Columns: {dataset_info['columns']}
- Data Types: {dataset_info['dtypes']}
- Shape: {dataset_info['shape']}

Analysis Goal: {analysis_goal}

Available Libraries (you can use these):
{imports_info}

Generate a step-by-step analysis plan with Python code for each step.
Return ONLY a JSON object with this structure:
{{
    "plan": "Brief description of the analysis approach",
    "steps": [
        {{
            "title": "Step title",
            "description": "What this step does",
            "code": "Python code for this step"
        }}
    ]
}}

Requirements:
- First step should load the dataset using: pd.read_excel(r'{dataset_info['path']}') or pd.read_csv(r'{dataset_info['path']}')
- Include data cleaning if needed
- Include exploratory data analysis (EDA)
- Include visualizations
- Keep each code block focused and executable
- Use clear variable names
- Add comments to explain the code
- You can use the libraries listed above
"""
        
        if provider == "ollama":
            url = self.metadata.config.get("base_url", "http://localhost:11434")
            temperature = self.metadata.config.get("temperature", 0.7)
            max_tokens = self.metadata.config.get("max_tokens", 2000)
            
            payload = {
                "model": model,
                "prompt": prompt,
                "stream": False,
                "format": "json",
                "options": {
                    "temperature": temperature,
                    "num_predict": max_tokens
                }
            }
            
            try:
                response = httpx.post(f"{url}/api/generate", json=payload, timeout=120.0)
                response.raise_for_status()
                result = response.json().get("response", "{}")
                return json.loads(result)
            except Exception as e:
                raise RuntimeError(f"Ollama code generation failed: {e}")
                
        elif provider == "openai":
            from openai import OpenAI
            import os
            api_key = self.metadata.config.get("api_key") or os.environ.get("OPENAI_API_KEY")
            temperature = self.metadata.config.get("temperature", 0.7)
            max_tokens = self.metadata.config.get("max_tokens", 2000)
            client = OpenAI(api_key=api_key)
            
            try:
                response = client.chat.completions.create(
                    model=model,
                    messages=[{"role": "user", "content": prompt}],
                    response_format={"type": "json_object"},
                    temperature=temperature,
                    max_tokens=max_tokens
                )
                return json.loads(response.choices[0].message.content)
            except Exception as e:
                raise RuntimeError(f"OpenAI code generation failed: {e}")
        else:
            raise ValueError(f"Provider {provider} not supported")
    
    def _execute_steps(self, steps: list) -> int:
        """Execute analysis steps and display output."""
        try:
            from IPython import get_ipython
            from IPython.display import display, Markdown, HTML
            
            ipython = get_ipython()
            if ipython is None:
                raise RuntimeError("Not running in Jupyter environment")
            
            # Execute each step
            for i, step in enumerate(steps, 1):
                # Display step header
                display(Markdown(f"## Step {i}: {step['title']}"))
                display(Markdown(f"*{step['description']}*"))
                display(HTML("<hr style='border: 1px solid #ddd; margin: 10px 0;'>"))
                
                # Display the code
                display(Markdown(f"```python\n{step['code']}\n```"))
                
                # Execute the code
                try:
                    result = ipython.run_cell(step['code'])
                    
                    # Check for errors
                    if result.error_in_exec:
                        display(HTML(f"<div style='color: red; padding: 10px; background: #fee; border-left: 4px solid red;'><b>Error:</b> {result.error_in_exec}</div>"))
                    elif result.success:
                        display(HTML("<div style='color: green; padding: 5px;'>✅ Executed successfully</div>"))
                    
                except Exception as e:
                    display(HTML(f"<div style='color: red; padding: 10px; background: #fee; border-left: 4px solid red;'><b>Execution Error:</b> {str(e)}</div>"))
                
                # Add spacing
                display(HTML("<br>"))
            
            return len(steps)
            
        except ImportError:
            raise RuntimeError("IPython not available. This node only works in Jupyter notebooks.")
    
    def execute(self, inputs: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        dataset_path = inputs.get("dataset_path", "")
        analysis_goal = inputs.get("analysis_goal", "")
        provider = inputs.get("provider") or self.metadata.config.get("provider", "ollama")
        model = inputs.get("model") or self.metadata.config.get("model", "llama3.2")
        
        if not dataset_path or not analysis_goal:
            raise ValueError("dataset_path and analysis_goal are required")
        
        # Get dataset info
        print(f"📊 Analyzing dataset: {dataset_path}")
        dataset_info = self._get_dataset_info(dataset_path)
        
        # Generate analysis plan
        print(f"🤖 Generating analysis plan using {provider}/{model}...")
        plan_data = self._generate_analysis_plan(dataset_info, analysis_goal, provider, model)
        
        # Execute steps and display output
        print(f"📝 Executing {len(plan_data['steps'])} analysis steps...\n")
        steps_executed = self._execute_steps(plan_data['steps'])
        
        print(f"\n✅ Executed {steps_executed} steps successfully!")
        
        return {
            "cells_created": steps_executed,
            "analysis_plan": plan_data.get("plan", "Analysis plan generated")
        }
