"""Code interpreter tool — runs Python code in AgentCore managed sandbox."""
import os
import base64
from strands import tool


def _parse_stream(result: dict) -> str:
    """Parse the streaming response from invoke_code_interpreter."""
    parts = []
    if "stream" not in result:
        return str(result)

    for event in result["stream"]:
        if "result" not in event:
            continue
        for item in event["result"].get("content", []):
            item_type = item.get("type", "")
            if item_type == "text":
                text = item.get("text", "")
                if text:
                    parts.append(text)
            elif item_type == "resource":
                resource = item.get("resource", {})
                if "text" in resource:
                    parts.append(resource["text"])
                elif "blob" in resource:
                    try:
                        parts.append(base64.b64decode(resource["blob"]).decode("utf-8"))
                    except Exception:
                        parts.append(f"<binary resource: {resource.get('uri', '?')}>")
            elif item_type == "image":
                # Base64-encoded image
                image_data = item.get("source", {}).get("data", "")
                mime = item.get("source", {}).get("mediaType", "image/png")
                parts.append(f"<image: {mime}, {len(image_data)} bytes base64>")

    return "\n".join(parts) if parts else "(no output)"


@tool
def run_code(code: str, packages: list[str] | None = None) -> str:
    """Execute Python code in a secure managed sandbox and return the output.
    Optionally install pip packages before running (e.g. ['pandas', 'numpy']).

    Args:
        code: Python code to execute.
        packages: Optional list of pip packages to install first.

    Returns:
        Execution output (stdout, results, errors).
    """
    try:
        from bedrock_agentcore.tools import CodeInterpreter, code_session

        region = os.environ.get('AWS_REGION', 'us-east-1')

        with code_session(region) as client:
            if packages:
                install_raw = client.install_packages(packages)
                install_out = _parse_stream(install_raw) if isinstance(install_raw, dict) else str(install_raw)
                print(f'[code_interpreter] install: {install_out[:200]}')

            raw = client.execute_code(code)
            return _parse_stream(raw)

    except Exception as e:
        import traceback
        return f'Code interpreter error: {type(e).__name__}: {e}\n{traceback.format_exc()[-500:]}'