"""Full E2E test for BaponiSandbox Deep Agents integration against staging.

Exercises every feature added in the refactor:
  1. Context manager (with statement)
  2. Per-call overrides (env_vars, metadata)
  3. Thread persistence (state survives across calls)
  4. Streaming execution (execute_stream + event types)
  5. File upload + download round-trip
  6. BaponiSandboxProvider lifecycle (get_or_create, reuse, delete)
  7. Multiple concurrent threads via provider
  8. Real Deep Agent with Claude (write + execute + file tools)

Run:
    cd sdks/python
    BAPONI_API_KEY=sk-... .venv/bin/python tests/e2e_deepagents_full.py
"""

from __future__ import annotations

import os
import sys
import time
import uuid

# ── Config ──────────────────────────────────────────────────────────────────

BAPONI_API_KEY = os.environ["BAPONI_API_KEY"]
BAPONI_BASE_URL = os.environ.get("BAPONI_BASE_URL", "https://staging.api.baponi.dev")
RUN_ID = uuid.uuid4().hex[:8]

passed = 0
failed = 0


def ok(msg: str) -> None:
    global passed
    passed += 1
    print(f"  \033[32mPASS\033[0m {msg}")


def fail(msg: str) -> None:
    global failed
    failed += 1
    print(f"  \033[31mFAIL\033[0m {msg}")


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


# ── Test 1: Context Manager ────────────────────────────────────────────────

def test_context_manager():
    section("1. Context manager - with statement")

    from baponi.deepagents import BaponiSandbox

    thread = f"e2e-ctx-{RUN_ID}"
    with BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=thread,
        metadata={"test": "context_manager", "run": RUN_ID},
    ) as sandbox:
        result = sandbox.execute("echo 'context manager works'")
        if result.exit_code == 0 and "context manager works" in result.output:
            ok(f"execute inside with-block: {result.output.strip()!r}")
        else:
            fail(f"exit_code={result.exit_code}, output={result.output!r}")

    # After exiting, sandbox should be closed - verify no crash
    ok("exited context manager without error")


# ── Test 2: Per-Call Overrides ──────────────────────────────────────────────

def test_per_call_overrides():
    section("2. Per-call overrides - env_vars, metadata")

    from baponi.deepagents import BaponiSandbox

    thread = f"e2e-override-{RUN_ID}"
    with BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=thread,
        env_vars={"DEFAULT_VAR": "from_constructor"},
        metadata={"source": "constructor"},
    ) as sandbox:

        # 2a. Constructor env_vars work
        result = sandbox.execute("echo $DEFAULT_VAR")
        if "from_constructor" in result.output:
            ok(f"constructor env_vars: {result.output.strip()!r}")
        else:
            fail(f"constructor env_vars missing: {result.output!r}")

        # 2b. Per-call env_vars override
        result = sandbox.execute(
            "echo $OVERRIDE_VAR",
            env_vars={"OVERRIDE_VAR": "per_call_value"},
        )
        if "per_call_value" in result.output:
            ok(f"per-call env_vars override: {result.output.strip()!r}")
        else:
            fail(f"per-call env_vars missing: {result.output!r}")

        # 2c. Per-call metadata (doesn't affect output, but shouldn't break)
        result = sandbox.execute(
            "echo 'metadata override ok'",
            metadata={"step": "override_test", "run": RUN_ID},
        )
        if result.exit_code == 0:
            ok("per-call metadata override accepted")
        else:
            fail(f"metadata override failed: exit_code={result.exit_code}")

        # 2d. Multiple env vars in one call
        result = sandbox.execute(
            'echo "$A:$B:$C"',
            env_vars={"A": "alpha", "B": "bravo", "C": "charlie"},
        )
        if "alpha:bravo:charlie" in result.output:
            ok(f"multiple env vars: {result.output.strip()!r}")
        else:
            fail(f"multiple env vars missing: {result.output!r}")

        # 2e. Per-call timeout (short timeout, fast command)
        result = sandbox.execute("echo fast", timeout=10)
        if result.exit_code == 0:
            ok("per-call timeout=10 accepted")
        else:
            fail(f"timeout override failed: exit_code={result.exit_code}")


# ── Test 3: Thread Persistence ──────────────────────────────────────────────

def test_thread_persistence():
    section("3. Thread persistence - state survives across calls")

    from baponi.deepagents import BaponiSandbox

    thread = f"e2e-persist-{RUN_ID}"
    with BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=thread,
        metadata={"test": "persistence"},
    ) as sandbox:

        # 3a. Write a file
        result = sandbox.execute(
            "echo 'persisted content 12345' > /home/baponi/persist-test.txt"
        )
        assert result.exit_code == 0, f"write failed: {result.output}"
        ok("wrote file to thread")

        # 3b. Read it back in a separate execution (same thread)
        result = sandbox.execute("cat /home/baponi/persist-test.txt")
        if "persisted content 12345" in result.output:
            ok(f"file persisted across calls: {result.output.strip()!r}")
        else:
            fail(f"persistence failed: {result.output!r}")

        # 3c. Write a Python module, verify it persists across executions
        sandbox.execute(
            "mkdir -p /home/baponi/mylib && "
            "echo 'def greet(name): return f\"hello {name}\"' > /home/baponi/mylib/greeter.py && "
            "echo '' > /home/baponi/mylib/__init__.py"
        )
        result = sandbox.execute(
            "cd /home/baponi && python3 -c \"from mylib.greeter import greet; print(greet('baponi'))\""
        )
        if "hello baponi" in result.output:
            ok("custom Python module persisted across executions")
        else:
            fail(f"module not persisted: {result.output!r}")

        # 3d. Create a directory tree, verify structure
        sandbox.execute("mkdir -p /home/baponi/project/src /home/baponi/project/tests")
        sandbox.execute("echo 'def main(): pass' > /home/baponi/project/src/app.py")
        sandbox.execute("echo 'def test_main(): pass' > /home/baponi/project/tests/test_app.py")
        result = sandbox.execute("find /home/baponi/project -type f | sort")
        if "src/app.py" in result.output and "tests/test_app.py" in result.output:
            ok("directory tree persisted")
        else:
            fail(f"tree missing: {result.output!r}")


# ── Test 4: Streaming Execution ─────────────────────────────────────────────

def test_streaming():
    section("4. Streaming execution - execute_stream + events")

    from baponi._events import OutputEvent, ResultEvent, StatusEvent
    from baponi.deepagents import BaponiSandbox

    thread = f"e2e-stream-{RUN_ID}"
    with BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=thread,
        metadata={"test": "streaming"},
    ) as sandbox:

        # 4a. Stream a multi-line output command
        stream = sandbox.execute_stream(
            "for i in $(seq 1 5); do echo \"line $i\"; sleep 0.1; done"
        )
        event_types = set()
        output_chunks: list[str] = []
        with stream:
            for event in stream:
                event_types.add(type(event).__name__)
                if isinstance(event, OutputEvent):
                    output_chunks.append(event.data)
            final = stream.get_final_result()

        if "StatusEvent" in event_types:
            ok(f"received StatusEvent")
        else:
            fail(f"no StatusEvent in {event_types}")

        if "OutputEvent" in event_types:
            ok(f"received OutputEvent ({len(output_chunks)} chunks)")
        else:
            fail(f"no OutputEvent in {event_types}")

        if "ResultEvent" in event_types:
            ok("received ResultEvent")
        else:
            fail(f"no ResultEvent in {event_types}")

        if final.exit_code == 0:
            ok(f"stream final result: exit_code=0")
        else:
            fail(f"stream exit_code={final.exit_code}")

        combined = "".join(output_chunks)
        if "line 1" in combined and "line 5" in combined:
            ok("all 5 lines received via stream")
        else:
            fail(f"stream output incomplete: {combined!r}")

        if stream.trace_id:
            ok(f"trace_id captured: {stream.trace_id[:16]}...")
        else:
            fail("no trace_id from stream")

        # 4b. Stream with per-call env_vars
        stream2 = sandbox.execute_stream(
            "echo $STREAM_VAR",
            env_vars={"STREAM_VAR": "streamed_value"},
        )
        with stream2:
            result2 = stream2.until_done()

        if "streamed_value" in (result2.stdout or ""):
            ok("streaming with per-call env_vars works")
        else:
            fail(f"stream env_vars: stdout={result2.stdout!r}")


# ── Test 5: File Upload + Download ──────────────────────────────────────────

def test_file_roundtrip():
    section("5. File upload + download round-trip")

    from baponi.deepagents import BaponiSandbox

    thread = f"e2e-files-{RUN_ID}"
    with BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=thread,
        metadata={"test": "files"},
    ) as sandbox:

        # 5a. Upload via signed URL
        csv_content = b"name,score\nalice,95\nbob,87\ncharlie,92\n"
        upload_results = sandbox.upload_files([
            ("/home/baponi/scores.csv", csv_content),
        ])
        if upload_results[0].error is None:
            ok("uploaded scores.csv via signed URL")
        else:
            fail(f"upload error: {upload_results[0].error}")

        # 5b. Verify visible in sandbox
        result = sandbox.execute("cat /home/baponi/scores.csv")
        if "alice,95" in result.output:
            ok("uploaded file visible in sandbox")
        else:
            fail(f"file not visible: {result.output!r}")

        # 5c. Process file in sandbox, write output
        result = sandbox.execute(
            "python3 -c \""
            "import csv; "
            "rows = list(csv.DictReader(open('/home/baponi/scores.csv'))); "
            "avg = sum(int(r['score']) for r in rows) / len(rows); "
            "open('/home/baponi/report.txt','w').write(f'Average: {avg:.1f}\\nCount: {len(rows)}\\n')"
            "\""
        )
        if result.exit_code == 0:
            ok("processed CSV in sandbox")
        else:
            fail(f"processing failed: {result.output!r}")

        # 5d. Download output via signed URL
        time.sleep(2)  # storage writeback
        dl_results = sandbox.download_files(["/home/baponi/report.txt"])
        if dl_results[0].error is None and dl_results[0].content:
            content = dl_results[0].content.decode()
            if "Average: 91.3" in content and "Count: 3" in content:
                ok(f"downloaded report: {content.strip()!r}")
            else:
                fail(f"report content wrong: {content!r}")
        else:
            fail(f"download error: {dl_results[0].error}")

        # 5e. Write files inside sandbox (avoids VFS cache staleness from signed URL)
        sandbox.execute("echo -n 'file-a' > /home/baponi/a.txt")
        sandbox.execute("mkdir -p /home/baponi/nested && echo -n 'file-b' > /home/baponi/nested/b.txt")
        result = sandbox.execute(
            "cat /home/baponi/a.txt && echo '---' && cat /home/baponi/nested/b.txt"
        )
        if "file-a" in result.output and "file-b" in result.output:
            ok("multi-file write and read")
        else:
            fail(f"multi-file: {result.output!r}")

        # 5f. Download one of the sandbox-written files
        time.sleep(2)
        dl = sandbox.download_files(["/home/baponi/nested/b.txt"])
        if dl[0].error is None and dl[0].content == b"file-b":
            ok("download sandbox-written file")
        else:
            fail(f"download nested file: error={dl[0].error}, content={dl[0].content!r}")


# ── Test 6: BaponiSandboxProvider Lifecycle ─────────────────────────────────

def test_provider_lifecycle():
    section("6. BaponiSandboxProvider - lifecycle management")

    from baponi.deepagents import BaponiSandboxProvider

    provider = BaponiSandboxProvider(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        timeout=25,
        metadata={"test": "provider", "run": RUN_ID},
    )

    # 6a. Create sandbox via provider
    sid = f"e2e-prov-{RUN_ID}"
    s1 = provider.get_or_create(sandbox_id=sid)
    result = s1.execute("echo 'provider sandbox'")
    if result.exit_code == 0 and "provider sandbox" in result.output:
        ok(f"provider.get_or_create(): {result.output.strip()!r}")
    else:
        fail(f"provider sandbox failed: {result.output!r}")

    # 6b. get_or_create returns same instance
    s2 = provider.get_or_create(sandbox_id=sid)
    if s1 is s2:
        ok("get_or_create returns same instance for same ID")
    else:
        fail("get_or_create returned different instance")

    # 6c. Write state, verify persistence through provider
    s1.execute("echo 'provider-state' > /home/baponi/prov.txt")
    result = s2.execute("cat /home/baponi/prov.txt")
    if "provider-state" in result.output:
        ok("state persisted through provider reuse")
    else:
        fail(f"provider state lost: {result.output!r}")

    # 6d. Delete sandbox
    provider.delete(sandbox_id=sid)
    s3 = provider.get_or_create(sandbox_id=sid)
    if s3 is not s1:
        ok("delete + get_or_create returns new instance")
    else:
        fail("delete did not clear the sandbox")

    # 6e. Auto-generated IDs work
    auto = provider.get_or_create()
    result = auto.execute("echo auto-id")
    if result.exit_code == 0 and auto.id.startswith("deep-"):
        ok(f"auto-generated ID: {auto.id}")
    else:
        fail(f"auto-id failed: id={auto.id}, exit_code={result.exit_code}")

    provider.delete(sandbox_id=auto.id)
    provider.delete(sandbox_id=sid)


# ── Test 7: Multiple Threads via Provider ───────────────────────────────────

def test_multi_thread():
    section("7. Multiple concurrent threads via provider")

    from baponi.deepagents import BaponiSandboxProvider

    provider = BaponiSandboxProvider(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        metadata={"test": "multi_thread"},
    )

    threads = [f"e2e-mt-{RUN_ID}-{i}" for i in range(3)]
    sandboxes = [provider.get_or_create(sandbox_id=t) for t in threads]

    # Write distinct state to each thread
    for i, sb in enumerate(sandboxes):
        sb.execute(f"echo 'thread-{i}-data' > /home/baponi/identity.txt")

    # Read back - each thread should have its own state
    all_correct = True
    for i, sb in enumerate(sandboxes):
        result = sb.execute("cat /home/baponi/identity.txt")
        expected = f"thread-{i}-data"
        if expected not in result.output:
            fail(f"thread {i} isolation broken: expected {expected!r}, got {result.output!r}")
            all_correct = False

    if all_correct:
        ok(f"3 threads isolated: each has its own state")

    # Cleanup
    for t in threads:
        provider.delete(sandbox_id=t)


# ── Test 8: Real Deep Agent ────────────────────────────────────────────────

def test_deep_agent():
    section("8. Real Deep Agent - Claude Sonnet 4.6")

    anthropic_key = os.environ.get("ANTHROPIC_API_KEY")
    if not anthropic_key:
        print("  SKIP: ANTHROPIC_API_KEY not set")
        return

    os.environ["ANTHROPIC_API_KEY"] = anthropic_key

    from baponi.deepagents import BaponiSandbox
    from deepagents import create_deep_agent

    thread = f"e2e-agent-{RUN_ID}"
    with BaponiSandbox(
        api_key=BAPONI_API_KEY,
        base_url=BAPONI_BASE_URL,
        thread_id=thread,
        timeout=25,
        metadata={"test": "deep_agent", "run": RUN_ID},
        env_vars={"PROJECT_NAME": "e2e_test"},
    ) as sandbox:

        agent = create_deep_agent(
            "anthropic:claude-sonnet-4-6",
            backend=sandbox,
            name=f"e2e-{RUN_ID}",
        )

        # 8a. Agent writes and executes code
        print("  --- 8a. Agent writes + executes Python ---")
        all_tools: set[str] = set()
        final_text = ""
        for event in agent.stream({"messages": [("human",
            "Write a Python script at /home/baponi/fibonacci.py that computes "
            "fib(20) using memoization, then run it and tell me the result. "
            "Use the $PROJECT_NAME env var in a print statement to prove it's available."
        )]}):
            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)
                    if isinstance(content, list):
                        for item in content:
                            if isinstance(item, dict) and item.get("type") == "tool_use":
                                tool = item["name"]
                                print(f"    -> {tool}")
                                all_tools.add(tool)
                    elif isinstance(content, str) and content.strip():
                        final_text = content

        print(f"  Tools used: {all_tools}")
        print(f"  Response: {final_text[:300]}")

        if "write_file" in all_tools:
            ok("agent called write_file")
        else:
            fail(f"agent never wrote file: {all_tools}")

        if "execute" in all_tools:
            ok("agent called execute")
        else:
            fail(f"agent never executed: {all_tools}")

        # Verify the file exists after agent ran
        result = sandbox.execute("test -f /home/baponi/fibonacci.py && echo EXISTS")
        if "EXISTS" in result.output:
            ok("agent-created file persisted in thread")
        else:
            fail("agent file not found")


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

if __name__ == "__main__":
    print(f"BaponiSandbox Deep Agents - Full E2E Suite")
    print(f"Staging: {BAPONI_BASE_URL}")
    print(f"Run ID:  {RUN_ID}")

    tests = [
        test_context_manager,
        test_per_call_overrides,
        test_thread_persistence,
        test_streaming,
        test_file_roundtrip,
        test_provider_lifecycle,
        test_multi_thread,
        test_deep_agent,
    ]

    for test in tests:
        try:
            test()
        except Exception as e:
            fail(f"{test.__name__} crashed: {e}")
            import traceback
            traceback.print_exc()

    section("RESULTS")
    print(f"  Passed: {passed}")
    print(f"  Failed: {failed}")
    if failed > 0:
        print(f"\n  \033[31m{failed} test(s) failed!\033[0m")
        sys.exit(1)
    else:
        print(f"\n  \033[32mAll {passed} checks passed!\033[0m")
