from enum import Enum
from dataclasses import dataclass
from agent_runtime_framework import (
    JourneyAgent, BaseJourneyState, BaseJourneyTools,
    AgentContext, ToolSchema, ToolSchemaBuilder
)

# 1. Define your journey steps
class OnboardingStep(str, Enum):
    WELCOME = "welcome"
    COLLECT_NAME = "collect_name"
    COLLECT_EMAIL = "collect_email"
    COMPLETE = "complete"

# 2. Define your state
@dataclass
class OnboardingState(BaseJourneyState[OnboardingStep]):
    step: OnboardingStep = OnboardingStep.WELCOME
    name: str = ""
    email: str = ""
    
    def is_complete(self) -> bool:
        return self.step == OnboardingStep.COMPLETE
    
    def to_dict(self) -> dict:
        return {
            "step": self.step.value,
            "name": self.name,
            "email": self.email,
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "OnboardingState":
        return cls(
            step=OnboardingStep(data.get("step", "welcome")),
            name=data.get("name", ""),
            email=data.get("email", ""),
        )

# 3. Define your tools
class OnboardingTools(BaseJourneyTools[OnboardingState]):
    async def save_name(self, name: str) -> str:
        self.state.name = name
        self.state.step = OnboardingStep.COLLECT_EMAIL
        await self._notify_state_change()
        return f"Great, {name}! Now, what's your email?"
    
    async def save_email(self, email: str) -> str:
        self.state.email = email
        self.state.step = OnboardingStep.COMPLETE
        await self._notify_state_change()
        return f"Perfect! You're all set, {self.state.name}!"

# 4. Define your agent
class OnboardingAgent(JourneyAgent[OnboardingState, OnboardingTools, OnboardingStep]):
    @property
    def key(self) -> str:
        return "onboarding-agent"
    
    def get_initial_state(self) -> OnboardingState:
        return OnboardingState()
    
    def get_system_prompt(self, state: OnboardingState) -> str:
        prompts = {
            OnboardingStep.WELCOME: "Welcome! Ask for the user's name.",
            OnboardingStep.COLLECT_NAME: "Collect the user's name using save_name tool.",
            OnboardingStep.COLLECT_EMAIL: "Collect the user's email using save_email tool.",
            OnboardingStep.COMPLETE: "Thank the user for completing onboarding.",
        }
        return prompts[state.step]
    
    def get_tool_schemas(self, state: OnboardingState) -> list[ToolSchema]:
        if state.step == OnboardingStep.COLLECT_NAME:
            return [
                ToolSchemaBuilder("save_name")
                .description("Save the user's name")
                .param("name", "string", "The user's name", required=True)
                .build()
            ]
        elif state.step == OnboardingStep.COLLECT_EMAIL:
            return [
                ToolSchemaBuilder("save_email")
                .description("Save the user's email")
                .param("email", "string", "The user's email", required=True)
                .build()
            ]
        return []
    
    def create_tools(self, state: OnboardingState, ctx: AgentContext) -> OnboardingTools:
        return OnboardingTools(state=state)
    
    async def execute_tool(self, tools: OnboardingTools, name: str, arguments: dict) -> str:
        method = getattr(tools, name)
        return await method(**arguments)

# 5. Run your agent
agent = OnboardingAgent(llm_client=my_llm_client)
result = await agent.run(context)

from agent_runtime_framework import configure

# Enable debug mode (development)
configure(debug=True)

# Enable production mode
configure(debug=False)

# Or use environment variable
# AGENT_RUNTIME_DEBUG=1 python your_app.py

from agent_runtime_framework import configure, JourneyAgent

# In debug mode, exceptions propagate
configure(debug=True)

class MyTools(BaseJourneyTools[MyState]):
    async def process_data(self, data: str) -> str:
        # This will raise immediately in debug mode
        result = int(data)  # ValueError if data is not a number
        return f"Processed: {result}"

# When this tool is called with invalid data, you'll get:
# ValueError: invalid literal for int() with base 10: 'abc'
# With full stack trace!

# In production mode, exceptions are caught
configure(debug=False)

# Same tool, but now returns:
# "Error executing process_data: invalid literal for int() with base 10: 'abc'"
# The agent continues running and can handle the error

from agent_runtime_framework import configure, FrameworkConfig

# Fine-grained control
configure(
    debug=True,                      # Enable debug mode
    swallow_tool_exceptions=False,   # Don't catch exceptions
    log_level="DEBUG"                # Verbose logging
)

# Or create a custom config
config = FrameworkConfig(
    debug=True,
    swallow_tool_exceptions=False,
    log_level="DEBUG"
)
from agent_runtime_framework import set_config
set_config(config)

# Environment variables
# AGENT_RUNTIME_DEBUG=1          # Enable debug mode
# AGENT_RUNTIME_LOG_LEVEL=DEBUG  # Set log level

from agent_runtime_framework import IntentRouter, RouteDefinition

class Journey(str, Enum):
    QUOTE = "quote"
    CLAIM = "claim"
    SUPPORT = "support"

router = IntentRouter[Journey]()
router.add_route(RouteDefinition(
    journey=Journey.QUOTE,
    name="start_quote",
    description="Get a new insurance quote",
))
router.add_route(RouteDefinition(
    journey=Journey.CLAIM,
    name="file_claim",
    description="File or check on a claim",
))

# Get routing tools for LLM
tools = router.get_tool_schemas()

# After LLM calls a routing tool
journey = router.resolve_tool_call("start_quote")  # Journey.QUOTE

from agent_runtime_framework import (
    StateStore, ConversationStore, MemoryManager
)

# Set up stores
state_store = StateStore()
conversation_store = ConversationStore()
manager = MemoryManager(state_store, conversation_store)

# Load context
context = await manager.load_context(conversation_id, "my-agent")

# Save after run
await manager.save_state(conversation_id, "my-agent", new_state)
await manager.save_messages(conversation_id, messages)

from agent_runtime_framework import PromptTemplate, StepPromptMapping

# Simple mapping
prompts = StepPromptMapping[MyStep](
    prompts={
        MyStep.WELCOME: "Welcome! How can I help?",
        MyStep.COLLECTING: "Please provide your information.",
    },
    default="I'm here to assist you.",
)

# With templates
template = PromptTemplate(
    "Hello $name! You are at step ${step}.",
    defaults={"name": "there"},
)
prompts.add(MyStep.WELCOME, template)

# Render
prompt = prompts.get(MyStep.WELCOME, name="Alice", step="welcome")

from agent_runtime_framework import LLMExecutor, ExecutorHooks, LoggingHooks

class MyHooks(ExecutorHooks):
    async def on_tool_start(self, name: str, arguments: dict) -> None:
        print(f"🔧 Calling tool: {name}")

    async def on_tool_end(self, name: str, result: str) -> None:
        print(f"✅ Tool completed: {name}")

executor = LLMExecutor(
    llm_client=my_llm,
    hooks=MyHooks(),
)

from agent_runtime_core.llm import get_llm_client

# Framework automatically uses agent_runtime_core's LLM client
agent = MyAgent()  # No need to pass llm_client explicitly

from agent_runtime_framework.adapters import DjangoRuntimeAdapter

class MyDjangoAgent(DjangoRuntimeAdapter[MyState, MyTools, MyStep]):
    @property
    def key(self) -> str:
        return "my-agent"

    def get_initial_state(self) -> MyState:
        return MyState()

    def get_system_prompt(self, state: MyState) -> str:
        return PROMPTS[state.step]

    def get_tool_schemas(self, state: MyState) -> list[ToolSchema]:
        return TOOLS[state.step]

    def create_tools(self, state: MyState, ctx, backend_client) -> MyTools:
        return MyTools(state=state, backend_client=backend_client)

    async def execute_tool(self, tools: MyTools, name: str, args: dict) -> str:
        method = getattr(tools, name)
        return await method(**args)

# Register with Django
from django_agent_runtime.runtime.registry import register_runtime
register_runtime(MyDjangoAgent())

schema = (
    ToolSchemaBuilder("my_tool")
    .description("What the tool does")
    .param("arg1", "string", "Description", required=True)
    .param("arg2", "number", "Description", required=False)
    .build()
)

executor = LLMExecutor(
    llm_client=my_llm,
    tool_executor=MethodToolExecutor(tools),
    config=LLMExecutorConfig(max_iterations=10),
    hooks=MyHooks(),
)

result = await executor.run(
    messages=[{"role": "user", "content": "Hello"}],
    tools=[tool_schema],
    system_prompt="You are helpful.",
)

router = IntentRouter[MyJourney]()
router.add_route(RouteDefinition(
    journey=MyJourney.QUOTE,
    name="start_quote",
    description="Start a quote journey",
))

# Get tool schemas for LLM
tools = router.get_tool_schemas()

# Resolve tool call to journey
journey = router.resolve_tool_call("start_quote")

manager = PromptManager[MyStep](
    mapping=StepPromptMapping(...),
    context_enricher=lambda state: {"user_name": state.name},
)

prompt = manager.get_prompt(state)