"""Real E2E test: Claude Sonnet 4 agent using Baponi SDK on Kind cluster.

This test verifies the full stack:
1. Baponi Python SDK → HTTP → Kind gateway → executor → nsjail sandbox
2. Claude Sonnet 4 generates Python code and executes it via Baponi
3. Results flow back through the SDK to the agent

Requires:
- Kind cluster running with port-forward on localhost:8000
- BAPONI_API_KEY set (API key for the Kind cluster)
- ANTHROPIC_API_KEY set (for Claude API calls)
"""

from __future__ import annotations

import json
import os
import sys
import time

import anthropic

from baponi import Baponi
from baponi.anthropic import code_sandbox_tool, create_code_sandbox

BAPONI_API_KEY = os.environ["BAPONI_API_KEY"]
ANTHROPIC_API_KEY = os.environ["ANTHROPIC_API_KEY"]
BAPONI_BASE_URL = os.environ.get("BAPONI_BASE_URL", "http://localhost:8000")

# Create a handler configured for our Kind cluster
_tool_def, handle_tool_call = create_code_sandbox(
    api_key=BAPONI_API_KEY,
    base_url=BAPONI_BASE_URL,
)


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


def test_direct_sdk() -> None:
    """Test 1: Direct SDK usage — no LLM, just verify the SDK works end-to-end."""
    separator("Test 1: Direct SDK Usage")

    client = Baponi(api_key=BAPONI_API_KEY, base_url=BAPONI_BASE_URL)

    # Simple Python
    result = client.execute("print('Hello from Baponi SDK!')")
    assert result.success, f"Expected success, got: {result}"
    assert "Hello from Baponi SDK!" in result.stdout
    print(f"  [PASS] Simple Python: stdout={result.stdout.strip()!r}")

    # Multi-line with computation
    result = client.execute("""
import math
print(f"pi = {math.pi:.10f}")
print(f"e  = {math.e:.10f}")
print(f"sqrt(2) = {math.sqrt(2):.10f}")
""")
    assert result.success
    assert "3.141592653" in result.stdout
    print(f"  [PASS] Math computation: {result.duration_ms}ms")

    # Bash execution
    result = client.execute("echo 'Hello from Bash' && uname -s", language="bash")
    assert result.success
    assert "Hello from Bash" in result.stdout
    print(f"  [PASS] Bash execution: stdout={result.stdout.strip()!r}")

    # Node.js execution
    result = client.execute(
        "console.log(JSON.stringify({status: 'ok', pi: Math.PI}))",
        language="node",
    )
    assert result.success
    data = json.loads(result.stdout.strip())
    assert data["status"] == "ok"
    print(f"  [PASS] Node.js execution: {data}")

    # Error handling — non-zero exit code returns result, doesn't raise
    result = client.execute("raise ValueError('intentional error')")
    assert not result.success
    assert result.exit_code != 0
    assert "ValueError" in result.stderr
    print(f"  [PASS] Error handling: exit_code={result.exit_code}")

    # Thread persistence
    thread_id = f"sdk-test-{int(time.time())}"
    client.execute(
        "with open('/home/safesandy/test.txt', 'w') as f: f.write('persisted!')",
        thread_id=thread_id,
    )
    result = client.execute(
        "with open('/home/safesandy/test.txt') as f: print(f.read())",
        thread_id=thread_id,
    )
    assert result.success
    assert "persisted!" in result.stdout
    print(f"  [PASS] Thread persistence: thread_id={thread_id}")

    # Metadata
    result = client.execute(
        "print('metadata test')",
        metadata={"test_run": "e2e", "env": "kind"},
    )
    assert result.success
    print(f"  [PASS] Metadata: accepted without error")

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


def test_claude_simple_agent() -> None:
    """Test 2: Claude generates and executes code via tool use."""
    separator("Test 2: Claude Simple Agent (Single Tool Call)")

    client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)

    messages = [
        {
            "role": "user",
            "content": (
                "Use the code_sandbox tool to write a Python program that calculates "
                "the first 20 Fibonacci numbers and prints them as a comma-separated list."
            ),
        }
    ]

    print("  Sending request to Claude Sonnet 4...")
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        tools=[code_sandbox_tool],
        messages=messages,
    )

    print(f"  Response stop_reason: {response.stop_reason}")
    assert response.stop_reason == "tool_use", f"Expected tool_use, got {response.stop_reason}"

    # Find the tool use block
    tool_block = None
    for block in response.content:
        if block.type == "tool_use":
            tool_block = block
            break

    assert tool_block is not None, "No tool_use block in response"
    assert tool_block.name == "code_sandbox"
    print(f"  Claude's code:\n    {tool_block.input['code'][:200]}...")

    # Execute the tool call via our SDK
    result = handle_tool_call(tool_block.name, tool_block.input)
    print(f"  Execution result: {result[:200]}")

    # Verify fibonacci numbers are in the output (Claude may start from 0 or 1)
    assert "55" in result and "89" in result and "144" in result, (
        f"Expected fibonacci numbers (55, 89, 144) in result: {result}"
    )
    print(f"  [PASS] Claude generated correct fibonacci code!")


def test_claude_multi_turn_agent() -> None:
    """Test 3: Multi-turn conversation with Claude using sandbox persistently."""
    separator("Test 3: Claude Multi-Turn Agent (Persistent Thread)")

    client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)
    thread_id = f"claude-agent-{int(time.time())}"

    # Override the handle to use thread persistence
    baponi_client = Baponi(api_key=BAPONI_API_KEY, base_url=BAPONI_BASE_URL)
    from baponi._base import result_to_llm_text

    def handle_with_thread(name: str, input_data: dict) -> str:
        """Handle tool calls with forced thread_id for persistence."""
        if name != "code_sandbox":
            raise ValueError(f"Unknown tool: {name}")
        result = baponi_client.execute(
            input_data["code"],
            language=input_data.get("language", "python"),
            timeout=input_data.get("timeout", 30),
            thread_id=input_data.get("thread_id", thread_id),
            metadata=input_data.get("metadata"),
        )
        return result_to_llm_text(result)

    messages: list[dict] = []

    # Turn 1: Ask Claude to create a data file
    messages.append({
        "role": "user",
        "content": (
            f"Use the code_sandbox tool with thread_id='{thread_id}' to create a CSV file at "
            "/home/safesandy/sales.csv with this data:\n"
            "month,revenue,expenses\n"
            "Jan,10000,8000\n"
            "Feb,12000,7500\n"
            "Mar,15000,9000\n"
            "Apr,11000,8500\n"
            "May,18000,10000\n"
            "Jun,20000,11000\n"
            "Then print 'File created successfully'."
        ),
    })

    print("  Turn 1: Asking Claude to create CSV data...")
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        tools=[code_sandbox_tool],
        messages=messages,
    )

    # Process tool calls
    tool_results = []
    for block in response.content:
        if block.type == "tool_use":
            result = handle_with_thread(block.name, block.input)
            tool_results.append({
                "type": "tool_result",
                "tool_use_id": block.id,
                "content": result,
            })
            print(f"  Turn 1 result: {result[:100]}")

    assert "File created successfully" in tool_results[0]["content"] or "created" in tool_results[0]["content"].lower()
    print(f"  [PASS] Turn 1: CSV file created")

    # Add assistant response + tool results to conversation
    messages.append({"role": "assistant", "content": response.content})
    messages.append({"role": "user", "content": tool_results})

    # Turn 2: Ask Claude to analyze the data from the SAME thread
    messages.append({
        "role": "user",
        "content": (
            f"Now use the code_sandbox tool with thread_id='{thread_id}' to read the CSV file "
            "you just created at /home/safesandy/sales.csv, calculate the total revenue, "
            "total expenses, and total profit. Print each on a separate line."
        ),
    })

    print("  Turn 2: Asking Claude to analyze the data...")
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        tools=[code_sandbox_tool],
        messages=messages,
    )

    for block in response.content:
        if block.type == "tool_use":
            result = handle_with_thread(block.name, block.input)
            print(f"  Turn 2 result: {result[:200]}")

            # Verify the calculations
            # Total revenue: 10000+12000+15000+11000+18000+20000 = 86000
            # Total expenses: 8000+7500+9000+8500+10000+11000 = 54000
            # Profit: 32000
            assert "86000" in result, f"Expected total revenue 86000 in: {result}"
            assert "54000" in result or "32000" in result, f"Expected expenses/profit in: {result}"
            print(f"  [PASS] Turn 2: Data analysis correct (revenue=86000)")

    baponi_client.close()


def test_claude_complex_code_generation() -> None:
    """Test 4: Claude generates complex code with imports and data processing."""
    separator("Test 4: Claude Complex Code Generation")

    client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)

    messages = [
        {
            "role": "user",
            "content": (
                "Use the code_sandbox tool to write a Python program that:\n"
                "1. Generates 100 random numbers between 1 and 1000 using random.seed(42)\n"
                "2. Calculates mean, median, standard deviation, min, max\n"
                "3. Finds all prime numbers in the list\n"
                "4. Prints a JSON object with keys: mean, median, std, min, max, prime_count, primes\n"
                "Use only the standard library (math, random, json, statistics)."
            ),
        }
    ]

    print("  Sending complex task to Claude...")
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=2048,
        tools=[code_sandbox_tool],
        messages=messages,
    )

    for block in response.content:
        if block.type == "tool_use":
            print(f"  Code length: {len(block.input['code'])} chars")
            result = handle_tool_call(block.name, block.input)
            print(f"  Raw result: {result[:300]}")

            # Parse the JSON output — Claude may use json.dumps(indent=2)
            # so we need to find the full JSON object, not just a single line
            text = result.strip()
            # Find the JSON object boundary
            start_idx = text.find("{")
            assert start_idx != -1, f"No JSON found in output: {result}"
            # Find matching closing brace
            brace_count = 0
            end_idx = start_idx
            for i in range(start_idx, len(text)):
                if text[i] == "{":
                    brace_count += 1
                elif text[i] == "}":
                    brace_count -= 1
                    if brace_count == 0:
                        end_idx = i + 1
                        break
            json_str = text[start_idx:end_idx]
            data = json.loads(json_str)

            assert "mean" in data, f"Missing 'mean' in output: {data}"
            assert "median" in data, f"Missing 'median' in output: {data}"
            assert "prime_count" in data, f"Missing 'prime_count' in output: {data}"
            assert isinstance(data["primes"], list), f"primes should be a list: {data}"
            assert data["prime_count"] > 0, f"Should have found some primes"

            print(f"  Stats: mean={data['mean']:.1f}, median={data['median']:.1f}")
            print(f"  Found {data['prime_count']} primes")
            print(f"  [PASS] Complex code generation worked!")


def test_claude_iterative_debugging() -> None:
    """Test 5: Claude handles execution errors and fixes its code."""
    separator("Test 5: Claude Iterative Debugging")

    client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)

    baponi_client = Baponi(api_key=BAPONI_API_KEY, base_url=BAPONI_BASE_URL)
    from baponi._base import result_to_llm_text

    messages: list[dict] = [
        {
            "role": "user",
            "content": (
                "Use the code_sandbox tool to write a Python program that sorts a list "
                "using a custom merge sort implementation. The list is: [38, 27, 43, 3, 9, 82, 10]. "
                "Print the sorted result."
            ),
        }
    ]

    print("  Turn 1: Asking Claude for merge sort implementation...")
    max_turns = 3
    final_result = None

    for turn in range(max_turns):
        response = client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=2048,
            tools=[code_sandbox_tool],
            messages=messages,
        )

        if response.stop_reason == "end_turn":
            print(f"  Claude finished (text response) at turn {turn + 1}")
            # Extract text
            for block in response.content:
                if hasattr(block, "text"):
                    print(f"  Claude says: {block.text[:200]}")
            break

        # Process tool calls
        tool_results = []
        for block in response.content:
            if block.type == "tool_use":
                result_text = handle_tool_call(block.name, block.input)
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": result_text,
                })
                print(f"  Turn {turn + 1} execution: {result_text[:150]}")

                # Check if we got the sorted output
                if "[3, 9, 10, 27, 38, 43, 82]" in result_text:
                    final_result = result_text
                    print(f"  [PASS] Correct sorted result at turn {turn + 1}!")

        messages.append({"role": "assistant", "content": response.content})
        messages.append({"role": "user", "content": tool_results})

        if final_result:
            break

    assert final_result is not None, "Claude never produced the correct sorted result"
    baponi_client.close()


def test_claude_node_execution() -> None:
    """Test 6: Claude uses Node.js in the sandbox."""
    separator("Test 6: Claude Node.js Execution")

    client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)

    messages = [
        {
            "role": "user",
            "content": (
                "Use the code_sandbox tool with language='node' to write a Node.js program that:\n"
                "1. Creates an array of objects representing users with name, age, and score fields\n"
                "2. Filters users with score > 80\n"
                "3. Sorts them by age descending\n"
                "4. Maps to a string like 'Name (age: X, score: Y)'\n"
                "5. Prints the result as JSON array\n"
                "Use at least 6 users."
            ),
        }
    ]

    print("  Asking Claude for Node.js code...")
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        tools=[code_sandbox_tool],
        messages=messages,
    )

    for block in response.content:
        if block.type == "tool_use":
            lang = block.input.get("language", "python")
            print(f"  Language: {lang}")
            result = handle_tool_call(block.name, block.input)
            print(f"  Result: {result[:300]}")

            # Claude may print extra text before/after the result
            # Just verify the execution worked and contains user data
            assert result and len(result) > 10, f"Expected substantial output: {result}"
            # Check for user-related content in the output
            result_lower = result.lower()
            assert any(w in result_lower for w in ["age", "score", "name", "alice", "bob"]), (
                f"Expected user data in output: {result}"
            )
            print(f"  [PASS] Node.js executed successfully with user data")


def main() -> None:
    """Run all E2E tests sequentially."""
    print("\n" + "=" * 70)
    print("  Baponi Python SDK — Real E2E Tests with Claude Sonnet 4")
    print("=" * 70)
    print(f"\n  Baponi endpoint: {BAPONI_BASE_URL}")
    print(f"  Anthropic model: claude-sonnet-4-20250514")

    start = time.time()
    tests = [
        ("Direct SDK", test_direct_sdk),
        ("Claude Simple Agent", test_claude_simple_agent),
        ("Claude Multi-Turn Agent", test_claude_multi_turn_agent),
        ("Claude Complex Code Generation", test_claude_complex_code_generation),
        ("Claude Iterative Debugging", test_claude_iterative_debugging),
        ("Claude Node.js Execution", test_claude_node_execution),
    ]

    passed = 0
    failed = 0
    errors: list[tuple[str, str]] = []

    for name, test_fn in tests:
        try:
            test_fn()
            passed += 1
        except Exception as e:
            failed += 1
            errors.append((name, str(e)))
            print(f"\n  [FAIL] {name}: {e}")

    elapsed = time.time() - start

    separator("Results")
    print(f"  Passed: {passed}/{len(tests)}")
    print(f"  Failed: {failed}/{len(tests)}")
    print(f"  Time:   {elapsed:.1f}s")

    if errors:
        print(f"\n  Failures:")
        for name, err in errors:
            print(f"    - {name}: {err}")

    print()

    if failed > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()
