session.chat_ctx

import asyncio
import json
from typing import Any, Dict

# --- Replace these with your actual LiveKit/agent imports ---
# from livekit_agent import Agent, AgentSession, ctx
# ----------------------------------------------------------------

# Simple async key-value store (replace with Redis/Postgres)
_STORE: Dict[str, str] = {}

async def save_to_store(key: str, value: Dict[str, Any]):
    _STORE[key] = json.dumps(value)

async def load_from_store(key: str) -> Dict[str, Any] | None:
    raw = _STORE.get(key)
    return json.loads(raw) if raw else None

# --- Example agent lifecycle handlers ---

async def on_agent_start(room_name: str):
    """
    Called when agent process starts / reconnects.
    Restore chat_ctx if present and create the Agent with it.
    """
    stored = await load_from_store(f"chatctx:{room_name}")
    if stored:
        # Construct agent with restored chat context
        agent = Agent(chat_ctx=stored)  # replace with your Agent class
    else:
        agent = Agent()  # fresh agent

    # Hook to save context on shutdown/disconnect
    async def save_context():
        # session referenced from agent; session.chat_ctx holds context/history
        ctx_to_save = agent.session.chat_ctx
        await save_to_store(f"chatctx:{room_name}", ctx_to_save)

    # Register your shutdown callback with your framework/context
    ctx.add_shutdown_callback(save_context)

    return agent

async def push_user_message(room_name: str, user_id: str, text: str):
    """
    Push a new user chat message into the agent's session so it continues
    from previous context and processes the new message.
    """
    # Load stored chat_ctx and ensure agent has it (if agent not running, you can
    # create a transient Agent instance with restored chat_ctx to process the message)
    stored = await load_from_store(f"chatctx:{room_name}")
    if stored:
        agent = Agent(chat_ctx=stored)
    else:
        agent = Agent()

    # Add the incoming user message into the session history/chat context.
    # How to add depends on your AgentSession API; example:
    agent.session.add_user_message(user_id=user_id, text=text)

    # Now run the agent's processing loop / trigger a response handler
    # This might be a single-step call or an async run depending on your API.
    await agent.process_once()  # replace with actual processing call

    # After processing, save updated chat_ctx back to store
    await save_to_store(f"chatctx:{room_name}", agent.session.chat_ctx)

    # Return agent response (example)
    return agent.session.last_system_message()

# --- Usage example ---
async def example_flow():
    room = "support-room-123"
    # Agent process starts (restore context if exists)
    agent = await on_agent_start(room)

    # Simulate agent disconnect: shutdown callback will fire and save context
    await ctx.shutdown()  # this triggers save_context callback in on_agent_start

    # Later: new user message arrives while agent reconnects
    response = await push_user_message(room, user_id="user_1", text="I still need help with billing.")
    print("Agent responded:", response)

# Run example (in real app, these are event-driven)
# asyncio.run(example_flow())

_STORE

agent.session.chat_ctx

agent.session.add_user_message(...)

process_once()

ctx.add_shutdown_callback(save_context)

userdata

chat_ctx