# ! pip install langgraph
# ! pip install langchain-core 
# ! pip install langchain-openai 
# ! pip install langchain-groq 
# ! pip install langchain-community 
# ! pip install python-dotenv 
# ! pip install pydantic 
# ! pip install typing-extensions 
# ! pip install chromadb 
# ! pip install langchain-text-splitters

from typing import Dict, TypedDict
from langgraph.graph import StateGraph, END
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
from IPython.display import display, Image
from langchain_core.runnables.graph import MermaidDrawMethod
from dotenv import load_dotenv
import os

import os, getpass

def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

_set_env("OPENAI_API_KEY")
_set_env("LANGCHAIN_API_KEY")
_set_env("GROQ_API_KEY")
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "Hackathon-Contextual-Quoting"

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
# llm = ChatGroq(model="llama3-groq-8b-8192-tool-use-preview", temperature=0)

import sqlite3

# Connect to SQLite (or create the database if it doesn't exist)
conn = sqlite3.connect('data/categories.db')

# Create a cursor object to execute SQL commands
cursor = conn.cursor()

# Create the category_rates table
cursor.execute('''
    CREATE TABLE IF NOT EXISTS category_rates (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        category TEXT NOT NULL,
        rate REAL NOT NULL
    )
''')

# Data to be inserted into the category_rates table
categories = [
    ('ADMINISTRATIVE', 0.61),
    ('ADVERTISING', 5.31),
    ('ARTISAN CONTRACTORS', 2.35),
    ('CLUBS/ASSOCIATIONS', 6.53),
    ('CONSULTANTS', 1.97),
    ('EDUCATION', 9.63),
    ('HEALTHCARE', 9.38),
    ('OTHER SERVICES', 6.79),
    ('PROFESSIONAL SERVICES', 8.22),
    ('RESTAURANTS/MOBILE FOOD SERVICES', 2.85),
    ('RETAIL', 5.80),
    ('TECHNOLOGY', 8.57),
    ('WHOLESALE/DISTRIBUTION', 5.61)
]

# Insert data into the category_rates table
cursor.executemany('''
    INSERT INTO category_rates (category, rate)
    VALUES (?, ?)
''', categories)

# Commit the changes and close the connection
conn.commit()
conn.close()

print("Database and table created successfully, and data inserted.")

from pydantic import BaseModel, Field
from typing import Optional, List

# Core business information model for initial data collection
class BusinessInformation(BaseModel):
    legal_name: Optional[str] = Field(default="Info not yet provided", description="The legal name of the business")
    primary_products_or_services: Optional[str] = Field(default="Info not yet provided", description="The primary products or services offered")
    secondary_or_ancillary_operations: Optional[str] = Field(default="Info not yet provided", description="Any secondary or ancillary operations")
    industries: Optional[str] = Field(default="Info not yet provided", description="In which industries does the business operate?")
    manufacturing_retail_wholesale_or_distribution: Optional[str] = Field(default="Info not yet provided", 
        description="Does the business engage in manufacturing, retail, wholesale, or distribution?")
    projected_revenue: Optional[str] = Field(default="Info not yet provided", description="The projected revenue for the upcoming year")

# Models for workflow routing between AI assistants
class ToUnderwritingAssistant(BaseModel):
    """Transfer work to the underwriting assistant"""

    request: str = Field(
        description="Any information that the underwriting assistant should know before proceeding."
     )

    class Config:
        json_schema_extra = {
            "example": {
                "request": "The business is a landscaping company that offers residential and commercial landscaping services, including lawn maintenance, garden design, tree trimming, and irrigation system installation."
            }
        }

# Models for classification and code assignment
class FinalClassificationsInfo(BaseModel):
    """Information about the final classifications and code that best describe the business."""
    
    category: str = Field(
        description="A string representing the category."
    )
    classification_description: str = Field(
        description="A string representing the category description."
    )
    code: int = Field(
        description="An integer representing the code assigned to the category."
    )

class ToQuoteAssistant(BaseModel):
    """Transfer work to the underwriting assistant"""

    request: str = Field(
        description="Any information that the underwriting assistant should know before proceeding."
     )
    
    final_classifications: list[FinalClassificationsInfo] = Field(
        description="A list of FinalClassificationsInfo objects, each containing a category, category description, and code."
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "request": "The business is a landscaping company that offers residential and commercial landscaping services, including lawn maintenance, garden design, tree trimming, and irrigation system installation.",
                "final_classifications": [
                    {"category": "CONSULTANTS", "classification_description": "Consultants--Business Services", "code": 74379},
                    {"category": "RESTAURANTS/MOBILE FOOD SERVICES", "classification_description": "Bakeries–Fast Food–Donut Shop.", "code": 10100}
                ]
            }
        }

# Workflow control models
class CompleteOrEscalate(BaseModel):
    """A tool to mark the current task as completed and/or to escalate control of the workflow to the main assistant,
    who can re-route the dialog based on the user's needs."""

    cancel: bool = True
    reason: str

    class Config:
        json_schema_extra = {
            "example": {
                "cancel": True,
                "reason": "Successfully obtained a quote.",
            },
            "example 2": {
                "cancel": True,
                "reason": "I have fully completed the task.",
            },
            "example 3": {
                "cancel": False,
                "reason": "I need to search the user's emails or calendar for more information.",
            },
        }

# Models for classification reasoning and rationale
class CategoryRationale(BaseModel):
    """Represents a category and its rationale or reasoning."""
    category: str = Field(..., description="The name of the category")
    cat_description: str = Field(..., description="The description of the category")
    code: int = Field(..., description="The unique code for the category")
    rate: float = Field(..., description="The rate for the category")
    rationalization: str = Field(..., description="A detailed rationalization for the category")

class CategoriesOutput(BaseModel):
    """Output structure for a list of categories with rationale explanations."""
    categories: List[CategoryRationale] = Field(..., description="List of categories with their rationale explanations")

# Models for classification accuracy assessment
class DescriptionGrade(BaseModel):
    """Schema to define a business description and its corresponding grade.
    
    Attributes:
        description (str): A string representing the business description.
        grade (float): A float representing the grade assigned to the business description.
    """
    category: str = Field(
        description="A string representing the category."
    )
    cat_description: str = Field(
        description="A string representing the category description."
    )
    code: int = Field(
        description="An integer representing the code assigned to the category."
    )
    grade: float = Field(
        description="A float representing the grade assigned to the business description."
    )

class DescriptionGradeSchema(BaseModel):
    """Schema to hold a list of business descriptions with their corresponding grades.
    
    Attributes:
        description_grades (list[DescriptionGrade]): A list of DescriptionGrade objects, each containing a business description and its corresponding grade.
    """
    description_grades: list[DescriptionGrade] = Field(
        description="A list of DescriptionGrade objects, each containing a business description and its corresponding grade."
    )

from typing import Annotated, Literal, Optional
from typing_extensions import TypedDict
from langgraph.graph.message import AnyMessage, add_messages

# State Reducer
def update_workflow_state(left: list[str], right: Optional[str]) -> list[str]:
    """Push or pop the workflow state, preventing duplicates."""
    if right is None or right == []:
        return left

    # Handle 'pop' case to remove the last element
    if right == "pop":  
        return left[:-1] if left else left  # Safeguard against empty lists
    
    # Handle cases where 'right' is a list (from the tutorial)
    if isinstance(right, list):
        # Avoid adding elements already in the list
        return left + [item for item in right if item not in left]

    # If 'right' is a single string, avoid duplication
    if right not in left:
        return left + [right]

    return left 

# Main State
class MainState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    business_information: dict[str, str | float | int]  # Allow values to be str, float, or int
    workflow_state: Annotated[
        list[
            Literal[
                "main_assistant",
                "underwriting_assistant",
                "quote_assistant",
            ]
        ],
        update_workflow_state
    ]

# RAG State
class RAGState(TypedDict):
    description: str
    documents: list[str]
    rationale: str
    graded_categories: list[dict[str, float]]

# Extra State
class ExtraState(TypedDict):
    tool_call_id: str
    final_classifications: dict

from langgraph.prebuilt import ToolNode
from langchain_core.runnables import RunnableLambda
from langchain_core.messages import ToolMessage

# Handle tool errors
def handle_tool_error(state) -> dict:
    error = state.get("error")
    tool_calls = state["messages"][-1].tool_calls
    return {
        "messages": [
            ToolMessage(
                content=f"Error: {repr(error)}\\n please fix your mistakes.",
                tool_call_id=tc["id"],
            )
            for tc in tool_calls
        ]
    }

# Create a tool node with fallback
def create_tool_node_with_fallback(tools: list) -> dict:
    return ToolNode(tools).with_fallbacks(
        [RunnableLambda(handle_tool_error)], exception_key="error"
    )

from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Load Categories
categories_path = "data/cat_codes_descriptions.csv"

loader = TextLoader(file_path=categories_path)

data_txt = loader.load()

# Split Categories
text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=150,  # Smaller chunk size to ensure one line per chunk
        chunk_overlap=0,  # No overlap needed for line-by-line splitting
        separators=["\n"],  # Primary separator is line breaks
        length_function=len,  # Length function based on character count
    )

splits = text_splitter.split_documents(data_txt)

from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings

# Add to vectorDB
vectorstore = Chroma.from_documents(documents=splits,
                                    collection_name="rag-chroma",
                                    embedding=OpenAIEmbeddings()
                                )
retriever = vectorstore.as_retriever(search_kwargs={"k": 10})

from langchain_core.tools import tool

# Business Info Gathering Tool
@tool
def info_tool(business_info: BusinessInformation) -> MainState:
    """
    A tool to process and store business information.
    """
    result = {
        "business_information": business_info.model_dump()
    }
    return result

# Quote Tool
@tool
def fetch_rates(categories: list[str], projected_revenue: float) -> str:
    """
    Fetches the rates for given categories and calculates the computed results based on the projected revenue.

    Args:
        categories (list[str]): A list of category names for which rates need to be fetched. Example categories include:
            - CATEGORY
            - ADMINISTRATIVE
            - ADVERTISING
            - ARTISAN CONTRACTORS
            - CLUBS/ASSOCIATIONS
            - CONSULTANTS
            - EDUCATION
            - HEALTHCARE
            - OTHER SERVICES
            - PROFESSIONAL SERVICES
            - RESTAURANTS/MOBILE FOOD SERVICES
            - RETAIL
            - TECHNOLOGY
            - WHOLESALE/DISTRIBUTION
        projected_revenue (float): The projected revenue to be distributed among the categories.

    Returns:
        str: A formatted message containing the quote breakdown
    """
    # Connect to SQLite database
    conn = sqlite3.connect('data/categories.db')
    cursor = conn.cursor()

    # Prepare parts to build the formatted message
    message_parts = ["Here is the requested quote:\n", "Quote Breakdown:"]

    # Initialize total premium
    total_premium = 0

    # Fetch rate for each classification
    for classification in categories:
        cursor.execute('''
            SELECT rate FROM category_rates WHERE category = ?
        ''', (classification,))
        result = cursor.fetchone()
        if result:
            rate = result[0]
            portion_revenue = projected_revenue / len(categories)
            computed_result = (portion_revenue / 1000) * rate
            total_premium += computed_result  # Accumulate total premium

            # Add formatted message
            message_parts.append(
                f"\nCategory: {classification}"
                f"\n  - Rate: ${rate:.2f} per $1,000"
                f"\n  - Portion of Projected Revenue: ${portion_revenue:,.2f}"
                f"\n  - Computed Premium: ${computed_result:,.2f}"
            )
        else:
            message_parts.append(f"\nWarning: No rate found for category: {classification}")

    # Add total premium to the message
    message_parts.append(f"\n\nTotal Annual Premium: ${total_premium:,.2f}")

    # Close the connection
    conn.close()

    return "\n".join(message_parts)

from langchain_core.runnables import Runnable, RunnableConfig

# Helper Assistant Class for handling LLM calls
class Assistant:
    def __init__(self, runnable: Runnable):
        self.runnable = runnable.with_config()
        
    def __call__(self, state: MainState, config: RunnableConfig):
        # Create input with both messages and business_information
        input_state = {
            "messages": state["messages"],
            "business_information": state.get("business_information", None),
        }
        
        while True:
            result = self.runnable.invoke(input_state)
            # If the LLM happens to return an empty response, we will re-prompt it
            if not result.tool_calls and (
                not result.content
                or isinstance(result.content, list)
                and not result.content[0].get("text")
            ):
                messages = state["messages"] + [("user", "Respond with a real output.")]
                state = {**state, "messages": messages}
                input_state["messages"] = messages
            else:
                break
        return {"messages": result}

# Main Assistant Prompt
main_assistant_prompt = ChatPromptTemplate.from_messages(
    [
        ("system",
        """
You are the main commercial underwriter assistant.\n
You are responsible for guiding the user through the quoting process and ensuring that the information they provide is accurate and consistent, to later provide a quote. 
When details and/or information about the business are provided, use "info_tool" to process the information and gather or update the required data.\n\n
Once you have gathered the information, you must first evaluate the prospects information with the underwriters by calling "ToUnderwritingAssistant".\n
Before proceeding to quote, you must provide the user with a summary of the information they have provided.\n
**Only call one tool at a time.**\n
Current business information:\n
{business_information}
""",
         ),
        ("placeholder", "{messages}"),
    ]
)

# Main Assistant Tools
main_assistant_tool = [info_tool]

# Runnable
main_assistant_runnable = main_assistant_prompt | llm.bind_tools(main_assistant_tool + [ToUnderwritingAssistant])

# Underwriting Assistant Prompt
underwriting_assistant_prompt = ChatPromptTemplate.from_messages(
    [
        ("system",
        """
You are the main commercial underwriter assistant.\n
You are responsible for evaluating the prospects information as he principal underwriter.\n
**Only call one tool at a time.**\n
When all the required information and categories have been confirmed, call "ToQuoteAssistant" to proceed to quote. 
Current business information:\n
{business_information}
""",
         ),
        ("placeholder", "{messages}"),
    ]
)

# Runnable
underwriting_assistant_runnable = underwriting_assistant_prompt | llm.bind_tools([ToQuoteAssistant])

# Quote Assistant Prompt
quote_prompt = ChatPromptTemplate.from_messages([
    ("system",
    """
You are an expert commercial underwriting assistant tasked with providing a quote for the business based on the classifications.

You will use the provided classifications and their corresponding rates to calculate the annual premium.

IMPORTANT: 
1. Call fetch_rates ONLY ONCE with all required categories and the projected revenue.
2. After receiving the rates and calculating the premium, call CompleteOrEscalate to end the process.
3. Do not call fetch_rates again unless explicitly asked to recalculate.

Business Description:
{business_information}

When you have the quote results, provide a clear summary to the user and then call CompleteOrEscalate with cancel=true to complete the process.
""",
         ),
        ("placeholder", "{messages}"),
    ]
)

# Quote Assistant Tools
quote_assistant_tools = [fetch_rates]

# Runnable
quote_assistant_runnable = quote_prompt | llm.bind_tools(quote_assistant_tools + [CompleteOrEscalate])

from langchain_core.output_parsers import StrOutputParser

# Retriever Assistant Prompt
retriever_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """
You are an expert commercial underwriting assistant specialized in analyzing business descriptions to extract key operational details relevant to the insurance underwriting process.

**Your task is to:**

1. **Analyze** the provided business description.
2. **Identify and summarize** the key aspects of the business operations, focusing on workflows, operational tasks, products/services offered, and any other relevant details.
3. **Highlight** any specific activities that are crucial for matching the business to the appropriate insurance code classification.

**Please provide the summary in a clear, concise, and structured format, using bullet points for each key aspect.**

---

**Business Information:**

{business_information}

---

**Example:**

**Business Description:**

"**ABC Landscaping** offers residential and commercial landscaping services, including lawn maintenance, garden design, tree trimming, and irrigation system installation. They also provide seasonal services like snow removal. The company operates a fleet of service vehicles and heavy machinery such as lawn mowers, trimmers, and excavators."

**Summary:**

- Provides residential and commercial landscaping services.
- Services include lawn maintenance, garden design, and tree trimming.
- Offers installation of irrigation systems.
- Provides seasonal snow removal services.
- Operates a fleet of service vehicles and heavy machinery (lawn mowers, trimmers, excavators).

""",
        )
    ]
)

# Runnable
retriever_runnable = retriever_prompt | llm | StrOutputParser()

# Reasoning Categories Assistant Prompt
reasoning_prompt = ChatPromptTemplate.from_messages(
    [
        ("system",
        """
You are an expert commercial underwriting assistant tasked with providing reasoning for the categories of the nature of operations of a business.\n
You will be given a description of the business and a list of categories that may apply to the business.\n
Use logical reasoning to determine which categories are relevant to the business based on the provided description.\n
Business Description:\n
{description}\n
Categories:\n
{documents}\n

---

**Example:**


{{
    "categories": [
      {{
        "category": "CONSULTANTS",
        "description": "Consultants--Business Services"
        "code": "74379",
        "rate": "1.97",
        "rationalization": "GABAS, LLC provides consulting services related to AI and technology, which falls under the category of business services. This includes advising businesses on AI implementation and strategy."
      }},
      {{
        "category": "RESTAURANTS/MOBILE FOOD SERVICES",
        "description": "Bakeries–Fast Food–Donut Shop."
        "code": "10100",
        "rate": "2.85",
        "rationalization": "A bakery that offers fast food and donut services is considered part of the food service industry."
      }}
    ]
}}

""",
         )
    ]
)

# Runnable
reasoning_runnable = reasoning_prompt | llm.with_structured_output(CategoriesOutput)

# Classification Grading Assistant Prompt
classification_grading_prompt = ChatPromptTemplate.from_messages(
    [
        ("system",
        """
You are an expert commercial underwriting assistant tasked with grading the accuracy between a given business description and the classifications assigned.\n
Apply logical reasoning, grade the accuracy between the business description and the categories assigned.\n
The objective is to match the business description with the most relevant categories.\n
Business Description:\n
{description}\n
Category Description and Rationalization:\n
{cat_and_rat}
""",
         )
    ]
)

# Runnable
classification_grading_runnable = classification_grading_prompt | llm.with_structured_output(DescriptionGradeSchema)

# Retrieve Node
def retrieve(state: MainState) -> RAGState:
    """
    The function will first generate a detailed description of the business and then retrieve relevant documents.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    print("---Rewrite Description---")
    nature_of_ops = state["business_information"]
    description = retriever_runnable.invoke({"business_information": nature_of_ops})
    print(f"Description:\n {description}")

    print("---RETRIEVE---")
    # Retrieval
    documents = retriever.invoke(description)
    print(f"Documents:\n {documents}")
    return {"description": description, "documents": documents}

# Reasoning Node
def reasoning(rag_state: RAGState) -> RAGState:
    """
    The function provides reasoning for the business categories based on the provided description and retrieved documents.

    Args:
        rag_state (RAGState): The current state containing the description and retrieved documents.

    Returns:
        RAGState: Updated state with reasoned categories of the business.
    """
    print("---REASONING---")
    description = rag_state["description"]
    documents = rag_state["documents"]

    # Reasoning
    rationale = reasoning_runnable.invoke({"description": description, "documents": documents})
    print(f"Rationale:\n {rationale}")
    return {"rationale": rationale}

# Classification Grading Node
def classification_grading(rag_state: RAGState) -> ExtraState:
    """
    Grade each category based on how well it matches the business description.
    """
    print("\n=== CLASSIFICATION GRADING START ===")
    
    # Extract core data
    description = rag_state["description"]
    rationale = rag_state["rationale"]
    categories = rationale.categories
    
    print(f"Processing {len(categories)} categories for grading...")
    
    try:
        # Format categories into a single string
        cat_and_rat = "\n\n".join([
            f"Category: {cat.category}\n"
            f"Description: {cat.cat_description}\n"
            f"Code: {cat.code}\n"
            f"Rationale: {cat.rationalization}\n"
            for cat in categories
        ])
        
        graded_categories = classification_grading_runnable.invoke({
            "description": description,
            "cat_and_rat": cat_and_rat
        })
        
        print(f"\nCompleted grading {len(graded_categories.description_grades)} categories")
        
    except Exception as e:
        print(f"Error during classification grading: {str(e)}")
        raise
    
    print("=== CLASSIFICATION GRADING END ===\n")
    return {"graded_categories": graded_categories.description_grades}

from langgraph.prebuilt import tools_condition
from typing import Literal, Callable

# Function to determine the next route for the main assistant
def route_main_assistant(
    state: MainState,
) -> Literal[
    "main_assistant_tools",
    "retrieve",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    if tool_calls:
        if tool_calls[0]["name"] == ToUnderwritingAssistant.__name__:
            return "pass_tool_call_id"
        else:
            return "main_assistant_tools"
    raise ValueError("Invalid route")

# Function to determine the next route for the underwriting assistant
def route_underwriting_assistant(
    state: MainState,
) -> Literal[

    "update_workflow_state",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls

    if tool_calls:
        if tool_calls[0]["name"] == ToQuoteAssistant.__name__:
            return "update_workflow_state"

    raise ValueError("Invalid route")

# Function to determine the next route for the quote assistant
def route_quote_assistant(
    state: MainState,
) -> Literal[
    "quote_assistant_tools",
    "reroute",
    "__end__",
]:
    route = tools_condition(state)
    if route == END:
        return END

    tool_calls = state["messages"][-1].tool_calls
    
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "reroute"
    
    # Get tool names using the name property for StructuredTools
    tool_names = [t.name for t in quote_assistant_tools]
    
    if all(tc["name"] in tool_names for tc in tool_calls):
        return "quote_assistant_tools"
        
    raise ValueError(f"Invalid route. Tool call {[tc['name'] for tc in tool_calls]} not found in available tools {tool_names}")

# Function to create an entry node for quote conversation
def create_entry_node_quote(assistant_name: str, new_workflow_state: str) -> Callable:
    def entry_node(state: MainState):
        tool_call_id = state["messages"][-1].tool_calls[0]["id"]
        final_classifications = state["messages"][-1].tool_calls[0]["args"]["final_classifications"]
        business_information = state.get("business_information", {})  # Changed from [] to {}
        
        # Create a new dictionary combining both pieces of information
        updated_business_info = {
            **business_information,
            "final_classifications": final_classifications
        }
        return {
            "messages": [
                ToolMessage(
                    content=f"The assistant is now {assistant_name}. You will proceed to provide a quote using the final classifications related to the business in evaluation."
                    f" Final Classifications/Categories: {final_classifications}."
                    " Do not mention who you are - just act as the proxy for the assistant.",
                    tool_call_id=tool_call_id,
                )
            ],
            "business_information": updated_business_info,
            "workflow_state": new_workflow_state
        }

    return entry_node

# Function to return to the workflow
def route_to_workflow(
    state: MainState,
) -> Literal[
    "main_assistant",
    "underwriting_assistant",
    "quote_assistant",
]:
    """If we are in a workflow state, route directly to the appropriate assistant."""
    workflow_state = state.get("workflow_state")
    if not workflow_state:
        return "main_assistant"
    return workflow_state

# Function to reroute the workflow to the main assistant
def reroute(state: MainState) -> dict:
    """Reroute the workflow to the main assistant."""
    messages = []
    if state["messages"][-1].tool_calls:
        messages.append(
            ToolMessage(
                content="Resuming dialog with the main assistant.",
                tool_call_id=state["messages"][-1].tool_calls[0]["id"],
            )
        )
    return {
        "workflow_state": "pop",
        "messages": messages,
    }

# Function to update the workflow state
def update_workflow_state(state: MainState) -> MainState:
    """Reroute the workflow to the main assistant."""
    return {
        "workflow_state": "pop",
    }

# Function to pass the tool call ID to ExtraState
def pass_tool_call_id(state: MainState) -> ExtraState:
    tool_call_id = state["messages"][-1].tool_calls[0]["id"]
    return {"tool_call_id": tool_call_id}

# Function to pass the final classifications to ExtraState
def pass_final_classifications(rag_state: RAGState) -> ExtraState:
    final_classifications = rag_state["graded_categories"]
    return {"final_classifications": final_classifications}

# Function to create a tool message and pass to MainState
def create_tool_message(extra_state: ExtraState) -> MainState:
    tool_call_id = extra_state["tool_call_id"]
    final_classifications = extra_state["final_classifications"]
    return {
            "messages": [
                ToolMessage(
                    content=""f"Assume the role of the Main Underwriting Assistant and reflect on the previous conversation between the host assistant and the user."
                    f" Based on the initial underwriting analysis, the following categories have been identified as most relevant to the business:"
                    f"{final_classifications}"
                    "Grades have been assigned based on the accuracy of the descriptions and the rationale provided for each category. KEEP THE GRADES PRIVATE AND DO NOT SHARE THEM WITH THE USER."
                    "To confirm their relevance to the business, please ask the user follow-up questions to finalize the validity of these categories.""",
                    tool_call_id=tool_call_id,
                )
            ],
            "workflow_state": "underwriting_assistant"
        }

# Function to update the state with the new business information from tool calls
def update_state(state: MainState) -> MainState:
    """
    The function updates the state with the new business information from tool calls.
    """
    messages = state.get("messages", [])
    
    # Find the most recent tool message
    tool_messages = [msg for msg in messages if msg.type == "tool"]
    if tool_messages:
        latest_tool_message = tool_messages[-1]
        # Parse the JSON content from the tool message
        import json
        try:
            tool_content = json.loads(latest_tool_message.content)
            business_info = tool_content.get("business_information", {})
        except json.JSONDecodeError:
            business_info = state.get("business_information", {})
    else:
        business_info = state.get("business_information", {})

    return {
        "messages": messages,
        "business_information": business_info
    }

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, StateGraph, END

# Graph State
graph = StateGraph(MainState)

## Add Nodes
# Main Assistant Node
graph.add_node("main_assistant", Assistant(main_assistant_runnable))
# Main Assistant Tools Node
graph.add_node("main_assistant_tools", create_tool_node_with_fallback(main_assistant_tool))
# Underwriting Assistant Node
graph.add_node("underwriting_assistant", Assistant(underwriting_assistant_runnable))
# Quote Assistant Node
graph.add_node("quote_assistant", Assistant(quote_assistant_runnable))
# Quote Assistant Tools Node
graph.add_node("quote_assistant_tools", create_tool_node_with_fallback(quote_assistant_tools))
# Entry Quote Assistant Node
graph.add_node(
    "entry_quote_assistant",
    create_entry_node_quote("Main Quote Assistant", "quote_assistant")
)
# Retrieve Node
graph.add_node("retrieve", retrieve)
# Reasoning Node
graph.add_node("reasoning", reasoning)
# Classification Grading Node
graph.add_node("classification_grading", classification_grading)
# Update State Node
graph.add_node("update_state", update_state)
# Reroute Node
graph.add_node("reroute", reroute)
# Pass Tool Call ID Node
graph.add_node("pass_tool_call_id", pass_tool_call_id)
# Pass Final Classifications Node
graph.add_node("pass_final_classifications", pass_final_classifications)
# Create Tool Message Node
graph.add_node("create_tool_message", create_tool_message)
# Update Workflow State Node
graph.add_node("update_workflow_state", update_workflow_state)

## Add Edges
# Set Conditional Entry Point
graph.set_conditional_entry_point(
    route_to_workflow,
    {
        "main_assistant": "main_assistant",
        "underwriting_assistant": "underwriting_assistant",
        "quote_assistant": "quote_assistant",
    }
)

# Conditional Edges from main assistant
graph.add_conditional_edges(
    "main_assistant",
    route_main_assistant,
    [
        "main_assistant_tools",
        "pass_tool_call_id",
        END,
    ],
)

# Conditional Edges from underwriting assistant
graph.add_conditional_edges(
    "underwriting_assistant",
    route_underwriting_assistant,
    [
        "reroute",
        "update_workflow_state",
        END,
    ],
)

# Conditional Edges from quote assistant
graph.add_conditional_edges(
    "quote_assistant",
    route_quote_assistant,
    [
        "quote_assistant_tools",
        "reroute",
        END,
    ],
)

# Edges from update workflow state to entry quote assistant
graph.add_edge("update_workflow_state", "entry_quote_assistant")
# Edges from pass tool call id to retrieve
graph.add_edge("pass_tool_call_id", "retrieve")
# Edges from retrieve to reasoning
graph.add_edge("retrieve", "reasoning")
# Edges from reasoning to classification grading
graph.add_edge("reasoning", "classification_grading")
# Edges from classification grading to pass final classifications
graph.add_edge("classification_grading", "pass_final_classifications")
# Edges from pass final classifications to create tool message  
graph.add_edge("pass_final_classifications", "create_tool_message")
# Edges from create tool message to underwriting assistant
graph.add_edge("create_tool_message", "underwriting_assistant")
# Edges from update state to main assistant
graph.add_edge("main_assistant_tools", "update_state")
# Edges from update state to main assistant
graph.add_edge("update_state", "main_assistant")
# Edges from reroute to main assistant
graph.add_edge("reroute", "main_assistant")
# Edges from entry quote assistant to quote assistant
graph.add_edge("entry_quote_assistant", "quote_assistant")
# Edges from quote assistant tools to quote assistant
graph.add_edge("quote_assistant_tools", "quote_assistant")

# Memory Saver
memory = MemorySaver()

# Compile Graph
final_graph = graph.compile(
    checkpointer=memory,
)

import uuid
from langchain.schema import AIMessage
from langchain_core.messages import ToolMessage 

# Function to print the event
def _print_event(event: dict, _printed: set, max_length=1500):
    current_state = event.get("dialog_state")
    if current_state:
        print(f"Currently in: ", current_state[-1])
    message = event.get("messages")
    if message:
        if isinstance(message, list):
            message = message[-1]
        if message.id not in _printed:
            msg_repr = message.pretty_repr(html=True)
            if len(msg_repr) > max_length:
                msg_repr = msg_repr[:max_length] + " ... (truncated)"
            print(msg_repr)
            _printed.add(message.id)

thread_id = str(uuid.uuid4())

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

_printed = set()

# Function to process questions
def process_questions():

    global _printed  # Ensure the use of the global _printed variable
    _printed = set()  # Re-initialize if needed for the function's scope
    input_count = 0
    max_inputs = 20

    while input_count < max_inputs:
        # Get user input for the question
        question = input("\nEnter your question (or type '/exit' to quit): ")
        if question.lower() == '/exit':
            break

        # Process the question
        events = final_graph.stream(
            {"messages": ("user", question)}, config=config, stream_mode="values"
        )
        for event in events:
            _print_event(event, _printed)
        snapshot = final_graph.get_state(config)

        while snapshot.next:
            # We have an interrupt! The agent is trying to use a tool, and the user can approve or deny it
            user_input = input(
                "Do you approve of the above actions? Type 'y' to continue; "
                "otherwise, explain your requested changes.\n\n"
            )
            if user_input.strip().lower() == "y":
                # Just continue
                result = final_graph.invoke(
                    None,
                    config
                )
            else:
                # Satisfy the tool invocation by providing instructions on the requested changes / change of mind
                result = final_graph.invoke(
                    {
                        "messages": [
                            ToolMessage(  # Ensure ToolMessage is imported or defined
                                tool_call_id=event["messages"][-1].tool_calls[0]["id"],
                                content=f"API call denied by user. Reasoning: '{user_input}'. Continue assisting, accounting for the user's input.",
                            )
                        ]
                    },
                    config,
                )
                
            # Filter and print the last AIMessage
            aimessages = [message for message in result["messages"] if isinstance(message, AIMessage)]
            last_aimessage = aimessages[-1] if aimessages else None
            if last_aimessage is not None:
                # print(last_aimessage.content)
                _print_event(result, _printed)
                
            snapshot = final_graph.get_state(config)

        input_count += 1
        if input_count >= max_inputs:
            print("Maximum number of inputs reached. Terminating the loop.")
            break

process_questions()