import argparse
from langchain.agents.openai_functions import create_openai_functions_agent
from langchain.agents import AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import Ollama
from src.utils import load_faq_to_chroma, search_course_docs, fetch_course_meta

# Initialize embeddings and LLM
llm = Ollama(model="llama3.1")

# Load or create Chroma collection
try:
    chroma = load_faq_to_chroma()
except Exception:
    chroma = None

# Define tools
from langchain.tools import tool

@tool
def search_course_docs_tool(query: str, k: int = 3) -> str:
    """Search local FAQ docs in ChromaDB."""
    docs = search_course_docs(query, k)
    return "\n".join([doc.page_content for doc in docs])

@tool
def fetch_course_meta_tool(query: str) -> str:
    """Fetch course metadata via MCP-style tool."""
    results = fetch_course_meta(query)
    return str(results)

tools = [search_course_docs_tool, fetch_course_meta_tool]

# Prompt template with source hint
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful FAQ assistant. Use the tools only when necessary. In your answer, include a line like 'source: chroma' or 'source: mcp_meta' to indicate which tool was used.")
])

agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt)
executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="FAQ bot CLI")
    parser.add_argument("--question", type=str, help="Question to ask the bot")
    args = parser.parse_args()
    if args.question:
        response = executor.invoke({"input": args.question})
        print(response["output"])
    else:
        # Interactive mode
        print("FAQ Bot. Type 'exit' to quit.")
        while True:
            q = input("> ")
            if q.lower() in ("exit", "quit"):
                break
            resp = executor.invoke({"input": q})
            print(resp["output"])