import asyncio

from agent_contracts import (
    BaseAgentState,
    ModularNode,
    NodeContract,
    NodeInputs,
    NodeOutputs,
    TriggerCondition,
)
from agent_contracts import get_node_registry, build_graph_from_registry

# 1. Define a Node
class HelloNode(ModularNode):
    CONTRACT = NodeContract(
        name="hello",
        description="Returns a hello message",
        reads=["request"],
        writes=["response"],
        supervisor="main",
        trigger_conditions=[TriggerCondition(priority=100)],  # Always trigger first
        is_terminal=True,  # End the flow after this node
    )

    async def execute(self, inputs: NodeInputs, config=None) -> NodeOutputs:
        return NodeOutputs(
            response={
                "response_type": "done",
                "response_message": "Hello World!",
            }
        )

async def main() -> None:
    # 2. Register & Build
    registry = get_node_registry()
    registry.register(HelloNode)

    graph = build_graph_from_registry(
        registry=registry,
        supervisors=["main"],
        state_class=BaseAgentState,
    )
    graph.set_entry_point("main_supervisor")  # required for LangGraph compilation
    compiled = graph.compile()

# 3. Run
    result = await compiled.ainvoke({"request": {"action": "start"}})
    print(result["response"])

if __name__ == "__main__":
    asyncio.run(main())

from agent_contracts import ModularNode, NodeContract, TriggerCondition

# Node A: Runs when user asks for "weather" (LLM semantic match)
class WeatherNode(ModularNode):
    CONTRACT = NodeContract(
        name="weather_agent",
        description="Handles weather-related requests",
        reads=["request"],
        writes=["response"],
        supervisor="main",
        requires_llm=True,
        trigger_conditions=[
            TriggerCondition(
                llm_hint="User is asking about the weather forecast",
                priority=10
            )
        ]
    )
    # ... implementation ...

# Node B: Runs when a strict flag is present (Rule match)
class UrgentNode(ModularNode):
    CONTRACT = NodeContract(
        name="urgent_agent",
        description="Handles urgent/high-priority requests",
        reads=["request"],
        writes=["response"],
        supervisor="main",
        trigger_conditions=[
            TriggerCondition(
                when={"request.priority": "high"},
                priority=20  # Checked BEFORE LLM
            )
        ]
    )
    # ... implementation ...

NodeContract(
    name="my_node",
    description="Calculates mortgage payments",
    reads=["user_profile", "loan_data"],
    writes=["payment_schedule"],
    requires_llm=True,                 # Whether LLM is required
    supervisor="main",                 # Which supervisor manages this node
    trigger_conditions=[
        TriggerCondition(llm_hint="User asks about monthly payments")
    ]
)

from agent_contracts import InteractiveNode

class InterviewNode(InteractiveNode):
    CONTRACT = NodeContract(
        name="interview",
        description="Conversational workflow node",
        reads=["request", "_internal"],
        writes=["response", "_internal"],
        supervisor="main",
        trigger_conditions=[
            TriggerCondition(priority=10, llm_hint="Use for conversational workflows"),
        ],
    )
    
    def prepare_context(self, inputs):
        """Extract context from inputs."""
        return {"interview_state": inputs.get_slice("interview")}
    
    def check_completion(self, context, inputs):
        """Check if interview is complete."""
        return context["interview_state"].get("complete", False)
    
    async def process_answer(self, context, inputs):
        """Process user's answer."""
        # Handle the answer logic
        return True
    
    async def generate_question(self, context, inputs):
        """Generate next question."""
        return NodeOutputs(
            response={
                "response_type": "question",
                "response_data": {"question": "..."},
            }
        )

from agent_contracts import Internal, reset_response

# Bad
user_id = state["profile"]["id"]

# Good (agent-contracts)
user_id = Internal.user_id.get(state)

# Writing (returns new state)
state = Internal.turn_count.set(state, 5)
state = reset_response(state)

supervisor = GenericSupervisor(
    supervisor_name="shopping",
    llm=llm,
    max_field_length=10000  # Default: 10000 characters
)

from agent_contracts import GenericSupervisor

def ecommerce_context_builder(state: dict, candidates: list[str]) -> dict:
    """Build context for e-commerce routing decisions."""
    cart = state.get("cart", {})
    inventory = state.get("inventory", {})
    
    return {
        "slices": {"request", "response", "_internal", "cart", "inventory"},
        "summary": {
            "cart_total": sum(item["price"] for item in cart.get("items", [])),
            "low_stock_count": len([i for i in inventory.get("items", [])
                                     if i["quantity"] < 10]),
            "user_tier": state.get("user", {}).get("tier", "standard"),
        },
    }

supervisor = GenericSupervisor(
    supervisor_name="checkout",
    llm=llm,
    registry=registry,
    context_builder=ecommerce_context_builder,
)

def conversation_context_builder(state: dict, candidates: list[str]) -> dict:
    """Build context with conversation history."""
    messages = state.get("conversation", {}).get("messages", [])
    user_messages = [m for m in messages if m.get("role") == "user"]
    
    return {
        "slices": {"request", "response", "_internal", "conversation"},
        "summary": {
            "total_turns": len(user_messages),
            "last_question": messages[-2].get("content") if len(messages) >= 2 else None,
            "last_answer": messages[-1].get("content") if messages else None,
        },
    }

supervisor = GenericSupervisor(
    supervisor_name="assistant",
    llm=llm,
    context_builder=conversation_context_builder,
)

from agent_contracts import AgentRuntime, RequestContext, InMemorySessionStore

runtime = AgentRuntime(
    graph=compiled_graph,
    session_store=InMemorySessionStore(),
)

result = await runtime.execute(RequestContext(
    session_id="abc123",
    action="answer",
    message="I like casual style",
    resume_session=True, # Loads state from store
))

print(result.response_type)  # "interview", "proposals", etc.
print(result.response_data)  # Response payload

from agent_contracts.runtime import StreamingRuntime

runtime = (
    StreamingRuntime()
    .add_node("search", search_node, "Searching...")
    .add_node("stylist", stylist_node, "Generating recommendations...")
)

async for event in runtime.stream(request):
    yield event.to_sse()

async for event in runtime.stream_with_graph(
    request,
    graph=compiled_graph,
    include_subgraphs=True,
):
    yield event.to_sse()

from agent_contracts import RuntimeHooks, SessionStore

class MyHooks(RuntimeHooks):
    async def prepare_state(self, state, request):
        # Normalize or enrich state before execution
        return state
    
    async def after_execution(self, state, result):
        # Persist session, log, etc.
        pass

from agent_contracts.config import load_config
config = load_config("agent_config.yaml")

from agent_contracts import ContractVisualizer
visualizer = ContractVisualizer(registry, graph=compiled)
doc = visualizer.generate_architecture_doc()

print(visualizer.generate_langgraph_flow())
print(visualizer.generate_state_slices_section())