from typing import TypedDict, List, Dict

# ==========================================
# 1. Define the Shared State
# ==========================================
class AgentState(TypedDict):
    user_input: str
    intents: List[str]
    research_plan: List[str]
    gathered_data: List[str]
    final_response: str

# ==========================================
# 2. The Planner Agent (Your Code, slightly integrated)
# ==========================================
class PlannerAgent:
    def __init__(self):
        self.intent_keywords = {
            "course": ["course", "btech", "cse", "syllabus", "duration"],
            "facilities": ["hostel", "library", "wifi", "canteen"],
            "events": ["event", "hackathon", "placement", "deadline"],
            "support": ["help", "scholarship", "counseling"]
        }

    def detect_intents(self, query: str) -> List[str]:
        query = query.lower()
        detected = [intent for intent, keywords in self.intent_keywords.items() 
                    if any(word in query for word in keywords)]
        return detected if detected else ["general"]

    def create_plan(self, intents: List[str]) -> List[str]:
        plan = []
        for intent in intents:
            if intent == "course":
                plan.extend(["Search Course DB", "Get Syllabus"])
            elif intent == "facilities":
                plan.extend(["Search Campus DB", "Get Hostel Rules"])
            elif intent == "events":
                plan.extend(["Search Calendar DB", "Get Placement Dates"])
            elif intent == "support":
                plan.extend(["Search Contact DB", "Get Scholarship Form"])
            else:
                plan.append("Search General FAQ")
        return plan

    def run(self, state: AgentState) -> AgentState:
        print("--- 🧠 PLANNER NODE RUNNING ---")
        intents = self.detect_intents(state["user_input"])
        plan = self.create_plan(intents)
        
        state["intents"] = intents
        state["research_plan"] = plan
        return state

# ==========================================
# 3. The Executor Agent (Gathers Information)
# ==========================================
class ExecutorAgent:
    """Simulates searching a university database based on the plan."""
    def __init__(self):
        # Mock database
        self.knowledge_base = {
            "Search Course DB": "B.Tech CSE is a 4-year undergraduate program.",
            "Get Syllabus": "Syllabus includes Data Structures, OS, and AI.",
            "Search Campus DB": "Campus has 24/7 Wi-Fi and a central library.",
            "Get Hostel Rules": "Hostel curfew is 10:00 PM for all students.",
            "Search Calendar DB": "Next Hackathon is on October 15th.",
            "Get Placement Dates": "Placement drives begin in the 7th semester.",
            "Search General FAQ": "Welcome to the University Enquiry Desk."
        }

    def run(self, state: AgentState) -> AgentState:
        print("--- 🔍 EXECUTOR NODE RUNNING ---")
        gathered_data = []
        
        for task in state["research_plan"]:
            # Retrieve data if it exists in our mock DB, otherwise return a default string
            data = self.knowledge_base.get(task, f"No data found for: {task}")
            gathered_data.append(data)
            
        state["gathered_data"] = gathered_data
        return state

# ==========================================
# 4. The Generator Agent (Drafts the Final Reply)
# ==========================================
class GeneratorAgent:
    """Takes the gathered data and formats it for the user."""
    
    def run(self, state: AgentState) -> AgentState:
        print("--- ✍️ GENERATOR NODE RUNNING ---")
        
        if not state["gathered_data"]:
            state["final_response"] = "I'm sorry, I couldn't find any information on that."
            return state
            
        # Format the gathered facts into a neat bulleted list
        response = "Here is the information I found for you:\n"
        for fact in state["gathered_data"]:
            response += f"• {fact}\n"
            
        state["final_response"] = response
        return state

# ==========================================
# 5. The Graph Orchestrator (The Main Loop)
# ==========================================
def run_pipeline(user_query: str):
    # Initialize the state dictionary
    state: AgentState = {
        "user_input": user_query,
        "intents": [],
        "research_plan": [],
        "gathered_data": [],
        "final_response": ""
    }
    
    # Initialize our agents
    planner = PlannerAgent()
    executor = ExecutorAgent()
    generator = GeneratorAgent()
    
    # Pass the state sequentially through the nodes
    state = planner.run(state)
    state = executor.run(state)
    state = generator.run(state)
    
    return state

# ==========================================
# Execution Block
# ==========================================
if __name__ == "__main__":
    print("Welcome to the University Bot. Type 'exit' to quit.\n")
    
    while True:
        query = input("You: ")
        if query.lower() in ['exit', 'quit']:
            break
            
        print("\nProcessing...\n")
        
        final_state = run_pipeline(query)
        
        print("\n==========================================")
        print("BOT RESPONSE")
        print("==========================================")
        print(final_state["final_response"])
        print("\n" + "-"*42 + "\n")