from livekit import Agent, pipeline, llm, StateStore
from livekit.tools import function_tool
from pydantic import BaseModel

# Simple in-memory state store for loop-like iteration tracking
state = StateStore()  # replace with your persistent store in production

# Tool that performs one step and returns whether more steps are needed
class StepResult(BaseModel):
    ok: bool
    continue_loop: bool
    message: str

@function_tool(name="perform_step", return_type=StepResult)
def perform_step(step_input: str, iteration: int) -> StepResult:
    # Strict validation and deterministic behavior
    if not step_input:
        return StepResult(ok=False, continue_loop=False, message="missing input")
    # Example deterministic rule: stop after 3 iterations
    if iteration >= 3:
        return StepResult(ok=True, continue_loop=False, message="done")
    # Perform action for this iteration
    return StepResult(ok=True, continue_loop=True, message=f"processed {iteration}")

# Pipeline node override to enforce rules before LLM runs
@pipeline.node("llm_node")
def llm_node_override(context, next_node):
    user_text = context.input_text.strip()
    # Enforce rule: reject empty user messages
    if not user_text:
        return llm.Response(content="Please provide input.", finished=True)

    # Load iteration counter from state (per-session key)
    session_id = context.session_id
    iteration = state.get(session_id, default=0)

    # Force the model to only call the tool and produce structured output
    context.tools = ["perform_step"]
    context.forced_structure = {"type": "tool_call", "tool": "perform_step"}

    # Attach iteration info to context so tool receives it
    context.metadata["iteration"] = iteration

    # Run next node (LLM) which should call the tool
    resp = next_node(context)

    # After LLM/tool run, read tool output from context.results
    tool_out = context.results.get("perform_step")
    if tool_out:
        iteration += 1
        state.set(session_id, iteration)

        if tool_out.get("continue_loop"):
            # Re-run the pipeline to emulate another loop iteration
            # Note: in real code, schedule or recursively invoke the flow
            return llm.Response(content=f"Iteration {iteration} done. Continuing...", finished=False)
        else:
            # Finalize
            return llm.Response(content=f"Finished after {iteration} iterations. {tool_out.get('message')}", finished=True)

    return resp

# Agent setup
agent = Agent(name="strict_agent")
agent.add_tool(perform_step)
agent.register_pipeline_node(llm_node_override)

# Example usage (framework-specific call)
# agent.handle_user_message(session_id="sess123", text="start process")