"""
Python Code Executor for Agent Tasks

Safely executes Python code generated by the agent, with:
- Sandboxed environment (restricted builtins)
- Pre-imported tool libraries
- Execution timeout
- Comprehensive error handling
- Result capture
- Tool execution tracking with descriptions
"""

import asyncio
import io
import sys
import time
import inspect
import uuid
import contextvars
from contextlib import redirect_stdout, redirect_stderr
from typing import Any, Dict, Optional, Tuple, Callable
import traceback
import json as _json
from pathlib import Path
from delfhos.errors import ToolExecutionError

from cortex._engine.utils.console import console
from cortex._engine.tools.tool_libraries import create_tool_libraries


# Initialize ContextVar to track parallel execution blocks
_parallel_group_id: contextvars.ContextVar[str] = contextvars.ContextVar('parallel_group_id', default=None)

# --- Safe JSON helpers for agent-generated code --------------------------------

_original_json_loads = _json.loads


def _safe_json_loads(obj, *args, **kwargs):
    """
    Safer json.loads for agent-generated code.
    
    Many generated snippets accidentally call json.loads() on objects that are
    already parsed (dict/list). Instead of raising TypeError, we just return
    the object as-is in those cases.
    """
    if isinstance(obj, (dict, list)):
        return obj
    return _original_json_loads(obj, *args, **kwargs)


def format_table(data, title=None, headers=None):
    """
    Format structured data as a Markdown table.
    
    Args:
        data: Can be:
            - List of dicts: [{"col1": "val1", "col2": "val2"}, ...]
            - List of lists: [["col1", "col2"], ["val1", "val2"], ...] (first row is headers)
            - Single dict: {"key1": "val1", "key2": "val2"}
            - List of values: ["val1", "val2", ...]
        title: Optional title to add above the table
        headers: Optional list of headers (if data is list of lists without headers)
    
    Returns:
        Markdown-formatted table string
    
    Examples:
        format_table([{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}])
        format_table([["Name", "Age"], ["Alice", 30], ["Bob", 25]])
        format_table([["Alice", 30], ["Bob", 25]], headers=["Name", "Age"])
        format_table({"Total": 100, "Count": 5})
    """
    if not data:
        return "*(No data)*"
    
    lines = []
    if title:
        lines.append(f"**{title}**\n")
        
    # Handle headers for list of lists
    if headers and isinstance(data, list) and len(data) > 0 and isinstance(data[0], (list, tuple)):
        # Prepend headers to data so existing logic handles it
        # Make sure we don't modify the original list if it's referenced elsewhere
        data = [headers] + list(data)
    
    # Handle list of dicts (most common case from SQL queries)
    if isinstance(data, list) and len(data) > 0 and isinstance(data[0], dict):
        # Get all unique keys from all dicts, deduplicating case variants
        # Keep only the first occurrence of each case-insensitive key (prefer original case)
        seen_keys_lower = {}
        all_keys = []
        for row in data:
            if isinstance(row, dict):
                for key in row.keys():
                    key_lower = str(key).lower()
                    if key_lower not in seen_keys_lower:
                        seen_keys_lower[key_lower] = key
                        all_keys.append(key)
        
        if not all_keys:
            return "*(Empty data)*"
        
        # Limit number of rows for large datasets (show first 100 rows)
        max_rows = 100
        display_data = data[:max_rows]
        total_rows = len(data)
        has_more = total_rows > max_rows
        
        # Header row
        headers = [str(key) for key in all_keys]
        lines.append("| " + " | ".join(headers) + " |")
        lines.append("| " + " | ".join(["---"] * len(headers)) + " |")
        
        # Data rows (limited)
        for row in display_data:
            if isinstance(row, dict):
                values = []
                for key in all_keys:
                    # Try to get value from the canonical key (first seen variant)
                    # But also check lowercase, capitalized, and uppercase variants
                    value = None
                    if key in row:
                        value = row[key]
                    else:
                        # Fallback: try case variants
                        key_lower = str(key).lower()
                        for variant_key in row.keys():
                            if str(variant_key).lower() == key_lower:
                                value = row[variant_key]
                                break
                    
                    # Convert value to string, handle None, escape pipes
                    if value is None:
                        value_str = ""
                    else:
                        value_str = str(value).replace("|", "\\|").replace("\n", " ")
                    values.append(value_str)
                lines.append("| " + " | ".join(values) + " |")
        
        # Add note if rows were truncated
        if has_more:
            lines.append(f"\n*Showing {max_rows} of {total_rows} rows. Use LIMIT in your query to see specific rows.*")
    
    # Handle list of lists (common from sheets - first row is headers)
    elif isinstance(data, list) and len(data) > 0 and isinstance(data[0], (list, tuple)):
        if len(data) == 0:
            return "*(Empty data)*"
        
        # Limit number of rows for large datasets (show first 100 rows)
        max_rows = 100
        total_rows = len(data) - 1  # Exclude header row
        display_data = data[:max_rows + 1]  # Include header + max_rows data rows
        has_more = total_rows > max_rows
        
        # First row is headers
        headers = [str(cell) for cell in data[0]]
        lines.append("| " + " | ".join(headers) + " |")
        lines.append("| " + " | ".join(["---"] * len(headers)) + " |")
        
        # Data rows (limited)
        for row in display_data[1:]:
            if isinstance(row, (list, tuple)):
                values = [str(cell).replace("|", "\\|").replace("\n", " ") if cell is not None else "" for cell in row]
                # Pad or truncate to match header length
                while len(values) < len(headers):
                    values.append("")
                values = values[:len(headers)]
                lines.append("| " + " | ".join(values) + " |")
        
        # Add note if rows were truncated
        if has_more:
            lines.append(f"\n*Showing {max_rows} of {total_rows} rows.*")
    
    # Handle single dict (key-value pairs)
    elif isinstance(data, dict):
        lines.append("| Key | Value |")
        lines.append("| --- | --- |")
        for key, value in data.items():
            key_str = str(key).replace("|", "\\|")
            value_str = str(value).replace("|", "\\|").replace("\n", " ") if value is not None else ""
            lines.append(f"| {key_str} | {value_str} |")
    
    # Handle list of simple values
    elif isinstance(data, list):
        lines.append("| Value |")
        lines.append("| --- |")
        for item in data:
            value_str = str(item).replace("|", "\\|").replace("\n", " ") if item is not None else ""
            lines.append(f"| {value_str} |")
    
    else:
        # Fallback: just convert to string
        return str(data)
    
    return "\n".join(lines)


class PythonExecutionError(ToolExecutionError):
    """Raised when Python code execution fails"""

    def __init__(self, message: str):
        super().__init__(tool_name="python_executor", detail=message)


class ToolExecutionTracker:
    """
    Tracks tool execution with descriptions for real-time frontend updates.
    
    Each tool call gets a unique entry with:
    - tool_name: The tool being used (sql, gmail, sheets, etc.)
    - description: User-provided description of what the tool is doing
    - start_time: When the tool call started
    - duration: How long it took
    - status: success/failure
    
    Supports parallel execution: Multiple concurrent calls to the same tool_name
    are tracked independently using unique call IDs (FIFO matching).
    """
    
    def __init__(self, orchestrator, task_id: str):
        self.orchestrator = orchestrator
        self.task_id = task_id
        self.active_calls = {}  # call_id -> {tool_name, start_time, description, ui_metadata}
        self.call_counter = 0  # Unique ID for each call (supports parallel execution)
    
    async def track_start(self, tool_name: str, description: str, ui_metadata: Optional[Dict[str, Any]] = None, model: str = None):
        """Track the start of a tool call and emit starting entry to frontend immediately
        
        Returns:
            call_id: Unique identifier for this call (for matching with track_end in parallel scenarios)
        """
        start_time = time.time()
        
        # Generate unique call ID for this execution (supports parallel calls to same tool)
        call_id = f"tool_{int(time.time()*1000)}_{len(self.active_calls)}"
        
        # Check if we are executing inside a parallel asyncio.gather block
        parallel_group = _parallel_group_id.get()
        if parallel_group:
            if not ui_metadata:
                ui_metadata = {}
            ui_metadata['parallel_group_id'] = parallel_group
            
        self.call_counter += 1
        
        self.active_calls[call_id] = {
            "tool_name": tool_name,
            "start_time": start_time,
            "description": description,
            "ui_metadata": ui_metadata or None,
            "model": model,
        }
        
        # Immediately emit starting entry to frontend for real-time updates
        if self.orchestrator:
            # Use async version to properly await the callback
            await self.orchestrator.track_tool_timing_async(
                self.task_id,
                tool_name,
                duration=None,  # None indicates ongoing/starting
                model=model,
                description=description,
                is_starting=True,
                metadata=ui_metadata,
            )
        
        return call_id
    
    async def track_end(self, tool_name: str, duration: float, success: bool = True, error: str = None, model: str = None, description: str = None, ui_metadata: Optional[Dict[str, Any]] = None):
        """Track the end of a tool call and emit timing event
        
        For parallel execution: Matches the oldest active call for this tool_name (FIFO)
        """
        # Find the oldest active call for this tool_name (FIFO matching)
        call_id = None
        call_info = None
        for cid in sorted(self.active_calls.keys()):  # Sort to maintain order
            info = self.active_calls[cid]
            if info.get("tool_name") == tool_name:
                call_id = cid
                call_info = info
                break
        
        # Remove the matched call
        if call_id:
            self.active_calls.pop(call_id, None)
        
        # Use provided description, or from call_info, or default
        final_description = description or (call_info.get("description") if call_info else None) or f"Using {tool_name}"
        
        # Merge the metadata tracking parallel groups with the metadata from the tool end
        final_metadata = {}
        if call_info and isinstance(call_info.get("ui_metadata"), dict):
            final_metadata.update(call_info.get("ui_metadata"))
        if ui_metadata and isinstance(ui_metadata, dict):
            final_metadata.update(ui_metadata)
            
        if not final_metadata:
            final_metadata = None
        
        if not call_info:
            # No matching start call, create a minimal entry
            call_info = {
                "description": final_description,
                "start_time": time.time() - duration,
                "ui_metadata": final_metadata,
            }
        
        # Track in orchestrator (this will trigger the callback to frontend)
        # Use async version to properly await the callback during Python execution
        if self.orchestrator:
            # This will find and update the starting entry with duration
            await self.orchestrator.track_tool_timing_async(
                self.task_id,
                tool_name,
                duration,
                model=model,
                description=final_description,
                is_starting=False,
                metadata=final_metadata,
            )


class PythonExecutor:
    """
    Executes agent-generated Python code in a controlled environment.
    """
    
    def __init__(self, tool_manager, task_id: str, agent_id: str, light_llm: str, heavy_llm: str, orchestrator=None,
                 vision_model: Optional[str] = None, search_model: Optional[str] = None):
        self.tool_manager = tool_manager
        self.task_id = task_id
        self.agent_id = agent_id
        self.light_llm = light_llm
        self.heavy_llm = heavy_llm
        self.vision_model = vision_model or self.heavy_llm
        self.search_model = search_model or self.light_llm
        self.orchestrator = orchestrator
        self.execution_timeout = 300  # 5 minutes max per task
        self.namespace = None
    
    async def execute(self, code: str) -> Dict[str, Any]:
        """
        Execute Python code and return results.
        
        Args:
            code: Python code string to execute
        
        Returns:
            Dict with:
                - success: bool
                - result: Any (final result/return value)
                - output: str (printed output)
                - error: str (error message if failed)
                - execution_time: float
        """
        start_time = time.time()
        
        # Initialize persistent namespace if not already created
        if self.namespace is None:
            # Create tool execution tracker
            tool_tracker = ToolExecutionTracker(self.orchestrator, self.task_id)
            
            # Create tool libraries with tracker
            libraries = create_tool_libraries(
                self.tool_manager,
                self.task_id,
                self.agent_id,
                self.light_llm,
                self.heavy_llm,
                tool_tracker=tool_tracker,
                vision_llm=self.vision_model,
                search_llm=self.search_model,
                memory=self.orchestrator.memory if self.orchestrator else None,
            )
            
            # Build safe execution namespace
            self.namespace = self._build_namespace(libraries)
        
        # Capture stdout/stderr
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        
        try:
            lib_keys = list(libraries.keys()) if 'libraries' in locals() else list(self.namespace.keys())
            console.debug(
                "Executing Python code",
                f"Code length: {len(code)} chars, Libraries: {lib_keys}",
                task_id=self.task_id,
                agent_id=self.agent_id
            )
            
            # Execute with timeout
            with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
                result = await asyncio.wait_for(
                    self._execute_async(code, self.namespace),
                    timeout=self.execution_timeout
                )
            
            execution_time = time.time() - start_time
            output = stdout_capture.getvalue()
            
            # Determine final result
            # Priority: explicit return > last expression > None (let orchestrator use output)
            # If code only uses print(), result will be None and orchestrator will use output
            # If code has explicit return, use that
            final_result = result if result is not None else None
            
            console.debug(
                "Python execution completed",
                f"Time: {execution_time:.2f}s, Output: {len(output)} chars",
                task_id=self.task_id,
                agent_id=self.agent_id
            )
            
            return {
                "success": True,
                "result": final_result,
                "output": output,
                "error": None,
                "execution_time": execution_time
            }
        
        except asyncio.TimeoutError:
            execution_time = time.time() - start_time
            error_msg = f"Execution timeout after {self.execution_timeout}s"
            console.error(
                "Python execution timeout",
                error_msg,
                task_id=self.task_id,
                agent_id=self.agent_id
            )
            return {
                "success": False,
                "result": None,
                "output": stdout_capture.getvalue(),
                "error": error_msg,
                "execution_time": execution_time
            }
        
        except Exception as e:
            execution_time = time.time() - start_time
            error_msg = self._format_error(e, stderr_capture.getvalue())
            console.error(
                "Python execution error",
                error_msg,
                task_id=self.task_id,
                agent_id=self.agent_id
            )
            return {
                "success": False,
                "result": None,
                "output": stdout_capture.getvalue(),
                "error": error_msg,
                "execution_time": execution_time
            }
    
    async def _execute_async(self, code: str, namespace: dict) -> Any:
        """
        Execute Python code asynchronously.
        
        Always wraps code in an async function for consistency and to handle
        both sync and async code uniformly.
        """
        # Always wrap in async function - this allows us to handle both sync and async code
        # Ensure code ends with a newline for proper indentation
        code_to_wrap = code.rstrip() + "\n"

        # Heuristic: if the LLM defined an async entrypoint but never called it,
        # automatically append an await so tools inside it actually run.
        if "async def main" in code_to_wrap and "await main(" not in code_to_wrap:
            code_to_wrap += "\nawait main()\n"
        elif "async def run" in code_to_wrap and "await run(" not in code_to_wrap:
            code_to_wrap += "\nawait run()\n"

        # Add a helper to export locals to the global namespace
        def __export_locals__(l):
            namespace.update({k: v for k, v in l.items() if not k.startswith('_')})
        namespace['__export_locals__'] = __export_locals__

        # Inject export at the end for successful runs
        code_to_wrap += "\n__export_locals__(locals())\n"

        wrapped_code = f"""async def __agent_task__():\n{self._indent_code(code_to_wrap)}"""
        
        # Compile and execute the async function definition
        try:
            compiled = compile(wrapped_code, "<string>", "exec")
            exec(compiled, namespace)
        except SyntaxError as e:
            # Log the wrapped code for debugging
            console.error(
                "Python code compilation failed",
                f"Syntax error at line {e.lineno}: {str(e)}\n\nWrapped code (first 30 lines):\n{chr(10).join(wrapped_code.split(chr(10))[:30])}",
                task_id=self.task_id,
                agent_id=self.agent_id
            )
            raise
        
        # Get the async function and await it
        async_func = namespace.get("__agent_task__")
        if async_func:
            try:
                result = await async_func()
                return result
            except Exception as e:
                # Execution crashed! Walk traceback to extract locals from the deepest user frame.
                # Since the LLM often creates inner functions (like `async def main():`), 
                # we want the deepest frame whose filename is "<string>".
                tb = e.__traceback__
                target_frame = None
                while tb:
                    if tb.tb_frame.f_code.co_filename == "<string>":
                        target_frame = tb.tb_frame
                    tb = tb.tb_next
                
                if target_frame:
                    namespace.update({k: v for k, v in target_frame.f_locals.items() if not k.startswith('_')})
                
                if isinstance(e, NameError):
                    # Log the full code when NameError occurs during execution
                    console.error(
                        "Python code NameError during execution",
                        f"NameError: {str(e)}\n\nFull generated code:\n{code}",
                        task_id=self.task_id,
                        agent_id=self.agent_id
                    )
                raise
        else:
            raise ToolExecutionError(tool_name="python_executor", detail="Failed to create async task function")

    
    def _build_namespace(self, libraries: dict) -> dict:
        """
        Build the execution namespace with safe builtins and tool libraries.
        """
        # Force matplotlib to use non-interactive Agg backend (file-only, no GUI).
        # This prevents NSWindow crashes when running from async background threads.
        try:
            import matplotlib
            matplotlib.use('Agg')
        except ImportError:
            pass

        # Monkey-patch json.loads with a safer version for the execution namespace
        # This prevents common LLM mistakes like json.loads(dict) from crashing tasks.
        try:
            _json.loads = _safe_json_loads  # type: ignore[assignment]
        except Exception:
            # If monkey-patch fails for any reason, continue with original behavior
            pass
        
        # Safe builtins - exclude dangerous or interactive functions like input()
        def _blocked_input(*args, **kwargs):
            raise ToolExecutionError(tool_name="python_executor", detail="input() is not available in this environment. Use approval.ask(...) for user consent instead.")

        # A safe print function that always flushes
        def _safe_print(*args, **kwargs):
            print(*args, **kwargs, flush=True)

        # Original modules to be proxied or used directly
        _asyncio = __import__("asyncio")
        _re = __import__("re")
        _datetime = __import__("datetime")
        _time = __import__("time")
        _math = __import__("math")
        _statistics = __import__("statistics")

        # Special interceptor for asyncio.gather to track parallel execution in the UI
        def _asyncio_getattr(name):
            if name == 'gather':
                async def gather_wrapper(*coros_or_futures, return_exceptions=False):
                    # Generate a unique ID for this parallel block
                    group_id = f"group_{uuid.uuid4().hex[:8]}"
                    
                    # We need to wrap each coroutine so it runs with the contextvar set
                    # asyncio.gather doesn't automatically propagate contextvars in all Python versions
                    async def coro_wrapper(coro):
                        token = _parallel_group_id.set(group_id)
                        try:
                            # Await the actual tool call/coroutine
                            return await coro
                        finally:
                            _parallel_group_id.reset(token)
                    
                    wrapped_coros = [coro_wrapper(c) for c in coros_or_futures]
                    return await _asyncio.gather(*wrapped_coros, return_exceptions=return_exceptions)
                
                return gather_wrapper
            return getattr(_asyncio, name)

        class AsyncioProxy:
            def __getattr__(self, name):
                return _asyncio_getattr(name)
            
        asyncio_proxy = AsyncioProxy()

        allowed_import_roots = {
            "asyncio",
            "datetime",
            "json",
            "math",
            "pathlib",
            "re",
            "statistics",
            "time",
        }

        def _safe_import(name, *args, **kwargs):
            root_name = name.split(".", 1)[0]
            if root_name not in allowed_import_roots:
                raise ToolExecutionError(tool_name="python_executor", detail=f"Import of module '{name}' is not allowed in this environment")
            if root_name == "asyncio":
                return asyncio_proxy
            return __import__(name, *args, **kwargs)
        
        # Build namespace
        namespace = {
            # Standard safe builtins
            "print": _safe_print,
            "int": int,
            "float": float,
            "bool": bool,
            "str": str,
            "list": list,
            "dict": dict,
            "tuple": tuple,
            "set": set,
            "len": len,
            "sum": sum,
            "min": min,
            "max": max,
            "abs": abs,
            "round": round,
            "any": any,
            "all": all,
            "enumerate": enumerate,
            "zip": zip,
            "isinstance": isinstance,
            "issubclass": issubclass,
            "type": type,
            "range": range,
            "json": _json, # Use the monkey-patched _json
            "asyncio": asyncio_proxy,
            "datetime": _datetime,
            "time": _time,
            "re": _re,
            "math": _math,
            "statistics": _statistics,
            "__import__": _safe_import,
            "Path": Path,
            "Exception": Exception,
            "ValueError": ValueError,
            "KeyError": KeyError,
            "TypeError": TypeError,
            "AttributeError": AttributeError,
            "IndexError": IndexError,
            "None": None,
            "True": True,
            "False": False,
            "locals": locals,
            "input": _blocked_input,
        }
        
        # Standard library modules (safe subset)
        safe_modules = {
            "json": _json,
            "re": __import__("re"),
            "datetime": __import__("datetime"),
            "time": __import__("time"),
            "math": __import__("math"),
            "statistics": __import__("statistics"),
            "asyncio": asyncio_proxy,
        }
        
        # Add safe modules to the namespace
        namespace.update(safe_modules)
        
        # Build namespace
        namespace.update({
            "__builtins__": namespace,
            "ctx": self.orchestrator.agent_context if self.orchestrator else {},
            **libraries,  # Tool libraries (sql, sheets, gmail, etc.)
            "task_id": self.task_id,  # Make task_id available for use in code (though files.save() handles paths automatically)
            "format_table": format_table,  # Helper function to format data as Markdown tables
            "__name__": "__agent_execution__",  # Define __name__ to prevent NameErrors and skip if __name__ == "__main__" blocks
            "__file__": "agent_script.py",  # Define __file__ for completeness
        })
        
        return namespace
    
    def _indent_code(self, code: str, spaces: int = 4) -> str:
        """Indent code for wrapping in a function"""
        indent = " " * spaces
        return "\n".join(indent + line if line.strip() else line for line in code.split("\n"))
    
    def _format_error(self, exception: Exception, stderr: str) -> str:
        """
        Format exception into user-friendly error message.
        """
        error_type = type(exception).__name__
        error_msg = str(exception)
        
        # Extract relevant traceback
        tb_lines = traceback.format_exception(type(exception), exception, exception.__traceback__)
        
        # Filter to show only user code (not internal executor code)
        filtered_tb = []
        for line in tb_lines:
            if "python_executor.py" not in line and "__agent_task__" not in line:
                filtered_tb.append(line)
        
        # Combine into readable error
        if filtered_tb:
            formatted_error = "".join(filtered_tb)
        else:
            formatted_error = f"{error_type}: {error_msg}"
        
        if stderr:
            formatted_error += f"\n\nStderr:\n{stderr}"
        
        return formatted_error.strip()


def parse_python_code(llm_response: str) -> str:
    """
    Extract Python code from LLM response.
    
    The LLM might return code wrapped in markdown code blocks or with explanations.
    This function extracts just the executable code.
    
    Handles thinking models that mix reasoning text with code output.
    
    If the response appears to be plain text (not code), wraps it in a print statement.
    """
    response = llm_response.strip()
    
    def _finalize_code(c: str) -> str:
        import re
        c = _clean_markdown_markers(c)
        # Prevent LLM from crashing the running event loop by replacing asyncio.run(func()) with await func()
        # Handles both asyncio.run(main()) and asyncio.run(main(args...))
        c = re.sub(r'asyncio\.run\s*\(\s*(.*?)\s*\)', r'await \1', c)
        return c

    # Check for markdown code blocks (most reliable extraction)
    if "```python" in response:
        # Extract the LAST ```python ... ``` block (thinking text might have earlier fragments)
        last_start = response.rfind("```python")
        start = last_start + 9
        end = response.find("```", start)
        if end != -1:
            code = response[start:end].strip()
            if _looks_like_python_code(code):
                return _finalize_code(code)
    
    elif "```" in response:
        # Generic code block - use the last one
        last_start = response.rfind("```")
        # Find the opening ``` before this closing one
        # Search backwards for the opening block
        search_area = response[:last_start]
        block_start = search_area.rfind("```")
        if block_start != -1:
            start = block_start + 3
            # Skip optional language identifier on the same line
            newline_pos = response.find("\n", start)
            if newline_pos != -1 and newline_pos < last_start:
                start = newline_pos + 1
            code = response[start:last_start].strip()
            if _looks_like_python_code(code):
                return _finalize_code(code)
    
    # Strip thinking/reasoning text: keep only lines that look like Python code
    # This handles thinking models that concatenate reasoning with code
    lines = response.split('\n')
    code_lines = []
    in_code = False
    
    # Python code indicators for line-level detection
    _line_indicators = (
        'import ', 'from ', 'await ', 'print(', 'async def', 'def ',
        'if ', 'for ', 'while ', 'try:', 'except', 'return ', 'class ',
        'elif ', 'else:', 'finally:', 'with ', 'raise ', 'pass',
        'break', 'continue', 'yield ', 'lambda ', 'global ', 'nonlocal ',
    )
    
    for line in lines:
        stripped = line.strip()
        # Once we start seeing code, keep everything (including blank lines, comments)
        if in_code:
            code_lines.append(line)
        elif stripped and (
            any(stripped.startswith(ind) or stripped.startswith(ind.lstrip()) for ind in _line_indicators)
            or (len(stripped) > 2 and '=' in stripped and not stripped.startswith('#') and not stripped.startswith('-') and not stripped.startswith('*'))
        ):
            in_code = True
            code_lines.append(line)
    
    if code_lines:
        code = '\n'.join(code_lines).strip()
        if _looks_like_python_code(code):
            return _finalize_code(code)
    
    # Fallback: check if raw response looks like Python code
    if _looks_like_python_code(response):
        return _finalize_code(response)
    
    # If it doesn't look like code, wrap it in print statements
    lines = response.split('\n')
    python_lines = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        escaped = line.replace('"', '\\"').replace("'", "\\'")
        python_lines.append(f'print("{escaped}")')
    
    return '\n'.join(python_lines)


def _clean_markdown_markers(code: str) -> str:
    """
    Remove markdown code block markers that might be present inside the code.
    
    Sometimes the LLM includes markdown markers (```python, ```) as part of the code content.
    This function strips them out.
    """
    lines = code.split('\n')
    cleaned_lines = []
    
    for line in lines:
        stripped = line.strip()
        # Skip lines that are just markdown markers
        if stripped in ['```python', '```', '```py']:
            continue
        # Remove trailing markdown markers from lines
        if stripped.endswith('```'):
            line = line.rstrip('`')
        cleaned_lines.append(line)
    
    return '\n'.join(cleaned_lines)


def _looks_like_python_code(text: str) -> bool:
    """
    Heuristic to check if text looks like Python code.
    
    Returns True if the text contains common Python patterns.
    """
    if not text:
        return False
    
    # Check for Python keywords/syntax
    python_indicators = [
        'import ', 'from ', 'await ', 'print(', 'async def', 'def ',
        'if ', 'for ', 'while ', 'try:', 'except', 'return ',
        'await sql.', 'await gmail.', 'await sheets.', 'await files.',
        'await llm.', 'await drive.', 'await calendar.',
        '=', '==', '!=', '+=', '-=', '*=', '/=',
        '(', ')', '[', ']', '{', '}'
    ]
    
    text_lower = text.lower()
    indicator_count = sum(1 for indicator in python_indicators if indicator.lower() in text_lower)
    
    # If we find multiple Python indicators, it's likely code
    if indicator_count >= 2:
        return True
    
    # If it starts with a Python keyword or import, it's code
    first_words = text.strip().split()[:3]
    if any(word in ['import', 'from', 'print', 'await', 'async', 'def', 'if', 'try', 'for', 'while'] 
           for word in first_words):
        return True
    
    # If it contains function calls with parentheses, it's likely code
    if '(' in text and ')' in text:
        return True
    
    # If it's just plain text (no code indicators), it's not code
    return False


async def execute_python_task(
    code: str,
    tool_manager,
    task_id: str,
    agent_id: str,
    light_llm: str,
    heavy_llm: str,
    orchestrator=None
) -> Dict[str, Any]:
    """
    Convenience function to execute Python code for a task.
    
    Args:
        code: Python code string
        tool_manager: Tool manager instance
        task_id: Task ID for logging
        agent_id: Agent ID for logging
        light_llm: Light LLM model name
        heavy_llm: Heavy LLM model name
        orchestrator: Optional orchestrator instance for tool tracking
    
    Returns:
        Execution result dict
    """
    executor = PythonExecutor(tool_manager, task_id, agent_id, light_llm, heavy_llm, orchestrator)
    return await executor.execute(code)

