from src.state import AgentState
from langchain_core.messages import HumanMessage, AIMessage

def check_clarification_needed(state: AgentState) -> AgentState:
    """Check if we need to ask for missing information before proceeding"""
    print("[clarification] Starting clarification check")
    
    user_input = state.get("user_input", "").lower()
    last_order_id = state.get("last_order_id")
    
    # Check if this is a response to a prior clarification request
    prior_clarification = state.get("session_context", {}).get("last_clarification_type")
    if prior_clarification == "order_id":
        # User is responding to our request for order ID
        # The entity_extractor should have already processed this
        if last_order_id:
            # Successfully resolved, clear the flag and continue
            state["session_context"]["last_clarification_type"] = None
            state["clarification_needed"] = False
            print(f"[clarification] Clarification resolved with order_id: {last_order_id}")
            return state
    
    # Check if user is asking about "their order" without providing ID
    order_keywords = ["my order", "the order", "my package", "my delivery", "my shipment", "where is my"]
    asking_about_order = any(keyword in user_input for keyword in order_keywords)
    
    # If asking about order without ID, request clarification
    if asking_about_order and not last_order_id:
        print("[clarification] Missing order ID detected")
        state["final_response"] = (
            "I'd be happy to help you track your order! "
            "Could you please provide your Order ID? "
            "It's a 32-character reference number that looks like this: "
            "e481f51cbdc54678b7cc49136f2d6af7\n\n"
            "You can find it in your order confirmation email or account dashboard."
        )
        state["clarification_needed"] = True
        state["clarification_type"] = "order_id"
        # Store in session context for next turn
        state["session_context"]["last_clarification_type"] = "order_id"
        # Return only new messages
        return {
            **state,
            "messages": [
                HumanMessage(content=state.get("user_input", "")),
                AIMessage(content=state["final_response"])
            ]
        }
    else:
        state["clarification_needed"] = False
        print("[clarification] No clarification needed")
        return {"clarification_needed": False}
    
    return state