"""Agent that tests actor_id isolation with real Strands Agent and SessionManager.

This agent:
1. Uses AgentCoreMemorySessionManager with set_actor_id_for_request()
2. Creates a Strands Agent with the session manager
3. Demonstrates concurrent request isolation
"""
import logging
import asyncio
import os

from bedrock_agentcore.runtime import BedrockAgentCoreApp
from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig
from bedrock_agentcore.memory.integrations.strands.session_manager import AgentCoreMemorySessionManager
from strands import Agent

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = BedrockAgentCoreApp()

# Memory config
MEMORY_ID = os.environ.get("MEMORY_ID", "CustomerSupport-pkP616GF9D")
REGION = os.environ.get("AWS_REGION", "us-west-2")

# Create session manager (shared across requests)
session_config = AgentCoreMemoryConfig(
    memory_id=MEMORY_ID,
    session_id="actor-isolation-test",
    actor_id="default-actor",
)
session_manager = AgentCoreMemorySessionManager(
    agentcore_memory_config=session_config,
    region_name=REGION,
)

# Create Strands Agent with the session manager
agent = Agent(
    system_prompt="You are a helpful assistant. Keep responses to 1 sentence.",
    session_manager=session_manager,
)


@app.entrypoint
async def invoke(payload, context):
    request_id = payload.get("request_id", "unknown")
    actor_id = payload.get("actor_id", "default")
    delay = payload.get("delay", 0)
    prompt = payload.get("prompt", "Say hello")

    logger.info(f"[{request_id}] START - actor_id={actor_id}, delay={delay}")

    # Set actor_id for this request
    token = session_manager.set_actor_id_for_request(actor_id)
    try:
        start_actor = session_manager.actor_id
        logger.info(f"[{request_id}] actor_id at start: {start_actor}")

        if delay > 0:
            logger.info(f"[{request_id}] Sleeping {delay}s...")
            await asyncio.sleep(delay)

        mid_actor = session_manager.actor_id
        logger.info(f"[{request_id}] actor_id after delay: {mid_actor}")

        # Invoke the Strands agent
        response = agent(prompt)
        response_text = response.message["content"][0]["text"] if response else "No response"

        end_actor = session_manager.actor_id
        logger.info(f"[{request_id}] actor_id at end: {end_actor}")

        isolated = start_actor == actor_id and mid_actor == actor_id and end_actor == actor_id
        logger.info(f"[{request_id}] END - isolated={isolated}")

        return {
            "request_id": request_id,
            "actor_id_set": actor_id,
            "actor_id_start": start_actor,
            "actor_id_mid": mid_actor,
            "actor_id_end": end_actor,
            "isolated": isolated,
            "response": response_text[:100],
        }
    finally:
        session_manager.reset_actor_id(token)


if __name__ == "__main__":
    app.run()

"""Test concurrent requests with different actor_ids against deployed AgentCore Runtime"""
import json
import time
import uuid
import boto3
from concurrent.futures import ThreadPoolExecutor

AGENT_ARN = "arn:aws:bedrock-agentcore:us-west-2:332634280070:runtime/actor_container-5BmRCO72R1"
REGION = "us-west-2"

def make_session_id():
    return f"actor-test-{uuid.uuid4()}"

def invoke(request_id: int, actor_id: str, delay: float, session_id: str):
    client = boto3.client("bedrock-agentcore", region_name=REGION)
    payload = {"request_id": request_id, "actor_id": actor_id, "delay": delay, "prompt": "Say hi"}
    
    start = time.time()
    print(f"[{request_id}] SEND actor_id={actor_id}, delay={delay}s")
    
    try:
        response = client.invoke_agent_runtime(
            agentRuntimeArn=AGENT_ARN,
            runtimeSessionId=session_id,
            payload=json.dumps(payload).encode(),
        )
        # Response is in 'response' key, not 'body'
        body = response.get("response")
        raw = body.read().decode("utf-8") if body else ""
        result = json.loads(raw) if raw else {}
        end = time.time()
        
        isolated = result.get("isolated", False)
        status = "✓" if isolated else "✗"
        print(f"[{request_id}] RECV in {end-start:.2f}s - isolated={isolated} {status}")
        return result
    except Exception as e:
        print(f"[{request_id}] ERROR: {e}")
        return {"error": str(e)}

if __name__ == "__main__":
    # Use a fixed session ID so both requests go to the SAME container
    session_id = "actor-test-fixed-session-12345678901234567890"
    print(f"Session: {session_id}")
    print(f"Agent: {AGENT_ARN}")
    print("\nTest: Request 1 (actor=user-A, delay=3s) vs Request 2 (actor=user-B, delay=0s)\n")
    
    # First, warm up the session with a quick request
    print("Warming up session...")
    warmup = invoke(0, "warmup", 0, session_id)
    if "error" in warmup:
        print("Warmup failed, waiting and retrying...")
        import time as t
        t.sleep(10)
        warmup = invoke(0, "warmup", 0, session_id)
    print(f"Warmup result: {warmup}\n")
    
    print("Now running concurrent test...")
    with ThreadPoolExecutor(max_workers=2) as executor:
        f1 = executor.submit(invoke, 1, "user-A", 3.0, session_id)
        time.sleep(0.5)
        f2 = executor.submit(invoke, 2, "user-B", 0, session_id)
        
        r1 = f1.result()
        r2 = f2.result()
    
    print("\n--- Results ---")
    if "error" not in r1 and "error" not in r2:
        print(f"Request 1: {r1}")
        print(f"Request 2: {r2}")
        if r1.get("isolated") and r2.get("isolated"):
            print("\n✓ ACTOR_ID ISOLATION WORKS!")
        else:
            print("\n✗ ACTOR_ID ISOLATION FAILED!")
    else:
        print("Test failed due to errors")