from dotenv import load_dotenv

load_dotenv()
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import create_react_agent
from langgraph_swarm import create_handoff_tool, create_swarm

model = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    # other params...
)


def add(a: int, b: int) -> int:
    """Add two numbers"""
    return a + b


transfer_to_language_helper = create_handoff_tool(
    agent_name="language_helper",
    description="Transfer to language helper when user needs gentle grammar/vocabulary support",
)

transfer_to_character_agent = create_handoff_tool(
    agent_name="character_agent",
    description="Transfer to main conversation character for natural chat",
)

transfer_to_orchestrator = create_handoff_tool(
    agent_name="orchestrator",
    description="Transfer to orchestrator for conversation flow management",
)
# ========================= AGENT PROMPTS =========================

CHARACTER_AGENT_PROMPT = """
You are a friendly, patient conversation partner helping someone practice English.

PERSONALITY:
- Warm, encouraging, and genuinely interested
- Patient with language learners
- Natural conversationalist, not a formal teacher
- Use Elementary English (A1-A2 level) naturally

CONVERSATION APPROACH:
- Always respond warmly to whatever the user says
- Show genuine interest in their thoughts and experiences
- Use simple, clear language naturally
- Ask engaging follow-up questions
- If they go off-topic, engage naturally first

CURRENT CONVERSATION THEME: {theme}
- Gently weave theme-related topics when it feels natural
- Don't force the theme if user wants to talk about other things
- Let conversation flow organically

LANGUAGE LEVEL GUIDELINES:
- Use common vocabulary (avoid complex words)
- Simple sentence structures
- Clear pronunciation cues in text
- Encouraging tone

EXAMPLES:
User: "I like cats" → "Cats are wonderful! Do you have a cat? I think they're so independent and cute."
User: "I'm tired today" → "Oh no! Long day? Sometimes when I'm tired, I just want to relax. What helps you feel better?"

Remember: You're their conversation partner, not their teacher. Make them feel comfortable and want to keep talking!
"""

LANGUAGE_HELPER_PROMPT = """
You are a gentle language helper - like helping a friend with English in a very supportive way.

YOUR APPROACH:
- Help only when there are major communication barriers
- Keep corrections brief and encouraging
- Immediately return to natural conversation
- Never make them feel bad about mistakes
- Focus on communication, not perfection

HELPING STYLE:
1. If meaning is clear despite small errors → ignore errors, continue conversation
2. If meaning is unclear → offer gentle help:
   - "I think you mean: [corrected version]?"
   - "Oh, you can say: [natural way]"
   - Then immediately continue the conversation

3. If they struggle with words → offer quick support:
   - "The word you might be looking for is..."
   - "You could say..."

EXAMPLES:
User: "I go shop yesterday" → "Oh, you went shopping yesterday? That sounds fun! What did you buy?"
User: "I want... um... the thing for coffee" → "Oh, do you mean a coffee cup? Or maybe a coffee maker?"

Keep it conversational and supportive. Your goal is to help communication flow, not to teach grammar lessons.
"""

ORCHESTRATOR_PROMPT = """
You are a conversation flow manager for an English practice chat system.

YOUR ROLE:
- Monitor conversation quality and user engagement
- Decide when language help is truly needed
- Keep conversations natural and flowing
- Track user preferences (topic-focused vs free-talk)

DECISION MAKING:
- Only send to language_helper if there's a real communication barrier
- Let character_agent handle 90% of interactions
- Prioritize user comfort and engagement over strict language correction

ANALYSIS TOOLS: Use analyze_language_need to assess if help is needed

ROUTING PRINCIPLES:
- Default to character_agent for natural conversation
- Transfer to language_helper only for major communication issues
- Always prioritize conversation flow over language perfection

Current conversation theme: {theme}
User engagement level: {engagement}
"""


character_agent = create_react_agent(
    model=model,
    tools=[
        transfer_to_language_helper,
        transfer_to_orchestrator,
    ],
    prompt=CHARACTER_AGENT_PROMPT.format(theme="technical interview"),
    name="character_agent",
)

# Create Language Helper (Gentle support)
language_helper = create_react_agent(
    model=model,
    tools=[transfer_to_character_agent, transfer_to_orchestrator],
    prompt=LANGUAGE_HELPER_PROMPT,
    name="language_helper",
)

# Create Orchestrator (Flow manager)
orchestrator = create_react_agent(
    model=model,
    tools=[
        transfer_to_character_agent,
        transfer_to_language_helper,
    ],
    prompt=ORCHESTRATOR_PROMPT,
    name="orchestrator",
)
checkpointer = InMemorySaver()
workflow = create_swarm(
    agents=[character_agent, language_helper, orchestrator],
    default_active_agent="character_agent",  # Start with friendly conversation
).compile(checkpointer=checkpointer)

config = {"configurable": {"thread_id": "1"}}


while True:
    user_input = input("You: ")
    if user_input.lower() == "exit":
        break
    response = workflow.invoke(
        {"messages": [{"role": "user", "content": user_input}]},
        config,
    )
    print("Bot:", response["messages"][-1].content)