"""LangGraph on AgentCore Runtime — syntax-checkable entrypoint sketch.

Versions / setup (pin in your project; this file does not install packages):
  - Python 3.12+
  - Official pattern: AWS AgentCore "Use any agent framework" LangGraph sample
    https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/using-any-agent-framework.html
  - init_chat_model(..., model_provider="bedrock_converse")
  - BedrockAgentCoreApp @entrypoint wrapping graph.invoke
  - ARM64 Runtime container in a supported AgentCore region

This stub does NOT call AWS. Imports of langgraph / bedrock_agentcore are
inside the live path so `python3 -m py_compile` succeeds without those
packages installed. Set DRY_RUN=0 only in an environment that has them.
"""

from __future__ import annotations

import os
from typing import Annotated, Any, TypedDict

DRY_RUN = os.environ.get("DRY_RUN", "1") == "1"
REGION = os.environ.get("AWS_REGION", "us-west-2")
# Swap to a model ID / inference profile your account allows.
MODEL_ID = os.environ.get("BEDROCK_MODEL_ID", "us.anthropic.claude-sonnet-4-6")


class GraphState(TypedDict):
    """Application workflow state — not the raw LLM context window."""

    messages: list[dict[str, Any]]
    ticket_id: str
    needs_approval: bool


def classify_node(state: GraphState) -> dict[str, Any]:
    """Deterministic stub: high-impact intents require HITL."""
    last = str(state["messages"][-1].get("content", "")).lower()
    needs_approval = any(token in last for token in ("refund", "delete", "production"))
    return {"needs_approval": needs_approval}


def model_node(state: GraphState) -> dict[str, Any]:
    """Live path calls Bedrock Converse through LangChain. DRY_RUN echoes state."""
    if DRY_RUN:
        return {
            "messages": state["messages"]
            + [{"role": "assistant", "content": "dry-run: classified, no model call"}]
        }

    from langchain.chat_models import init_chat_model
    from langgraph.graph.message import add_messages  # noqa: F401 — used by live graphs

    llm = init_chat_model(
        MODEL_ID,
        model_provider="bedrock_converse",
        max_tokens=1024,
        region_name=REGION,
    )
    reply = llm.invoke(state["messages"])
    return {"messages": [{"role": "assistant", "content": reply.content}]}


def route_after_classify(state: GraphState) -> str:
    return "human_gate" if state.get("needs_approval") else "model"


def build_graph() -> Any:
    """Compile a tiny StateGraph: START → classify → (HITL | model) → END."""
    if DRY_RUN:
        return None

    from langgraph.graph import END, START, StateGraph

    builder = StateGraph(GraphState)
    builder.add_node("classify", classify_node)
    builder.add_node("model", model_node)
    builder.add_node("human_gate", lambda s: s)
    builder.add_edge(START, "classify")
    builder.add_conditional_edges(
        "classify",
        route_after_classify,
        {"human_gate": "human_gate", "model": "model"},
    )
    builder.add_edge("model", END)
    builder.add_edge("human_gate", END)
    return builder.compile()


def agent_invocation(payload: dict[str, Any], _context: Any = None) -> dict[str, str]:
    """AgentCore Runtime entrypoint shape — wrap with @app.entrypoint when live."""
    prompt = str(payload.get("prompt", "")).strip() or "No prompt found in input."
    state: GraphState = {
        "messages": [{"role": "user", "content": prompt}],
        "ticket_id": str(payload.get("ticket_id", "unknown")),
        "needs_approval": False,
    }
    updates = classify_node(state)
    state = {**state, **updates}
    if DRY_RUN:
        return {
            "result": "dry-run",
            "needs_approval": str(state["needs_approval"]).lower(),
            "note": "Set DRY_RUN=0 with langgraph + bedrock-agentcore installed to invoke Bedrock",
        }

    from bedrock_agentcore.runtime import BedrockAgentCoreApp  # noqa: F401

    graph = build_graph()
    output = graph.invoke(state)
    last = output["messages"][-1]
    content = last.get("content") if isinstance(last, dict) else str(last)
    return {"result": str(content)}


if __name__ == "__main__":
    print(
        agent_invocation(
            {"prompt": "Refund order 1842 for a damaged shipment", "ticket_id": "T-1842"}
        )
    )