import os
from typing import List, Dict, Any

from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tools import BaseTool
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.schema import HumanMessage, SystemMessage

from .tools import search_course_docs, fetch_course_meta

# Load tools
TOOLS: List[BaseTool] = [search_course_docs, fetch_course_meta]

# System prompt guiding the agent
SYSTEM_PROMPT = """
You are a helpful assistant for a machine learning course. Your job is to answer user questions.

- If the question is about course materials, lecture slides, assignments, or any content that can be found in the FAQ documents, use the tool `search_course_docs`.
- If the question is about course schedule, instructor information, or other metadata, use the tool `fetch_course_meta`.
- Do not use both tools unless absolutely necessary.
- In your answer, always include a source tag: `source: chroma` if you used the FAQ tool, or `source: mcp_meta` if you used the metadata tool.
"""

def build_agent() -> AgentExecutor:
    """
    Build and return a LangChain AgentExecutor with the defined tools and system prompt.
    """
    llm = Ollama(model="llama3", temperature=0.0)

    # Prompt template
    prompt = ChatPromptTemplate.from_messages(
        [
            SystemMessage(content=SYSTEM_PROMPT),
            MessagesPlaceholder(variable_name="history"),
            HumanMessage(content="{input}"),
        ]
    )

    # Create the agent
    agent = create_openai_tools_agent(llm=llm, tools=TOOLS, prompt=prompt)

    # Wrap with AgentExecutor
    agent_executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True, handle_parsing_errors=True)
    return agent_executor