"""Code-use agent service - Jupyter notebook-like code execution for browser automation."""

import asyncio
import datetime
import html
import json
import logging
import re
import traceback
from pathlib import Path
from typing import Any

from uuid_extensions import uuid7str

from noesium.subagents.bu.adapters.llm_adapter import (
    AssistantMessage,
    BaseChatModel,
    BaseMessage,
    ContentPartImageParam,
    ContentPartTextParam,
    ImageURL,
    UserMessage,
)
from noesium.subagents.bu.browser import BrowserSession
from noesium.subagents.bu.browser.profile import BrowserProfile
from noesium.subagents.bu.dom.service import DomService
from noesium.subagents.bu.filesystem.file_system import FileSystem
from noesium.subagents.bu.screenshots.service import ScreenshotService
from noesium.subagents.bu.tokens.service import TokenCost
from noesium.subagents.bu.tokens.views import UsageSummary
from noesium.subagents.bu.tools.service import CodeAgentTools, Tools
from noesium.subagents.bu.utils import get_browser_use_version

from .formatting import format_browser_state_for_llm
from .namespace import EvaluateError, create_namespace
from .utils import detect_token_limit_issue, extract_code_blocks, extract_url_from_task, truncate_message_content
from .views import (
    CellType,
    CodeAgentHistory,
    CodeAgentHistoryList,
    CodeAgentModelOutput,
    CodeAgentResult,
    CodeAgentState,
    CodeAgentStepMetadata,
    ExecutionStatus,
    NotebookSession,
)

logger = logging.getLogger(__name__)


class CodeAgent:
    """
    Agent that executes Python code in a notebook-like environment for browser automation.

    This agent provides a Jupyter notebook-like interface where the LLM writes Python code
    that gets executed in a persistent namespace with browser control functions available.
    """

    def __init__(
        self,
        task: str,
        # Optional parameters
        llm: BaseChatModel | None = None,
        browser_session: BrowserSession | None = None,
        browser: BrowserSession | None = None,  # Alias for browser_session
        tools: Tools | None = None,
        controller: Tools | None = None,  # Alias for tools
        # Agent settings
        page_extraction_llm: BaseChatModel | None = None,
        file_system: FileSystem | None = None,
        available_file_paths: list[str] | None = None,
        sensitive_data: dict[str, str | dict[str, str]] | None = None,
        max_steps: int = 100,
        max_failures: int = 8,
        max_validations: int = 0,
        use_vision: bool = True,
        calculate_cost: bool = False,
        demo_mode: bool | None = None,
        **kwargs,
    ):
        """
        Initialize the code-use agent.

        Args:
                task: The task description for the agent
                browser_session: Optional browser session (will be created if not provided) [DEPRECATED: use browser]
                browser: Optional browser session (cleaner API)
                tools: Optional Tools instance (will create default if not provided)
                controller: Optional Tools instance
                page_extraction_llm: Optional LLM for page extraction
                file_system: Optional file system for file operations
                available_file_paths: Optional list of available file paths
                sensitive_data: Optional sensitive data dictionary
                max_steps: Maximum number of execution steps
                max_failures: Maximum consecutive errors before termination (default: 8)
                max_validations: Maximum number of times to run the validator agent (default: 0)
                use_vision: Whether to include screenshots in LLM messages (default: True)
                calculate_cost: Whether to calculate token costs (default: False)
                demo_mode: Enable the in-browser demo panel for live logging (default: False)
                llm: Optional ChatBrowserUse LLM instance (will create default if not provided)
                **kwargs: Additional keyword arguments for compatibility (ignored)
        """
        # Log and ignore unknown kwargs for compatibility
        if kwargs:
            logger.debug(f"Ignoring additional kwargs for CodeAgent compatibility: {list(kwargs.keys())}")

        if llm is None:
            # For noesium, we require an LLM to be provided
            raise RuntimeError("CodeAgent requires an LLM to be provided. Please pass an llm parameter.")

        # No longer restrict to ChatBrowserUse - work with any LLM adapter

        # Handle browser vs browser_session parameter (browser takes precedence)
        if browser and browser_session:
            raise ValueError(
                'Cannot specify both "browser" and "browser_session" parameters. Use "browser" for the cleaner API.'
            )
        browser_session = browser or browser_session

        # Handle controller vs tools parameter (controller takes precedence)
        if controller and tools:
            raise ValueError(
                'Cannot specify both "controller" and "tools" parameters. Use "controller" for the cleaner API.'
            )
        tools = controller or tools

        # Store browser_profile for creating browser session if needed
        self._demo_mode_enabled = False
        if browser_session is None:
            profile_kwargs: dict[str, Any] = {}
            if demo_mode is not None:
                profile_kwargs["demo_mode"] = demo_mode
            self._browser_profile_for_init = BrowserProfile(**profile_kwargs)
        else:
            self._browser_profile_for_init = None

        self.task = task
        self.llm = llm
        self.browser_session = browser_session
        if self.browser_session:
            if demo_mode is not None and self.browser_session.browser_profile.demo_mode != demo_mode:
                self.browser_session.browser_profile = self.browser_session.browser_profile.model_copy(
                    update={"demo_mode": demo_mode}
                )
            self._demo_mode_enabled = bool(self.browser_session.browser_profile.demo_mode)
        self.tools = tools or CodeAgentTools()
        self.page_extraction_llm = page_extraction_llm
        self.file_system = file_system if file_system is not None else FileSystem(base_dir="./")
        self.available_file_paths = available_file_paths or []
        self.sensitive_data = sensitive_data
        self.max_steps = max_steps
        self.max_failures = max_failures
        self.max_validations = max_validations
        self.use_vision = use_vision

        self.session = NotebookSession()
        self.namespace: dict[str, Any] = {}
        self._llm_messages: list[BaseMessage] = []  # Internal LLM conversation history
        self.complete_history: list[CodeAgentHistory] = []  # Type-safe history with model_output and result
        self.dom_service: DomService | None = None
        self._last_browser_state_text: str | None = None  # Track last browser state text
        self._last_screenshot: str | None = None  # Track last screenshot (base64)
        self._consecutive_errors = 0  # Track consecutive errors for auto-termination
        self._validation_count = 0  # Track number of validator runs
        self._last_llm_usage: Any | None = None  # Track last LLM call usage stats
        self._step_start_time = 0.0  # Track step start time for duration calculation
        self.usage_summary: UsageSummary | None = None  # Track usage summary across run for history property
        self._sample_output_added = False  # Track whether preview cell already created

        # Initialize screenshot service for eval tracking
        self.id = uuid7str()
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        base_tmp = Path("/tmp")
        self.agent_directory = base_tmp / f"browser_use_code_agent_{self.id}_{timestamp}"
        self.screenshot_service = ScreenshotService(agent_directory=self.agent_directory)

        # Initialize token cost service for usage tracking
        self.token_cost_service = TokenCost(include_cost=calculate_cost)
        self.token_cost_service.register_llm(llm)
        if page_extraction_llm:
            self.token_cost_service.register_llm(page_extraction_llm)

        # Set version and source for telemetry
        self.version = get_browser_use_version()
        try:
            package_root = Path(__file__).parent.parent.parent
            repo_files = [".git", "README.md", "docs", "examples"]
            if all(Path(package_root / file).exists() for file in repo_files):
                self.source = "git"
            else:
                self.source = "pip"
        except Exception:
            self.source = "unknown"

    async def run(self, max_steps: int | None = None) -> NotebookSession:
        """
        Run the agent to complete the task.

        Args:
                max_steps: Optional override for maximum number of steps (uses __init__ value if not provided)

        Returns:
                The notebook session with all executed cells
        """
        # Use override if provided, otherwise use value from __init__
        steps_to_run = max_steps if max_steps is not None else self.max_steps
        self.max_steps = steps_to_run
        # Start browser if not provided
        if self.browser_session is None:
            assert self._browser_profile_for_init is not None
            self.browser_session = BrowserSession(browser_profile=self._browser_profile_for_init)
            await self.browser_session.start()

        if self.browser_session:
            self._demo_mode_enabled = bool(self.browser_session.browser_profile.demo_mode)
            if self._demo_mode_enabled and getattr(self.browser_session.browser_profile, "headless", False):
                logger.warning(
                    "Demo mode is enabled but the browser is headless=True; set headless=False to view the panel."
                )
            if self._demo_mode_enabled:
                await self._demo_mode_log(f"Started CodeAgent task: {self.task}", "info", {"tag": "task"})

        # Initialize DOM service with cross-origin iframe support enabled
        self.dom_service = DomService(
            browser_session=self.browser_session,
            cross_origin_iframes=True,  # Enable for code-use agent to access forms in iframes
        )

        # Create namespace with all tools
        self.namespace = create_namespace(
            browser_session=self.browser_session,
            tools=self.tools,
            page_extraction_llm=self.page_extraction_llm,
            file_system=self.file_system,
            available_file_paths=self.available_file_paths,
            sensitive_data=self.sensitive_data,
        )

        # Initialize conversation with task
        self._llm_messages.append(UserMessage(content=f"Task: {self.task}"))

        # Track agent run error for telemetry
        agent_run_error: str | None = None
        should_delay_close = False

        # Extract URL from task and navigate if found
        initial_url = extract_url_from_task(self.task)
        if initial_url:
            try:
                logger.info(f"Extracted URL from task, navigating to: {initial_url}")
                # Use the navigate action from namespace
                await self.namespace["navigate"](initial_url)
                # Wait for page load
                await asyncio.sleep(2)

                # Record this navigation as a cell in the notebook
                nav_code = f"await navigate('{initial_url}')"
                cell = self.session.add_cell(source=nav_code)
                cell.status = ExecutionStatus.SUCCESS
                cell.execution_count = self.session.increment_execution_count()
                cell.output = f"Navigated to {initial_url}"

                # Get browser state after navigation for the cell
                if self.dom_service:
                    try:
                        browser_state_text, _ = await self._get_browser_state()
                        cell.browser_state = browser_state_text
                    except Exception as state_error:
                        logger.debug(f"Failed to capture browser state for initial navigation cell: {state_error}")

            except Exception as e:
                logger.warning(f"Failed to navigate to extracted URL {initial_url}: {e}")
                # Record failed navigation as error cell
                nav_code = f"await navigate('{initial_url}')"
                cell = self.session.add_cell(source=nav_code)
                cell.status = ExecutionStatus.ERROR
                cell.execution_count = self.session.increment_execution_count()
                cell.error = str(e)

        # Get initial browser state before first LLM call
        if self.browser_session and self.dom_service:
            try:
                browser_state_text, screenshot = await self._get_browser_state()
                self._last_browser_state_text = browser_state_text
                self._last_screenshot = screenshot
            except Exception as e:
                logger.warning(f"Failed to get initial browser state: {e}")

        # Main execution loop
        for step in range(self.max_steps):
            logger.info(f"\n\n\n\n\n\n\nStep {step + 1}/{self.max_steps}")
            await self._demo_mode_log(f"Starting step {step + 1}/{self.max_steps}", "info", {"step": step + 1})

            # Start timing this step
            self._step_start_time = datetime.datetime.now().timestamp()

            # Check if we're approaching the step limit or error limit and inject warning
            steps_remaining = self.max_steps - step - 1
            errors_remaining = self.max_failures - self._consecutive_errors

            should_warn = (
                steps_remaining <= 1  # Last step or next to last
                or errors_remaining <= 1  # One more error will terminate
                or (steps_remaining <= 2 and self._consecutive_errors >= 2)  # Close to both limits
            )

            if should_warn:
                warning_message = (
                    f"\n\n⚠️ CRITICAL WARNING: You are approaching execution limits!\n"
                    f"- Steps remaining: {steps_remaining + 1}\n"
                    f"- Consecutive errors: {self._consecutive_errors}/{self.max_failures}\n\n"
                    f"YOU MUST call done() in your NEXT response, even if the task is incomplete:\n"
                    f"- Set success=False if you couldn't complete the task\n"
                    f"- Return EVERYTHING you found so far (partial data is better than nothing)\n"
                    f"- Include any variables you've stored (products, all_data, etc.)\n"
                    f"- Explain what worked and what didn't\n\n"
                    f"Without done(), the user will receive NOTHING."
                )
                self._llm_messages.append(UserMessage(content=warning_message))

            try:
                # Fetch fresh browser state right before LLM call (only if not already set)
                if not self._last_browser_state_text and self.browser_session and self.dom_service:
                    try:
                        logger.debug("🔍 Fetching browser state before LLM call...")
                        browser_state_text, screenshot = await self._get_browser_state()
                        self._last_browser_state_text = browser_state_text
                        self._last_screenshot = screenshot

                        # # Log browser state
                        # if len(browser_state_text) > 2000:
                        # 	logger.info(
                        # 		f'Browser state (before LLM):\n{browser_state_text[:2000]}...\n[Truncated, full state {len(browser_state_text)} chars sent to LLM]'
                        # 	)
                        # else:
                        # 	logger.info(f'Browser state (before LLM):\n{browser_state_text}')
                    except Exception as e:
                        logger.warning(f"Failed to get browser state before LLM call: {e}")

                # Get code from LLM (this also adds to self._llm_messages)
                try:
                    code, full_llm_response = await self._get_code_from_llm(step_number=step + 1)
                except Exception as llm_error:
                    # LLM call failed - count as consecutive error and retry
                    self._consecutive_errors += 1
                    logger.warning(
                        f"LLM call failed (consecutive errors: {self._consecutive_errors}/{self.max_failures}), retrying: {llm_error}"
                    )
                    await self._demo_mode_log(
                        f"LLM call failed: {llm_error}",
                        "error",
                        {"step": step + 1},
                    )

                    # Check if we've hit the consecutive error limit
                    if self._consecutive_errors >= self.max_failures:
                        logger.error(f"Terminating: {self.max_failures} consecutive LLM failures")
                        break

                    await asyncio.sleep(1)  # Brief pause before retry
                    continue

                if not code or code.strip() == "":
                    # If task is already done, empty code is fine (LLM explaining completion)
                    if self._is_task_done():
                        logger.info("Task already marked as done, LLM provided explanation without code")
                        # Add the text response to history as a non-code step
                        await self._add_step_to_complete_history(
                            model_output_code="",
                            full_llm_response=full_llm_response,
                            output=full_llm_response,  # Treat the explanation as output
                            error=None,
                            screenshot_path=await self._capture_screenshot(step + 1),
                        )
                        break  # Exit the loop since task is done

                    logger.warning("LLM returned empty code")
                    self._consecutive_errors += 1

                    # new state
                    if self.browser_session and self.dom_service:
                        try:
                            browser_state_text, screenshot = await self._get_browser_state()
                            self._last_browser_state_text = browser_state_text
                            self._last_screenshot = screenshot
                        except Exception as e:
                            logger.warning(f"Failed to get new browser state: {e}")
                    continue

                # Execute code blocks sequentially if multiple python blocks exist
                # This allows JS/bash blocks to be injected into namespace before Python code uses them
                all_blocks = self.namespace.get("_all_code_blocks", {})
                python_blocks = [k for k in sorted(all_blocks.keys()) if k.startswith("python_")]

                if len(python_blocks) > 1:
                    # Multiple Python blocks - execute each sequentially
                    output = None
                    error = None

                    for i, block_key in enumerate(python_blocks):
                        logger.info(f"Executing Python block {i + 1}/{len(python_blocks)}")
                        block_code = all_blocks[block_key]
                        block_output, block_error, _ = await self._execute_code(block_code)

                        # Accumulate outputs
                        if block_output:
                            output = (output or "") + block_output
                        if block_error:
                            error = block_error
                            # Stop on first error
                            break
                else:
                    # Single Python block - execute normally
                    output, error, _ = await self._execute_code(code)

                # Track consecutive errors
                if error:
                    self._consecutive_errors += 1
                    logger.warning(f"Consecutive errors: {self._consecutive_errors}/{self.max_failures}")

                    # Check if we've hit the consecutive error limit
                    if self._consecutive_errors >= self.max_failures:
                        logger.error(
                            f"Terminating: {self.max_failures} consecutive errors reached. The agent is unable to make progress."
                        )
                        await self._demo_mode_log(
                            f"Terminating after {self.max_failures} consecutive errors without progress.",
                            "error",
                            {"step": step + 1},
                        )
                        # Add termination message to complete history before breaking
                        await self._add_step_to_complete_history(
                            model_output_code=code,
                            full_llm_response=f"[Terminated after {self.max_failures} consecutive errors]",
                            output=None,
                            error=f"Auto-terminated: {self.max_failures} consecutive errors without progress",
                            screenshot_path=None,
                        )
                        break
                else:
                    # Reset consecutive error counter on success
                    self._consecutive_errors = 0

                # Check if task is done - validate completion first if not at limits
                if self._is_task_done():
                    # Get the final result from namespace (from done() call)
                    final_result: str | None = self.namespace.get("_task_result")  # type: ignore[assignment]

                    # Check if we should validate (not at step/error limits and under max validations)
                    steps_remaining = self.max_steps - step - 1
                    should_validate = (
                        self._validation_count < self.max_validations  # Haven't exceeded max validations
                        and steps_remaining >= 4  # At least 4 steps away from limit
                        and self._consecutive_errors < 3  # Not close to error limit (8 consecutive)
                    )

                    if should_validate:
                        self._validation_count += 1
                        logger.info("Validating task completion with LLM...")
                        from .namespace import validate_task_completion

                        is_complete, reasoning = await validate_task_completion(
                            task=self.task,
                            output=final_result,
                            llm=self.llm,
                        )

                        if not is_complete:
                            # Task not truly complete - inject feedback and continue
                            logger.warning("Validator: Task not complete, continuing...")
                            validation_feedback = (
                                f"\n\n⚠️ VALIDATOR FEEDBACK:\n"
                                f"Your done() call was rejected. The task is NOT complete yet.\n\n"
                                f"Validation reasoning:\n{reasoning}\n\n"
                                f"You must continue working on the task. Analyze what is missing and complete it.\n"
                                f"Do NOT call done() again until the task is truly finished."
                            )

                            # Clear the done flag so execution continues
                            self.namespace["_task_done"] = False
                            self.namespace.pop("_task_result", None)
                            self.namespace.pop("_task_success", None)

                            # Add validation feedback to LLM messages
                            self._llm_messages.append(UserMessage(content=validation_feedback))

                            # Don't override output - let execution continue normally
                        else:
                            logger.info("Validator: Task complete")
                            # Override output with done message for final step
                            if final_result:
                                output = final_result
                    else:
                        # At limits - skip validation and accept done()
                        if self._validation_count >= self.max_validations:
                            logger.info(
                                f"Reached max validations ({self.max_validations}) - skipping validation and accepting done()"
                            )
                        else:
                            logger.info("At step/error limits - skipping validation")
                        if final_result:
                            output = final_result

                if output:
                    # Check if this is the final done() output
                    if self._is_task_done():
                        # Show done() output more prominently
                        logger.info(
                            f"✓ Task completed - Final output from done():\n{output[:300] if len(output) > 300 else output}"
                        )
                        # Also show files_to_display if they exist in namespace
                        attachments: list[str] | None = self.namespace.get("_task_attachments")  # type: ignore[assignment]
                        if attachments:
                            logger.info(f'Files displayed: {", ".join(attachments)}')
                    else:
                        logger.info(f"Code output:\n{output}")

                # Browser state is now only logged when fetched before LLM call (not after execution)

                # Take screenshot for eval tracking
                screenshot_path = await self._capture_screenshot(step + 1)

                # Add step to complete_history for eval system
                await self._add_step_to_complete_history(
                    model_output_code=code,
                    full_llm_response=full_llm_response,
                    output=output,
                    error=error,
                    screenshot_path=screenshot_path,
                )

                # Check if task is done (after validation)
                if self._is_task_done():
                    # Get the final result from namespace
                    final_result: str | None = self.namespace.get("_task_result", output)  # type: ignore[assignment]
                    logger.info("Task completed successfully")
                    if final_result:
                        logger.info(f"Final result: {final_result}")
                        self._add_sample_output_cell(final_result)
                    if self._demo_mode_enabled:
                        await self._demo_mode_log(
                            f'Final Result: {final_result or "Task completed"}',
                            "success",
                            {"tag": "task"},
                        )
                    should_delay_close = True
                    break
                # If validation rejected done(), continue to next iteration
                # The feedback message has already been added to _llm_messages

                # Add result to LLM messages for next iteration (without browser state)
                result_message = self._format_execution_result(code, output, error, current_step=step + 1)
                truncated_result = truncate_message_content(result_message)
                self._llm_messages.append(UserMessage(content=truncated_result))

            except Exception as e:
                logger.error(f"Error in step {step + 1}: {e}")
                traceback.print_exc()
                break
        else:
            # Loop completed without break - max_steps reached
            logger.warning(f"Maximum steps ({self.max_steps}) reached without task completion")
            await self._demo_mode_log(
                f"Maximum steps ({self.max_steps}) reached without completing the task.",
                "error",
                {"tag": "task"},
            )

        # If task is not done, capture the last step's output as partial result
        if not self._is_task_done() and self.complete_history:
            # Get the last step's output/error and use it as final extracted_content
            last_step = self.complete_history[-1]
            last_result = last_step.result[0] if last_step.result else None
            last_output = last_result.extracted_content if last_result else None
            last_error = last_result.error if last_result else None

            # Build a partial result message from the last step
            partial_result_parts = []
            partial_result_parts.append(f"Task incomplete - reached step limit ({self.max_steps} steps).")
            partial_result_parts.append("Last step output:")

            if last_output:
                partial_result_parts.append(f"\nOutput: {last_output}")
            if last_error:
                partial_result_parts.append(f"\nError: {last_error}")

            # Add any accumulated variables that might contain useful data
            data_vars = []
            for var_name in sorted(self.namespace.keys()):
                if not var_name.startswith("_") and var_name not in {
                    "json",
                    "asyncio",
                    "csv",
                    "re",
                    "datetime",
                    "Path",
                }:
                    var_value = self.namespace[var_name]
                    # Check if it's a list or dict that might contain collected data
                    if isinstance(var_value, (list, dict)) and var_value:
                        data_vars.append(f"  - {var_name}: {type(var_value).__name__} with {len(var_value)} items")

            if data_vars:
                partial_result_parts.append("\nVariables in namespace that may contain partial data:")
                partial_result_parts.extend(data_vars)

            partial_result = "\n".join(partial_result_parts)

            # Update the last step's extracted_content with this partial result
            if last_result:
                last_result.extracted_content = partial_result
                last_result.is_done = False
                last_result.success = False

            logger.info(f"\nPartial result captured from last step:\n{partial_result}")
            if self._demo_mode_enabled:
                await self._demo_mode_log(f"Partial result:\n{partial_result}", "error", {"tag": "task"})

        # Log final summary if task was completed
        if self._is_task_done():
            logger.info("\n" + "=" * 60)
            logger.info("TASK COMPLETED SUCCESSFULLY")
            logger.info("=" * 60)
            final_result: str | None = self.namespace.get("_task_result")  # type: ignore[assignment]
            if final_result:
                logger.info(f"\nFinal Output:\n{final_result}")
                self._add_sample_output_cell(final_result)

            attachments: list[str] | None = self.namespace.get("_task_attachments")  # type: ignore[assignment]
            if attachments:
                logger.info(f"\nFiles Attached:\n{chr(10).join(attachments)}")
            logger.info("=" * 60 + "\n")
            if self._demo_mode_enabled and not should_delay_close:
                await self._demo_mode_log(
                    f'Final Result: {final_result or "Task completed"}',
                    "success",
                    {"tag": "task"},
                )
                should_delay_close = True

        # Auto-close browser if keep_alive is False
        if should_delay_close and self._demo_mode_enabled:
            await asyncio.sleep(30)
        await self.close()

        # Store usage summary for history property
        self.usage_summary = await self.token_cost_service.get_usage_summary()

        # Log token usage summary
        await self.token_cost_service.log_usage_summary()

        # Store history data in session for history property
        self.session._complete_history = self.complete_history
        self.session._usage_summary = self.usage_summary

        return self.session

    async def _get_code_from_llm(self, step_number: int | None = None) -> tuple[str, str]:
        """Get Python code from the LLM.

        Returns:
                Tuple of (extracted_code, full_llm_response)
        """
        # Prepare messages for this request
        # Include browser state as separate message if available (not accumulated in history)
        messages_to_send = self._llm_messages.copy()

        if self._last_browser_state_text:
            # Create message with optional screenshot
            if self.use_vision and self._last_screenshot:
                # Build content with text + screenshot
                content_parts: list[ContentPartTextParam | ContentPartImageParam] = [
                    ContentPartTextParam(text=self._last_browser_state_text)
                ]

                # Add screenshot
                content_parts.append(
                    ContentPartImageParam(
                        image_url=ImageURL(
                            url=f"data:image/png;base64,{self._last_screenshot}",
                            media_type="image/png",
                            detail="auto",
                        ),
                    )
                )

                messages_to_send.append(UserMessage(content=content_parts))
            else:
                # Text only
                messages_to_send.append(UserMessage(content=self._last_browser_state_text))

            # Clear browser state after including it so it's only in this request
            self._last_browser_state_text = None
            self._last_screenshot = None

        # Call LLM with message history (including temporary browser state message)
        response = await self.llm.ainvoke(messages_to_send)

        # Store usage stats from this LLM call
        self._last_llm_usage = response.usage

        # Log the LLM's raw output for debugging
        logger.info(f"LLM Response:\n{response.completion}")
        await self._demo_mode_log(
            f"LLM Response:\n{response.completion}",
            "thought",
            {"step": step_number} if step_number else None,
        )

        # Check for token limit or repetition issues
        max_tokens = getattr(self.llm, "max_tokens", None)
        completion_tokens = response.usage.completion_tokens if response.usage else None
        is_problematic, issue_message = detect_token_limit_issue(
            completion=response.completion,
            completion_tokens=completion_tokens,
            max_tokens=max_tokens,
            stop_reason=response.stop_reason,
        )

        if is_problematic:
            logger.warning(f"Token limit issue detected: {issue_message}")
            # Don't add the bad response to history
            # Instead, inject a system message prompting recovery
            recovery_prompt = (
                f"Your previous response hit a token limit or became repetitive: {issue_message}\n\n"
                "Please write a SHORT plan (2 sentences) for what to do next, then execute ONE simple action."
            )
            self._llm_messages.append(UserMessage(content=recovery_prompt))
            # Return a controlled error message instead of corrupted code
            return "", f"[Token limit error: {issue_message}]"

        # Store the full response
        full_response = response.completion

        # Extract code blocks from response
        # Support multiple code block types: python, js, bash, markdown
        code_blocks = extract_code_blocks(response.completion)

        # Inject non-python blocks into namespace as variables
        # Track which variables are code blocks for browser state display
        if "_code_block_vars" not in self.namespace:
            self.namespace["_code_block_vars"] = set()

        for block_type, block_content in code_blocks.items():
            if not block_type.startswith("python"):
                # Store js, bash, markdown blocks (and named variants) as variables in namespace
                self.namespace[block_type] = block_content
                self.namespace["_code_block_vars"].add(block_type)
                print(f"→ Code block variable: {block_type} (str, {len(block_content)} chars)")
                logger.debug(f"Injected {block_type} block into namespace ({len(block_content)} chars)")

        # Store all code blocks for sequential execution
        self.namespace["_all_code_blocks"] = code_blocks

        # Get Python code if it exists
        # If no python block exists and no other code blocks exist, return empty string to skip execution
        # This prevents treating plain text explanations as code
        code = code_blocks.get("python", response.completion)

        # Add to LLM messages (truncate for history to save context)
        truncated_completion = truncate_message_content(response.completion)
        self._llm_messages.append(AssistantMessage(content=truncated_completion))

        return code, full_response

    def _print_variable_info(self, var_name: str, value: Any) -> None:
        """Print compact info about a variable assignment."""
        # Skip built-in modules and known imports
        skip_names = {
            "json",
            "asyncio",
            "csv",
            "re",
            "datetime",
            "Path",
            "pd",
            "np",
            "plt",
            "requests",
            "BeautifulSoup",
            "PdfReader",
            "browser",
            "file_system",
        }
        if var_name in skip_names:
            return

        # Skip code block variables (already printed)
        if "_code_block_vars" in self.namespace and var_name in self.namespace.get("_code_block_vars", set()):
            return

        # Print compact variable info
        if isinstance(value, (list, dict)):
            preview = str(value)[:100]
            print(f"→ Variable: {var_name} ({type(value).__name__}, len={len(value)}, preview={preview}...)")
        elif isinstance(value, str) and len(value) > 50:
            print(f"→ Variable: {var_name} (str, {len(value)} chars, preview={value[:50]}...)")
        elif callable(value):
            print(f"→ Variable: {var_name} (function)")
        else:
            print(f"→ Variable: {var_name} ({type(value).__name__}, value={repr(value)[:50]})")

    async def _execute_code(self, code: str) -> tuple[str | None, str | None, str | None]:
        """
        Execute Python code in the namespace.

        Args:
                code: The Python code to execute

        Returns:
                Tuple of (output, error, browser_state)
        """
        # Create new cell
        cell = self.session.add_cell(source=code)
        cell.status = ExecutionStatus.RUNNING
        cell.execution_count = self.session.increment_execution_count()

        output = None
        error = None

        try:
            # Capture output
            import ast
            import io
            import sys

            old_stdout = sys.stdout
            sys.stdout = io.StringIO()

            try:
                # Add asyncio to namespace if not already there
                if "asyncio" not in self.namespace:
                    self.namespace["asyncio"] = asyncio

                # Store the current code in namespace for done() validation
                self.namespace["_current_cell_code"] = code
                # Store consecutive errors count for done() validation
                self.namespace["_consecutive_errors"] = self._consecutive_errors

                # Check if code contains await expressions - if so, wrap in async function
                # This mimics how Jupyter/IPython handles top-level await
                try:
                    tree = ast.parse(code, mode="exec")
                    has_await = any(
                        isinstance(node, (ast.Await, ast.AsyncWith, ast.AsyncFor)) for node in ast.walk(tree)
                    )
                except SyntaxError:
                    # If parse fails, let exec handle the error
                    has_await = False

                if has_await:
                    # When code has await, we must wrap in async function
                    # To make variables persist naturally (like Jupyter without needing 'global'):
                    # 1. Extract all assigned variable names from the code
                    # 2. Inject 'global' declarations for variables that already exist in namespace
                    # 3. Extract user's explicit global declarations and pre-define those vars
                    # 4. Return locals() so we can update namespace with new variables

                    # Find all variable names being assigned + user's explicit globals
                    try:
                        assigned_names = set()
                        user_global_names = set()

                        for node in ast.walk(tree):
                            if isinstance(node, ast.Assign):
                                for target in node.targets:
                                    if isinstance(target, ast.Name):
                                        assigned_names.add(target.id)
                            elif isinstance(node, ast.AugAssign) and isinstance(node.target, ast.Name):
                                assigned_names.add(node.target.id)
                            elif isinstance(node, (ast.AnnAssign, ast.NamedExpr)):
                                if hasattr(node, "target") and isinstance(node.target, ast.Name):
                                    assigned_names.add(node.target.id)
                            elif isinstance(node, ast.Global):
                                # Track user's explicit global declarations
                                user_global_names.update(node.names)

                        # Pre-define any user-declared globals that don't exist yet
                        # This prevents NameError when user writes "global foo" before "foo = ..."
                        for name in user_global_names:
                            if name not in self.namespace:
                                self.namespace[name] = None

                        # Filter to only existing namespace vars (like Jupyter does)
                        # Include both: assigned vars that exist + user's explicit globals
                        existing_vars = {
                            name for name in (assigned_names | user_global_names) if name in self.namespace
                        }
                    except Exception:
                        existing_vars = set()

                    # Build global declaration if needed
                    global_decl = ""
                    has_global_decl = False
                    if existing_vars:
                        vars_str = ", ".join(sorted(existing_vars))
                        global_decl = f"    global {vars_str}\n"
                        has_global_decl = True

                    indented_code = "\n".join("    " + line if line.strip() else line for line in code.split("\n"))
                    wrapped_code = f"""async def __code_exec__():
{global_decl}{indented_code}
    # Return locals so we can update the namespace
    return locals()

__code_exec_coro__ = __code_exec__()
"""
                    # Store whether we added a global declaration (needed for error line mapping)
                    self.namespace["_has_global_decl"] = has_global_decl

                    # Compile and execute wrapper at module level
                    compiled_code = compile(wrapped_code, "<code>", "exec")
                    exec(compiled_code, self.namespace, self.namespace)

                    # Get and await the coroutine, then update namespace with new/modified variables
                    coro = self.namespace.get("__code_exec_coro__")
                    if coro:
                        result_locals = await coro
                        # Update namespace with all variables from the function's locals
                        # This makes variable assignments persist across cells
                        if result_locals:
                            for key, value in result_locals.items():
                                if not key.startswith("_"):
                                    self.namespace[key] = value
                                    # Variable info is tracked in "Available" section, no need for verbose inline output

                        # Clean up temporary variables
                        self.namespace.pop("__code_exec_coro__", None)
                        self.namespace.pop("__code_exec__", None)
                else:
                    # No await - execute directly at module level for natural variable scoping
                    # This means x = x + 10 will work without needing 'global x'

                    # Track variables before execution
                    vars_before = set(self.namespace.keys())

                    compiled_code = compile(code, "<code>", "exec")
                    exec(compiled_code, self.namespace, self.namespace)

                    # Track newly created/modified variables (info shown in "Available" section)
                    vars_after = set(self.namespace.keys())
                    vars_after - vars_before

                # Get output
                output_value = sys.stdout.getvalue()
                if output_value:
                    output = output_value

            finally:
                sys.stdout = old_stdout

            # Wait 2 seconds for page to stabilize after code execution
            await asyncio.sleep(0.5)

            # Note: Browser state is now fetched right before LLM call instead of after each execution
            # This reduces unnecessary state fetches for operations that don't affect the browser

            cell.status = ExecutionStatus.SUCCESS
            cell.output = output
            cell.browser_state = None  # Will be captured in next iteration before LLM call

        except Exception as e:
            # Handle EvaluateError specially - JavaScript execution failed
            if isinstance(e, EvaluateError):
                error = str(e)
                cell.status = ExecutionStatus.ERROR
                cell.error = error
                logger.error(f"Code execution error: {error}")

                await asyncio.sleep(1)

                # Browser state will be fetched before next LLM call
                # Return immediately - do not continue executing code
                return output, error, None

            # Handle NameError specially - check for code block variable confusion
            if isinstance(e, NameError):
                error_msg = str(e)
                cell.status = ExecutionStatus.ERROR
                cell.error = error

                # Browser state will be fetched before next LLM call
                await asyncio.sleep(0.5)
                return output, error, None

            # For syntax errors and common parsing errors, show just the error message
            # without the full traceback to keep output clean
            if isinstance(e, SyntaxError):
                error_msg = e.msg if e.msg else str(e)
                error = f"{type(e).__name__}: {error_msg}"

                # Detect common f-string issues with JSON/JavaScript code
                if "unterminated" in error_msg.lower() and "string" in error_msg.lower() and code:
                    # Check if code contains f-strings with potential JSON/JS content
                    has_fstring = bool(re.search(r'\bf["\']', code))
                    has_json_pattern = bool(
                        re.search(r'json\.dumps|"[^"]*\{[^"]*\}[^"]*"|\'[^\']*\{[^\']*\}[^\']*\'', code)
                    )
                    has_js_pattern = bool(re.search(r"evaluate\(|await evaluate", code))

                    if has_fstring and (has_json_pattern or has_js_pattern):
                        error += (
                            "\n\n💡 TIP: Detected f-string with JSON/JavaScript code containing {}.\n"
                            "   Use separate ```js or ```markdown blocks instead of f-strings to avoid escaping issues.\n"
                            "   If your code block needs ``` inside it, wrap with 4+ backticks: ````markdown code`\n"
                        )

                # Detect and provide helpful hints for common string literal errors
                if "unterminated" in error_msg.lower() and "string" in error_msg.lower():
                    # Detect what type of string literal is unterminated
                    is_triple = "triple-quoted" in error_msg.lower()
                    msg_lower = error_msg.lower()

                    # Detect prefix type from error message
                    if "f-string" in msg_lower and "raw" in msg_lower:
                        prefix = "rf or fr"
                        desc = "raw f-string"
                    elif "f-string" in msg_lower:
                        prefix = "f"
                        desc = "f-string"
                    elif "raw" in msg_lower and "bytes" in msg_lower:
                        prefix = "rb or br"
                        desc = "raw bytes"
                    elif "raw" in msg_lower:
                        prefix = "r"
                        desc = "raw string"
                    elif "bytes" in msg_lower:
                        prefix = "b"
                        desc = "bytes"
                    else:
                        prefix = ""
                        desc = "string"

                    # Build hint based on triple-quoted vs single/double quoted
                    if is_triple:
                        if prefix:
                            hint = f"Hint: Unterminated {prefix}'''...''' or {prefix}\"\"\"...\"\" ({desc}). Check for missing closing quotes or unescaped quotes inside."
                        else:
                            hint = "Hint: Unterminated '''...''' or \"\"\"...\"\" detected. Check for missing closing quotes or unescaped quotes inside."
                        hint += "\n      If you need ``` inside your string, use a ````markdown varname` code block with 4+ backticks instead."
                    else:
                        if prefix:
                            hint = f"Hint: Unterminated {prefix}'...' or {prefix}\"...\" ({desc}). Check for missing closing quote or unescaped quotes inside."
                        else:
                            hint = "Hint: Unterminated '...' or \"...\" detected. Check for missing closing quote or unescaped quotes inside the string."
                    error += f"\n{hint}"

                # Show the problematic line from the code
                if e.text:
                    error += f"\n{e.text}"
                elif e.lineno and code:
                    # If e.text is empty, extract the line from the code
                    lines = code.split("\n")
                    if 0 < e.lineno <= len(lines):
                        error += f"\n{lines[e.lineno - 1]}"

            else:
                # For other errors, try to extract useful information
                error_str = str(e)
                error = f"{type(e).__name__}: {error_str}" if error_str else f"{type(e).__name__} occurred"

                # For RuntimeError or other exceptions, try to extract traceback info
                # to show which line in the user's code actually failed
                if hasattr(e, "__traceback__"):
                    # Walk the traceback to find the frame with '<code>' filename
                    tb = e.__traceback__
                    while tb is not None:
                        frame = tb.tb_frame
                        if frame.f_code.co_filename == "<code>":
                            # Found the frame executing user code
                            # Get the line number from the traceback
                            tb.tb_lineno
                            break
                        tb = tb.tb_next

            cell.status = ExecutionStatus.ERROR
            cell.error = error
            logger.error(f"Code execution error: {error}")

            await asyncio.sleep(1)

            # Browser state will be fetched before next LLM call

        return output, error, None

    async def _get_browser_state(self) -> tuple[str, str | None]:
        """Get the current browser state as text with ultra-minimal DOM structure for code agents.

        Returns:
                Tuple of (browser_state_text, screenshot_base64)
        """
        if not self.browser_session or not self.dom_service:
            return "Browser state not available", None

        try:
            # Get full browser state including screenshot if use_vision is enabled
            include_screenshot = True
            state = await self.browser_session.get_browser_state_summary(include_screenshot=include_screenshot)

            # Format browser state with namespace context
            browser_state_text = await format_browser_state_for_llm(
                state=state, namespace=self.namespace, browser_session=self.browser_session
            )

            screenshot = state.screenshot if include_screenshot else None
            return browser_state_text, screenshot

        except Exception as e:
            logger.error(f"Failed to get browser state: {e}")
            return f"Error getting browser state: {e}", None

    def _format_execution_result(
        self, code: str, output: str | None, error: str | None, current_step: int | None = None
    ) -> str:
        """Format the execution result for the LLM (without browser state)."""
        result = []

        # Add step progress header if step number provided
        if current_step is not None:
            progress_header = f"Step {current_step}/{self.max_steps} executed"
            # Add consecutive failure tracking if there are errors
            if error and self._consecutive_errors > 0:
                progress_header += f" | Consecutive failures: {self._consecutive_errors}/{self.max_failures}"
            result.append(progress_header)

        if error:
            result.append(f"Error: {error}")

        if output:
            # Truncate output if too long
            if len(output) > 10000:
                output = output[:9950] + "\n[Truncated after 10000 characters]"
            result.append(f"Output: {output}")
        if len(result) == 0:
            result.append("Executed")
        return "\n".join(result)

    def _is_task_done(self) -> bool:
        """Check if the task is marked as done in the namespace."""
        # Check if 'done' was called by looking for a special marker in namespace
        return self.namespace.get("_task_done", False)

    async def _capture_screenshot(self, step_number: int) -> str | None:
        """Capture and store screenshot for eval tracking."""
        if not self.browser_session:
            return None

        try:
            # Get browser state summary which includes screenshot
            state = await self.browser_session.get_browser_state_summary(include_screenshot=True)
            if state and state.screenshot:
                # Store screenshot using screenshot service
                screenshot_path = await self.screenshot_service.store_screenshot(state.screenshot, step_number)
                return str(screenshot_path) if screenshot_path else None
        except Exception as e:
            logger.warning(f"Failed to capture screenshot for step {step_number}: {e}")
            return None

    async def _add_step_to_complete_history(
        self,
        model_output_code: str,
        full_llm_response: str,
        output: str | None,
        error: str | None,
        screenshot_path: str | None,
    ) -> None:
        """Add a step to complete_history using type-safe models."""
        # Get current browser URL and title for state
        url: str | None = None
        title: str | None = None
        if self.browser_session:
            try:
                url = await self.browser_session.get_current_page_url()
                # Get title from browser
                cdp_session = await self.browser_session.get_or_create_cdp_session()
                result = await cdp_session.cdp_client.send.Runtime.evaluate(
                    params={"expression": "document.title", "returnByValue": True},
                    session_id=cdp_session.session_id,
                )
                title = result.get("result", {}).get("value")
            except Exception as e:
                logger.debug(f"Failed to get browser URL/title for history: {e}")

        # Check if this is a done result
        is_done = self._is_task_done()

        # Get self-reported success from done() call if task is done
        self_reported_success: bool | None = None
        if is_done:
            task_success = self.namespace.get("_task_success")
            self_reported_success = task_success if isinstance(task_success, bool) else None

        # Create result entry using typed model
        result_entry = CodeAgentResult(
            extracted_content=output if output else None,
            error=error if error else None,
            is_done=is_done,
            success=self_reported_success,
        )

        # Create state entry using typed model
        state_entry = CodeAgentState(url=url, title=title, screenshot_path=screenshot_path)

        # Create metadata entry using typed model
        step_end_time = datetime.datetime.now().timestamp()
        metadata_entry = CodeAgentStepMetadata(
            input_tokens=self._last_llm_usage.prompt_tokens if self._last_llm_usage else None,
            output_tokens=self._last_llm_usage.completion_tokens if self._last_llm_usage else None,
            step_start_time=self._step_start_time,
            step_end_time=step_end_time,
        )

        # Create model output entry using typed model (if there's code to track)
        model_output_entry: CodeAgentModelOutput | None = None
        if model_output_code or full_llm_response:
            model_output_entry = CodeAgentModelOutput(
                model_output=model_output_code if model_output_code else "",
                full_response=full_llm_response if full_llm_response else "",
            )

        # Create history entry using typed model
        history_entry = CodeAgentHistory(
            model_output=model_output_entry,
            result=[result_entry],
            state=state_entry,
            metadata=metadata_entry,
            screenshot_path=screenshot_path,  # Keep for backward compatibility
        )

        self.complete_history.append(history_entry)
        await self._demo_mode_log_step(history_entry)

    async def _demo_mode_log(self, message: str, level: str = "info", metadata: dict[str, Any] | None = None) -> None:
        if not (self._demo_mode_enabled and message and self.browser_session):
            return
        try:
            await self.browser_session.send_demo_mode_log(
                message=message,
                level=level,
                metadata=metadata or {},
            )
        except Exception as exc:
            logger.debug(f"[DemoMode] Failed to send log: {exc}")

    async def _demo_mode_log_step(self, history_entry: CodeAgentHistory) -> None:
        if not self._demo_mode_enabled:
            return
        step_number = len(self.complete_history)
        result = history_entry.result[0] if history_entry.result else None
        if not result:
            return
        level = "error" if result.error else "success" if result.success else "info"
        message_parts = [f"Step {step_number}:"]
        if result.error:
            message_parts.append(f"Error: {result.error}")
        if result.extracted_content:
            message_parts.append(result.extracted_content)
        elif result.success:
            message_parts.append("Marked done.")
        else:
            message_parts.append("Executed.")
        await self._demo_mode_log(
            " ".join(message_parts).strip(),
            level,
            {"step": step_number, "url": history_entry.state.url if history_entry.state else None},
        )

    def _add_sample_output_cell(self, final_result: Any | None) -> None:
        if self._sample_output_added or final_result is None:
            return

        sample_content: str | None = None

        def _extract_sample(data: Any) -> Any | None:
            if isinstance(data, list) and data:
                return data[0]
            if isinstance(data, dict) and data:
                first_key = next(iter(data))
                return {first_key: data[first_key]}
            return data if isinstance(data, (str, int, float, bool)) else None

        data: Any | None = None
        if isinstance(final_result, str):
            try:
                data = json.loads(final_result)
            except Exception:
                sample_content = final_result.strip()
        elif isinstance(final_result, (list, dict)):
            data = final_result

        if data is not None:
            sample = _extract_sample(data)
            if isinstance(sample, (dict, list)):
                try:
                    sample_content = json.dumps(sample, indent=2, ensure_ascii=False)
                except Exception:
                    sample_content = str(sample)
            elif sample is not None:
                sample_content = str(sample)

        if not sample_content:
            return

        sample_cell = self.session.add_cell(source="# Sample output preview")
        sample_cell.cell_type = CellType.MARKDOWN
        sample_cell.status = ExecutionStatus.SUCCESS
        sample_cell.execution_count = None
        escaped = html.escape(sample_content)
        sample_cell.output = f"<pre>{escaped}</pre>"

        self._sample_output_added = True

    def screenshot_paths(self, n_last: int | None = None) -> list[str | None]:
        """
        Get screenshot paths from complete_history for eval system.

        Args:
                n_last: Optional number of last screenshots to return

        Returns:
                List of screenshot file paths (or None for missing screenshots)
        """
        paths = [step.screenshot_path for step in self.complete_history]

        if n_last is not None:
            return paths[-n_last:] if len(paths) > n_last else paths

        return paths

    @property
    def message_manager(self) -> Any:
        """
        Compatibility property for eval system.
        Returns a mock object with last_input_messages attribute.
        """

        class MockMessageManager:
            def __init__(self, llm_messages: list[BaseMessage]) -> None:
                # Convert code-use LLM messages to format expected by eval system
                self.last_input_messages = llm_messages

        return MockMessageManager(self._llm_messages)

    @property
    def history(self) -> CodeAgentHistoryList:
        """
        Compatibility property for eval system.
        Returns a CodeAgentHistoryList object with history attribute containing complete_history.
        This is what the eval system expects when it does: agent_history = agent.history
        """
        return CodeAgentHistoryList(self.complete_history, self.usage_summary)

    async def close(self) -> None:
        """Close the browser session."""
        if self.browser_session:
            # Check if we should close the browser based on keep_alive setting
            if not self.browser_session.browser_profile.keep_alive:
                await self.browser_session.kill()
            else:
                logger.debug("Browser keep_alive is True, not closing browser session")

    async def __aenter__(self) -> "CodeAgent":
        """Async context manager entry."""
        return self

    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
        """Async context manager exit."""
        await self.close()
