"""LangChain + Amazon Bedrock — syntax-checkable agent sketch.

Versions / setup (pin in your project; this file does not install packages):
  - Python 3.12+
  - langchain-aws ChatBedrockConverse (Bedrock Converse API)
  - langchain.agents.create_agent (model + tools + system_prompt)
  - Optional host: bedrock_agentcore.runtime.BedrockAgentCoreApp
  - Always set max_tokens explicitly

This stub does NOT call AWS. Live imports sit behind DRY_RUN so
`python3 -m py_compile` succeeds without those packages installed.
"""

from __future__ import annotations

import os
from typing import Any

DRY_RUN = os.environ.get("DRY_RUN", "1") == "1"
REGION = os.environ.get("AWS_REGION", "us-west-2")
MODEL_ID = os.environ.get("BEDROCK_MODEL_ID", "us.anthropic.claude-sonnet-4-6")


def lookup_order_status(order_id: str) -> str:
    """Narrow read tool — not a database handle.

    In production this should call a Gateway-backed API with IAM + Policy,
    not a warehouse connection string the model can see.
    """
    if not order_id or len(order_id) > 32:
        return "invalid order_id"
    if DRY_RUN:
        return f"dry-run: order {order_id} is IN_TRANSIT"
    return f"order {order_id} status unavailable in stub"


def build_model() -> Any:
    from langchain_aws import ChatBedrockConverse

    return ChatBedrockConverse(
        model=MODEL_ID,
        region_name=REGION,
        max_tokens=1024,
        temperature=0,
    )


def build_agent() -> Any:
    from langchain.agents import create_agent

    return create_agent(
        model=build_model(),
        tools=[lookup_order_status],
        system_prompt=(
            "You are a support assistant. Use lookup_order_status for order questions. "
            "Never invent shipment facts. Escalate refunds instead of calling write tools."
        ),
    )


def agent_invocation(payload: dict[str, Any], _context: Any = None) -> dict[str, str]:
    """Shape used as an AgentCore Runtime entrypoint when hosted on Runtime."""
    prompt = str(payload.get("prompt", "")).strip() or "No prompt found in input."
    if DRY_RUN:
        preview = lookup_order_status(str(payload.get("order_id", "ORD-0")))
        return {
            "result": "dry-run",
            "tool_preview": preview,
            "note": "Set DRY_RUN=0 with langchain-aws + langchain installed to invoke Bedrock",
        }

    from bedrock_agentcore.runtime import BedrockAgentCoreApp  # noqa: F401

    agent = build_agent()
    result = agent.invoke({"messages": [{"role": "user", "content": prompt}]})
    messages = result.get("messages", [])
    last = messages[-1] if messages else {"content": str(result)}
    content = last.get("content") if isinstance(last, dict) else str(last)
    return {"result": str(content)}


if __name__ == "__main__":
    print(
        agent_invocation(
            {
                "prompt": "Where is order ORD-1842?",
                "order_id": "ORD-1842",
            }
        )
    )