"""
Cloud agent (steps 3 + 7)  --  runs on Amazon Bedrock AgentCore Runtime.

It receives an ALREADY de-identified prompt from Fusion (contains a token like
CUST_000123, never a name), reasons with a Bedrock model, calls the RAG tool
through AgentCore Gateway (MCP), and returns de-identified talking points.
Fusion restores the real identity on the way back out -- not this agent.

Framework: Strands.  Deploy target: AgentCore Runtime (see scripts/agentcore_setup.sh).
Pin versions in requirements.txt; SDK surfaces move quickly.

Auth to the Gateway: our Gateway uses AWS_IAM, so every MCP request is SigV4-signed
with the Runtime's execution-role credentials (service `bedrock-agentcore`). If a
GATEWAY_TOKEN is provided instead, we fall back to bearer auth (CUSTOM_JWT gateways).
"""
import os

import httpx
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.session import Session
from bedrock_agentcore.runtime import BedrockAgentCoreApp
from strands import Agent
from strands.models import BedrockModel
from strands.tools.mcp import MCPClient
from mcp.client.streamable_http import streamablehttp_client

app = BedrockAgentCoreApp()

GATEWAY_URL = os.environ["AGENTCORE_GATEWAY_URL"]     # set by agentcore_setup.sh
GATEWAY_TOKEN = os.environ.get("AGENTCORE_GATEWAY_TOKEN", "")
REGION = os.environ.get("AWS_REGION", "ap-southeast-1")
# In non-US regions Claude 3.5 Sonnet v2 is INFERENCE_PROFILE-only, so default to
# the APAC cross-region profile. Override with BEDROCK_MODEL_ID at launch.
MODEL_ID = os.environ.get("BEDROCK_MODEL_ID", "apac.anthropic.claude-3-5-sonnet-20241022-v2:0")

SYSTEM_PROMPT = (
    "You are a financial-advisor assistant. You will be given a customer reference "
    "that is an opaque TOKEN (e.g. CUST_000123). Treat it as an opaque identifier: "
    "never invent a name, and always pass the token verbatim to tools. Use the "
    "get_customer_activity_summary tool to fetch a de-identified evidence package, "
    "then produce concise, numbered visit talking points based only on that evidence. "
    "Do not include the raw token in your final talking points."
)


class _SigV4Auth(httpx.Auth):
    """SigV4-sign each MCP request with the Runtime's execution-role credentials."""

    requires_request_body = True

    def __init__(self, service, region):
        self._creds = Session().get_credentials()
        self._service = service
        self._region = region

    def auth_flow(self, request):
        aws_req = AWSRequest(
            method=request.method,
            url=str(request.url),
            data=request.content,
            headers=dict(request.headers),
        )
        SigV4Auth(self._creds, self._service, self._region).add_auth(aws_req)
        request.headers.update(dict(aws_req.headers))
        yield request


def _mcp_client():
    if GATEWAY_TOKEN:
        headers = {"Authorization": f"Bearer {GATEWAY_TOKEN}"}
        return MCPClient(lambda: streamablehttp_client(GATEWAY_URL, headers=headers))
    auth = _SigV4Auth("bedrock-agentcore", REGION)
    return MCPClient(lambda: streamablehttp_client(GATEWAY_URL, auth=auth))


@app.entrypoint
def invoke(payload):
    """payload == {'prompt': '<de-identified prompt from Fusion>'}"""
    prompt = payload.get("prompt", "")
    with _mcp_client() as client:
        agent = Agent(
            model=BedrockModel(model_id=MODEL_ID),
            system_prompt=SYSTEM_PROMPT,
            tools=client.list_tools_sync(),
        )
        result = agent(prompt)
    return {"result": str(result)}


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