from langchain_core.tools import tool
from loguru import logger

from universal_mcp.agents.codeact0.config import CodeActConfig, get_config
from universal_mcp.agents.Sandbox.subprocess_sandbox import SubprocessSandbox

# Truncation constants
MAX_CHARS_FULL = 5000
MAX_LINES_HEADTAIL = 50
SUMMARY_THRESHOLD = 10000


@tool(parse_docstring=True)
def execute_python_code(step_title: str, snippet: str, display_output_to_user: bool = False) -> str:
    """
    Executes a Python code snippet in a sandbox with retained context (top level defined functions, variables, loaded functions using `load_functions` are retained). Briefly describe (in one line) what you intend to do before each tool call. Do not put this line as a comment inside the snippet, it should be before the tool call.


    **Design Principles**:
    - Write concise code and avoid repeating lines from previous snippets that have already executed.
    - Break logic into multiple small helper functions (with names starting with _, and having max 20 lines each), that each do a single atomic task.
    - ALL functions MUST be declared with `async def`**. Never use `def` for any function. Use `await` when calling async functions.
    - Keep large constants (e.g., multiline strings, dicts, json schemas) global or in a dedicated helper function. Do not declare them inside a function responsible for performing another task.
    - Modify only the relevant helper during debugging—context persists across executions.
    - ALWAYS use keyword arguments ONLY when calling functions (both pre-defined and your own helper functions) to prevent positional argument errors and improve readability.
    - ALWAYS use smart_print() to examine external function outputs if they are dict/list[dict] before using them, or displaying them to the user.
    - Example:
        async def _get_json_schema():
            return {"key1":"many details"...}
        async def _helper_function_1(...):
            ...

        async def _helper_function_2(...):
            ...
        result1 = await _helper_function_1(..., await _get_json_schema())
        smart_print(result1[:1]) #As an example, to check if it has been processed correctly. Taking care not to print large lists
        result2 = await _helper_function_2(...)
        smart_print(result2[:1])
        final_result = ...
        smart_print(final_result)
        - Thus, while debugging, if you face an error in result2, you do not need to rewrite _helper_function_1() or _get_json_schema().
    - You have preloaded functions, including a web_search and intelligent language processing and generation functions (llm). Follow the following with respect to LLM functions-
        **CRITICAL INSTRUCTION VIOLATIONS TO CHECK BEFORE EXECUTION:**
        - [ ] Am I using regex/manual parsing/string manipulation for data extraction? -> STOP, use llm__extract_data
        - [ ] Am I hardcoding patterns for text analysis? -> STOP, use llm__classify_data
        - [ ] Am I writing textual/report or any large static text content myself-> STOP, use llm__generate_text

        **MANDATORY Pre-execution Checklist:**
        Before writing ANY code that processes or generates text content:
        1. Is this data extraction? → Use llm__extract_data
        2. Is this classification/comparison? → Use llm__classify_data
        3. Is this text analysis? → Use LLM tools
        4. Is this text generation (e.g. a markdown report, content for a document, HTML report, large multiline strings etc) → use llm__generate_text
        5. Only use manual parsing for: file paths, URLs, structured data formats

    - You can only import libraries that come pre-installed with Python, or pandas. However, do consider using preloaded functions or searching for external functions first, using the search and load tools to access them in the code.
    - For handling csv, spreadsheets, and Excel files, import pandas and any of its dependencies. Use engines `pyarrow` for reading csv files and `calamine` for excel files.
    - Do not handle powerpoint files or similar complex file formats using code. Politely refuse such requests and help as much as possible using external functions.
    - Do not use this to create large-scale, multi-file applications or web-apps. Politely refuse requests beyond a simple website.
    - Use loops to process multiple items—test once before scaling.
    - Do not use this tool to think or communicate to the user (through comments or otherwise), except when using this to summarize your work done as the final step with display_output_to_user=True. Communication to the user should be in a normal message right before this tool call, briefly informing them about your plan for the tool call.

    Args:
        step_title (str): A 5-10 word non-technical title for the user explaining the snippet (e.g. "Fetching data from the product's website", "Web searching for competitors in the market")
        snippet (str): Python code to execute.
        display_output_to_user (bool, default - False): To show the output of a snippet directly to the user. Use this when displaying final results. Only use this when you are sure that the code will execute correctly (i.e. functions and variables have already been used/populated), and the output will be in Markdown format (i.e. Markdown formatting of titles, sections and no dictionaries or lists of dictionaries). Prefer using existing variables with f-string for displaying details (e.g. links, data). This markdown must follow the general output rules.

    Returns:
        Execution result or error as a string.

    Raises:
        ValueError: If snippet is empty.
    """
    # This is a dummy implementation - actual execution happens in Sandbox.handle_execute_python_code
    if not snippet or not snippet.strip():
        raise ValueError("Parameter 'snippet' is required and cannot be empty or whitespace")
    return f"Successfully executed {len(snippet)} characters of Python code"


class Sandbox(SubprocessSandbox):
    """CodeAct sandbox that extends SubprocessSandbox with smart truncation."""

    def __init__(self, sandbox_timeout: int | None = None, config: CodeActConfig | None = None):
        self.config = config or get_config()
        timeout = sandbox_timeout if sandbox_timeout is not None else self.config.sandbox.timeout
        super().__init__(timeout=timeout)
        # Add predefined type aliases directly to namespace
        self.namespace["File"] = str
        self.namespace["OnlineResource"] = str
        self.namespace["LongText"] = str

    @staticmethod
    def smart_truncate(
        output: str,
        max_chars_full: int | None = None,
        max_lines_headtail: int | None = None,
        summary_threshold: int | None = None,
        config: CodeActConfig | None = None,
    ) -> str:
        """
        Truncates or summarizes output intelligently to avoid filling the context too fast.

        Args:
            output (str): The string output from code execution.
            max_chars_full (int): Max characters to keep full output.
            max_lines_headtail (int): Number of lines to keep from head and tail for medium outputs.
            summary_threshold (int): If truncated output exceeds this, hard-truncate.
            config (CodeActConfig): Optional config to override defaults.

        Returns:
            str: Truncated or summarized output.
        """
        max_chars_full = max_chars_full if max_chars_full is not None else MAX_CHARS_FULL
        max_lines_headtail = max_lines_headtail if max_lines_headtail is not None else MAX_LINES_HEADTAIL
        summary_threshold = summary_threshold if summary_threshold is not None else SUMMARY_THRESHOLD

        if len(output) <= max_chars_full:
            return output  # Small output, include fully

        lines = output.splitlines()
        if len(lines) > 2 * max_lines_headtail:
            # Medium-large output: take head + tail
            head = "\n".join(lines[:max_lines_headtail])
            tail = "\n".join(lines[-max_lines_headtail:])
            output = f"{head}\n... [truncated {len(lines) - 2 * max_lines_headtail} lines] ...\n{tail}"

        if len(output) > summary_threshold:
            truncate_msg = "\n... [output truncated to fit context] ..."
            if summary_threshold > len(truncate_msg):
                output = output[: summary_threshold - len(truncate_msg)] + truncate_msg
            else:
                output = output[:summary_threshold]

        return output

    async def handle_execute_python_code(self, code: str, display_output_to_user: bool = False) -> tuple[str, str]:
        """
        Execute Python code in the sandbox.

        Context is already loaded in self.namespace by route_entry.
        After execution, serializes and returns the updated namespace.

        Args:
            code: Python code to execute

        Returns:
            tuple[str, str]: (output, new_context_serialized)
                - output: stdout or error message from code execution
                - new_context_serialized: Base64-encoded cloudpickle string of updated context
        """
        logger.info("Executing code with context keys: %s", list(self.namespace.keys()))

        # Execute code (context already in self.namespace)
        result = await self.run(code)

        # Format output
        if result["exit_code"] == 0:
            output = result["stdout"] or "<code ran, no output printed to stdout>"
        else:
            output = result["error_message"] or result["stderr"] or "Unknown error"

        # Truncate output only if not for user
        if not display_output_to_user:
            output = self.smart_truncate(output)

        # Get serialized context from sandbox
        new_context_serialized = await self.get_context()

        logger.info("After execution, context keys: %s", list(self._filter_picklable(self.namespace).keys()))

        return output, new_context_serialized
