"""E2E tests for BaponiSandbox deep agents integration against staging.

Tests:
1. Direct SDK: execute, upload_files, download_files
2. Real deep agent with Claude Sonnet 4.6: observe tool calls
"""

import os
import sys
import time
import uuid
from pathlib import Path

from dotenv import load_dotenv

# ── Config ──────────────────────────────────────────────────────────────────
# Load from sdks/python/.env (gitignored)
load_dotenv(Path(__file__).resolve().parent.parent / ".env")

BAPONI_API_KEY = os.environ["BAPONI_API_KEY"]
BAPONI_BASE_URL = os.environ.get("BAPONI_BASE_URL", "https://api.baponi.ai")
ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"]

os.environ["ANTHROPIC_API_KEY"] = ANTHROPIC_API_KEY

THREAD_ID = f"e2e-deep-{uuid.uuid4().hex[:12]}"

# ── Helpers ─────────────────────────────────────────────────────────────────

def ok(msg: str) -> None:
    print(f"  [PASS] {msg}")

def fail(msg: str) -> None:
    print(f"  [FAIL] {msg}")
    sys.exit(1)

def section(title: str) -> None:
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}")


# ── Test 1: Direct SDK ──────────────────────────────────────────────────────

def test_direct_sdk():
    section("Test 1: Direct SDK - execute, upload, download")

    from baponi.deepagents import BaponiSandbox

    with BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=THREAD_ID,
        timeout=25,
    ) as sandbox:
        print(f"  Thread ID: {sandbox.id}")

        # 1a. Execute
        print("\n  --- 1a. Execute 'echo hello' ---")
        result = sandbox.execute("echo hello")
        assert result.exit_code == 0, f"exit_code={result.exit_code}"
        assert "hello" in result.output, f"output={result.output!r}"
        ok(f"execute: exit_code=0, output={result.output.strip()!r}")

        # 1b. Upload file via signed URL
        print("\n  --- 1b. Upload file via signed URL ---")
        test_content = b"Hello from E2E test!\nLine 2.\n"
        upload_result = sandbox.upload_files([("/home/baponi/e2e-test.txt", test_content)])
        assert len(upload_result) == 1
        if upload_result[0].error:
            fail(f"upload error: {upload_result[0].error}")
        ok(f"upload_files: path={upload_result[0].path}, no error")

        # 1c. Verify file is visible in sandbox (tests vfs_forget fix!)
        print("\n  --- 1c. Verify uploaded file visible in sandbox ---")
        result = sandbox.execute("cat /home/baponi/e2e-test.txt")
        if result.exit_code != 0:
            print(f"  stdout: {result.output!r}")
            fail(f"cat failed with exit_code={result.exit_code} - vfs_forget may not be working")
        if "Hello from E2E test!" not in result.output:
            fail(f"File content mismatch: {result.output!r}")
        ok(f"File visible in sandbox after signed URL upload (vfs_forget works!)")

        # 1d. Write file inside sandbox, then download via signed URL
        print("\n  --- 1d. Write file in sandbox, download via signed URL ---")
        result = sandbox.execute("echo 'sandbox-written-content' > /home/baponi/from-sandbox.txt")
        assert result.exit_code == 0, f"write failed: {result.output}"

        # Small delay for storage writeback
        time.sleep(2)

        download_result = sandbox.download_files(["/home/baponi/from-sandbox.txt"])
        assert len(download_result) == 1
        if download_result[0].error:
            fail(f"download error: {download_result[0].error}")
        content = download_result[0].content
        if content is None:
            fail("download returned None content")
        decoded = content.decode("utf-8", errors="replace")
        if "sandbox-written-content" not in decoded:
            fail(f"Downloaded content mismatch: {decoded!r}")
        ok(f"download_files: content={decoded.strip()!r}")

        # 1e. Upload multiple files, verify both visible
        print("\n  --- 1e. Upload multiple files ---")
        results = sandbox.upload_files([
            ("/home/baponi/multi-a.txt", b"file-a-content"),
            ("/home/baponi/sub/multi-b.txt", b"file-b-content"),
        ])
        errors = [r for r in results if r.error]
        if errors:
            fail(f"multi-upload errors: {[(r.path, r.error) for r in errors]}")
        ok(f"upload 2 files: no errors")

        result = sandbox.execute("cat /home/baponi/multi-a.txt && cat /home/baponi/sub/multi-b.txt")
        if "file-a-content" not in result.output or "file-b-content" not in result.output:
            fail(f"multi-file content mismatch: {result.output!r}")
        ok(f"Both uploaded files visible in sandbox")

        # 1f. Per-call env_vars override
        print("\n  --- 1f. Per-call env_vars override ---")
        result = sandbox.execute("echo $TEST_VAR", env_vars={"TEST_VAR": "hello_from_override"})
        if result.exit_code != 0:
            fail(f"per-call env_vars failed: exit_code={result.exit_code}")
        if "hello_from_override" not in result.output:
            fail(f"per-call env_vars not visible: {result.output!r}")
        ok(f"Per-call env_vars override works: output={result.output.strip()!r}")

        print(f"\n  All direct SDK tests passed.")


# ── Test 2: Real Deep Agent with Claude Sonnet 4.6 ─────────────────────────

def _extract_tool_calls(event: dict) -> tuple[set[str], str | None]:
    """Extract tool call names and final text from a stream event."""
    tools = set()
    text = None
    for key, val in event.items():
        if not isinstance(val, dict):
            continue
        msgs = val.get("messages")
        if not isinstance(msgs, list):
            continue
        for msg in msgs:
            content = getattr(msg, "content", None)
            # Tool calls are in content as list of dicts with type='tool_use'
            if isinstance(content, list):
                for item in content:
                    if isinstance(item, dict) and item.get("type") == "tool_use":
                        tools.add(item["name"])
            # Final text response
            elif isinstance(content, str) and content.strip():
                text = content
    return tools, text


def _run_agent_prompt(agent, prompt: str) -> tuple[set[str], str]:
    """Stream agent with prompt, return (all_tools_used, final_text)."""
    all_tools = set()
    final_text = ""
    for event in agent.stream({"messages": [("human", prompt)]}):
        tools, text = _extract_tool_calls(event)
        for t in tools:
            print(f"    -> Tool call: {t}")
        all_tools.update(tools)
        if text:
            final_text = text
    return all_tools, final_text


def test_deep_agent():
    section("Test 2: Real Deep Agent - Claude Sonnet 4.6")

    from baponi.deepagents import BaponiSandbox
    from deepagents import create_deep_agent

    agent_thread = f"e2e-agent-{uuid.uuid4().hex[:12]}"

    sandbox = BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=agent_thread,
        timeout=25,
    )
    print(f"  Agent thread ID: {sandbox.id}")

    agent = create_deep_agent(
        "anthropic:claude-sonnet-4-6",
        backend=sandbox,
        name="e2e-test-agent",
    )

    # 2a. Ask agent to write a file and read it back
    print("\n  --- 2a. Agent writes and reads a Python file ---")
    tools, text = _run_agent_prompt(
        agent,
        "Write a Python file at /home/baponi/fibonacci.py that defines a function "
        "fib(n) returning the nth fibonacci number. Then execute it with python3 to "
        "compute fib(10) and tell me the result.",
    )
    print(f"  Tools used: {tools}")
    print(f"  Agent response: {text[:200]}")
    if "write_file" not in tools:
        fail("Agent never called write_file")
    if "execute" not in tools:
        fail("Agent never called execute")
    ok("Agent used write_file + execute to solve fibonacci task")

    # 2b. Ask agent to explore files it created
    print("\n  --- 2b. Agent lists and reads files ---")
    tools2, text2 = _run_agent_prompt(
        agent,
        "List the files in /home/baponi/ and read the contents of fibonacci.py. "
        "Tell me what's in it.",
    )
    print(f"  Tools used: {tools2}")
    print(f"  Agent response: {text2[:300]}")
    if "ls" not in tools2 and "glob" not in tools2:
        fail("Agent never listed files")
    if "read_file" not in tools2:
        fail("Agent never read the file")
    ok("Agent used ls/glob + read_file to explore sandbox")

    # 2c. Ask agent to install a package and use it
    print("\n  --- 2c. Agent installs package and runs code ---")
    tools3, text3 = _run_agent_prompt(
        agent,
        "Install the 'cowsay' pip package and then write and run a Python script that "
        "uses cowsay to print 'Baponi rules'. Show me the output.",
    )
    print(f"  Tools used: {tools3}")
    print(f"  Agent response: {text3[:300]}")
    if "execute" not in tools3:
        fail("Agent never called execute (expected pip install + python run)")
    ok("Agent installed package and ran code in sandbox")

    print(f"\n  All deep agent tests passed.")


# ── Main ────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("BaponiSandbox Deep Agents E2E Test Suite")
    print(f"Staging: {BAPONI_BASE_URL}")
    print(f"Thread: {THREAD_ID}")

    test_direct_sdk()
    test_deep_agent()

    section("ALL E2E TESTS PASSED")
