import os
from typing import Dict, List

from langchain.agents import AgentExecutor, create_agent
from langchain_ollama import Ollama
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage
from langchain.tools import Tool

from src.utils import search_course_docs
from src.mcp_tool import fetch_course_meta

# Define tool functions

def chroma_tool(query: str) -> List[str]:
    """Search FAQ documents in Chroma. Use when question about course content."""
    return search_course_docs(query)


def mcp_tool(query: str) -> Dict:
    """Get course metadata such as schedule. Use when question about schedule or instructor."""
    return fetch_course_meta(query)

# Create Tool objects
TOOL_DEFINITIONS = [
    Tool.from_function(
        func=chroma_tool,
        name="search_course_docs",
        description="Search FAQ documents in Chroma. Use when question about course content.",
    ),
    Tool.from_function(
        func=mcp_tool,
        name="fetch_course_meta",
        description="Get course metadata such as schedule. Use when question about schedule or instructor.",
    ),
]

# Prompt template
SYSTEM_PROMPT = (
    "You are a FAQ bot for the course. Use only the tools provided. Do not call both tools unless necessary."
    " Indicate source in your answer: 'source: chroma' or 'source: mcp_meta'."
)

PROMPT = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("human", "{input}"),
])

# Create agent
llm = Ollama(model="llama2")
agent = create_agent(llm, TOOL_DEFINITIONS, PROMPT)
executor = AgentExecutor(agent=agent, tools=TOOL_DEFINITIONS, verbose=True)

# CLI entry point
if __name__ == "__main__":
    import argparse
    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:
        result = executor.invoke({"input": args.question})
        print(result["output"])
    else:
        print("Enter questions (Ctrl-D to exit):")
        try:
            while True:
                q = input("Q: ")
                res = executor.invoke({"input": q})
                print("A:", res["output"])
        except EOFError:
            pass