#pip install langchain_google_genai langgraph gradio
import os
import sys
import typing
from typing import Annotated, Literal, Iterable
from typing_extensions import TypedDict

from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langchain_core.tools import tool
from langchain_core.messages import AIMessage, ToolMessage, HumanMessage, BaseMessage, SystemMessage
from random import randint

import wikipedia

import gradio as gr
import logging

class OrderState(TypedDict):
    """State representing the customer's order conversation."""
    messages: Annotated[list, add_messages]
    order: list[str]
    finished: bool

# System instruction for the BaristaBot
SYSINT = (
    "system",
    "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: "
    "FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings."
    "If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise."
    "If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise."
    "If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."
    "If a tool required for task completion is unavailable after multiple tries, return 0."
    )

WELCOME_MSG = "Welcome to the BaristaBot cafe. Type `q` to quit. How may I serve you today?"

# Initialize the Google Gemini LLM
llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest")

@tool
def wikipedia_search(title: str) -> str:
    """Provides a short snippet from a Wikipedia article with the given title"""
    page = wikipedia.page(title)
    return page.content[:300]

def agent_node(state: OrderState) -> OrderState:
    """agent with tool handling."""
    print(f"Messagelist sent to agent node: {[msg.content for msg in state.get('messages', [])]}")
    defaults = {"order": [], "finished": False}

    # Ensure we always have at least a system message
    if not state.get("messages", []):
        return defaults | state | {"messages": []}

    try:
        # Prepend system instruction if not already present
        messages_with_system = [
            SystemMessage(content=SYSINT)
        ] + state.get("messages", [])

        # Process messages through the LLM
        new_output = llm_with_tools.invoke(messages_with_system)

        return defaults | state | {"messages": [new_output]}
    except Exception as e:
        # Fallback if LLM processing fails
        return defaults | state | {"messages": [AIMessage(content=f"I'm having trouble processing that. {str(e)}")]}

def maybe_route_to_tools(state: OrderState) -> str:
    """Route between chat and tool nodes."""
    if not (msgs := state.get("messages", [])):
        raise ValueError(f"No messages found when parsing state: {state}")

    msg = msgs[-1]

    if state.get("finished", False):
        print("from agent GOTO End node")
        return END

    elif hasattr(msg, "tool_calls") and len(msg.tool_calls) > 0:
        if any(tool["name"] in tool_node.tools_by_name.keys() for tool in msg.tool_calls):
            print("from agent GOTO tools node")
            return "tools"

    print("tool call failed, letting agent try again")
    return "human"

def human_node(state: OrderState) -> OrderState:
    """Handle user input."""
    logging.info(f"Messagelist sent to human node: {[msg.content for msg in state.get('messages', [])]}")
    last_msg = state["messages"][-1]

    if last_msg.content.lower() in {"q", "quit", "exit", "goodbye"}:
        state["finished"] = True

    return state

def maybe_exit_human_node(state: OrderState) -> Literal["agent", "__end__"]:
    """Determine if conversation should continue."""
    if state.get("finished", False):
        logging.info("from human GOTO End node")
        return END
    last_msg = state["messages"][-1]
    if isinstance(last_msg, AIMessage):
        logging.info("Chatbot response obtained, ending conversation")
        return END
    else:
        logging.info("from human GOTO agent node")
        return "agent"

# Prepare tools
auto_tools = []
tool_node = ToolNode(auto_tools)

interactive_tools = [wikipedia_search]

# Bind all tools to the LLM
llm_with_tools = llm.bind_tools(auto_tools + interactive_tools)

# Build the graph
graph_builder = StateGraph(OrderState)

# Add nodes
graph_builder.add_node("agent", agent_node)
graph_builder.add_node("human", human_node)
graph_builder.add_node("tools", tool_node)

# Add edges and routing
graph_builder.add_conditional_edges("agent", maybe_route_to_tools)
graph_builder.add_conditional_edges("human", maybe_exit_human_node)
graph_builder.add_edge("tools", "agent")
graph_builder.add_edge("ordering", "agent")
graph_builder.add_edge(START, "human")

# Compile the graph
chat_graph = graph_builder.compile()

def convert_history_to_messages(history: list) -> list[BaseMessage]:
    """
    Convert Gradio chat history to a list of Langchain messages.
    
    Args:
    - history: Gradio's chat history format
    
    Returns:
    - List of Langchain BaseMessage objects
    """
    messages = []
    for human, ai in history:
        if human:
            messages.append(HumanMessage(content=human))
        if ai:
            messages.append(AIMessage(content=ai))
    return messages

def gradio_chat(message: str, history: list) -> str:
    """
    Gradio-compatible chat function that manages the conversation state.
    
    Args:
    - message: User's input message
    - history: Gradio's chat history
    
    Returns:
    - Bot's response as a string
    """
    logging.info(f"{len(history)} history so far: {history}")
    # Ensure non-empty message
    if not message or message.strip() == "":
        message = "Hello, how can I help you today?"

    # Convert history to Langchain messages
    conversation_messages = []
    for old_message in history:
        if old_message["content"].strip():
            if old_message["role"] == "user":
                conversation_messages.append(HumanMessage(content=old_message["content"]))
            if old_message["role"] == "assistant":
                conversation_messages.append(AIMessage(content=old_message["content"]))
    
    # Add current message
    conversation_messages.append(HumanMessage(content=message))
    
    # Create initial state with conversation history
    conversation_state = {
        "messages": conversation_messages, 
        "order": [], 
        "finished": False
    }
    logging.info(f"Conversation so far: {str(conversation_state)}")
    try:
        # Process the conversation through the graph
        conversation_state = chat_graph.invoke(conversation_state, {"recursion_limit": 10})
        
        # Extract the latest bot message
        latest_message = conversation_state["messages"][-1]
        
        # Return the bot's response content
        logging.info(f"return: {latest_message.content}")
        return latest_message.content
    
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Gradio interface
def launch_baristabot():
    gr.ChatInterface(
        gradio_chat,
        type="messages",
        title="BaristaBot",
        description="Your friendly AI cafe assistant",
        theme="ocean"
    ).launch()

if __name__ == "__main__":
    # initiate logging tool
    logging.basicConfig(
                    stream=sys.stdout,
                    level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')
    launch_baristabot()