"""
Main entry point for the FAQ‑bot assignment.

The bot answers questions about course materials stored in a local ChromaDB
vector store and, when necessary, calls an external MCP‑style tool that
returns metadata such as the course schedule.  The implementation uses
LangChain v1+ (``langchain>=1.0.0``) together with ``langchain-chroma``
and ``langchain-ollama`` for embeddings.

The file contains:
* Utility functions to load markdown files into ChromaDB.
* Two tools – one that searches the vector store and another that mocks an
  MCP call.
* A simple LangChain agent that decides which tool to use based on the
  content of the user query.
* Three example queries executed when the script is run as ``__main__``.
"""

from __future__ import annotations

import os
import json
import pathlib
from typing import List, Dict

# ---------------------------------------------------------------------------
# Dependencies – all are declared in requirements.txt
# ---------------------------------------------------------------------------
from langchain_ollama import ChatOllama
from langchain_chroma import Chroma
from langchain_core.messages import HumanMessage
from langchain.agents import create_agent
from langchain.tools import tool
from langchain.embeddings.ollama import OllamaEmbeddings

# ---------------------------------------------------------------------------
# Configuration constants
# ---------------------------------------------------------------------------
DATA_DIR = pathlib.Path("data")  # directory with .md files
CHROMA_PATH = pathlib.Path("./chroma_faq")  # persistent store location
MCP_ENDPOINT = "http://localhost:8000/course_meta"  # mock endpoint

# ---------------------------------------------------------------------------
# Helper – load markdown files into ChromaDB
# ---------------------------------------------------------------------------
@tool(name="load_faq_to_chroma", description="Load all .md files from data/ into a persistent Chroma vector store.")
def load_faq_to_chroma() -> str:
    """Create or update the Chroma collection with local markdown files.

    The function is idempotent – it will overwrite existing documents
    with the same IDs.  It returns a short JSON string describing how many
    documents were added.
    """
    # Ensure data directory exists
    if not DATA_DIR.exists():
        return json.dumps({"error": f"{DATA_DIR} does not exist."})

    # Gather all markdown files
    md_files = list(DATA_DIR.glob("*.md"))
    if not md_files:
        return json.dumps({"error": "No .md files found in data/"})

    # Prepare documents for Chroma
    docs: List[Dict[str, str]] = []
    for file_path in md_files:
        content = file_path.read_text(encoding="utf-8")
        docs.append({"id": file_path.stem, "content": content})

    # Create embeddings and store
    embedding = OllamaEmbeddings(model="nomic-embed-text")
    chroma = Chroma.from_documents(
        documents=[doc["content"] for doc in docs],
        ids=[doc["id"] for doc in docs],
        embedding=embedding,
        persist_directory=str(CHROMA_PATH),
    )

    return json.dumps({"status": "loaded", "documents": len(docs)})

# ---------------------------------------------------------------------------
# Tool – search the vector store
# ---------------------------------------------------------------------------
@tool(name="search_course_docs", description="Search local course documents for a query. Returns top k results.")
def search_course_docs(query: str, k: int = 3) -> str:
    """Return the most relevant chunks from the Chroma store.

    The function loads the persistent collection and performs a similarity
    search.  Results are returned as a numbered list of snippets.
    """
    if not CHROMA_PATH.exists():
        return json.dumps({"error": f"Chroma store {CHROMA_PATH} not found."})

    embedding = OllamaEmbeddings(model="nomic-embed-text")
    chroma = Chroma(persist_directory=str(CHROMA_PATH), embedding_function=embedding)
    docs = chroma.similarity_search(query, k=k)
    if not docs:
        return json.dumps({"result": "No relevant documents found."})

    snippets = [f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)]
    return json.dumps({"source": "chroma", "snippets": snippets})

# ---------------------------------------------------------------------------
# Tool – mock MCP call (could be replaced with a real HTTP request)
# ---------------------------------------------------------------------------
@tool(name="fetch_course_meta", description="Retrieve course metadata such as schedule. Returns JSON.")
def fetch_course_meta(query: str) -> str:
    """Simulate an external MCP tool by returning static data.

    In a real deployment this would perform an HTTP GET to the MCP server.
    For the purposes of the assignment we simply return a hard‑coded JSON
    payload that mimics what an MCP endpoint might provide.
    """
    # Static mock data – in practice replace with httpx.get(MCP_ENDPOINT)
    mock_data = {
        "schedule": [
            {"day": "Monday", "time": "10:00-12:00", "topic": "Introduction"},
            {"day": "Wednesday", "time": "14:00-16:00", "topic": "Advanced Topics"},
        ],
        "instructor": "Prof. Alexei Petrov",
    }
    return json.dumps({"source": "mcp_meta", "data": mock_data})

# ---------------------------------------------------------------------------
# Agent – decides which tool to use based on the query content
# ---------------------------------------------------------------------------
# System prompt guides the agent not to call both tools unnecessarily.
SYSTEM_PROMPT = (
    "You are an assistant that answers questions about course materials."
    " If the question is about lecture notes or FAQs, use the search_course_docs tool."
    " If the question asks for schedule, instructor info, or other metadata, use fetch_course_meta."
)

# Create the agent once – it will be reused in examples.
llm = ChatOllama(model="llama3", temperature=0.0)
agent = create_agent(
    llm=llm,
    tools=[search_course_docs, fetch_course_meta],
    system_prompt=SYSTEM_PROMPT,
)

# ---------------------------------------------------------------------------
# Example usage – three predefined queries and an interactive prompt
# ---------------------------------------------------------------------------
EXAMPLES: List[str] = [
    "What topics are covered in the first lecture?",
    "When is the next class scheduled?",
    "Can you give me a summary of the course material?",
]

def run_example(query: str) -> None:
    print(f"\nUser: {query}")
    result = agent.invoke(
        {"messages": [HumanMessage(content=query)]},
        {"configurable": {"thread_id": "example-session"}},
    )
    # The agent returns a list of messages – the last one is the answer.
    answer = result["messages"][-1].content
    print(f"Assistant: {answer}")

if __name__ == "__main__":
    # Ensure data directory exists for demo purposes
    if not DATA_DIR.exists():
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        # Create a tiny FAQ file as an example.
        (DATA_DIR / "faq.md").write_text(
            "# Course FAQ\n" \
            "## What is the course about?\n" \
            "The course covers advanced topics in AI and machine learning."
        )
    # Load data into Chroma (idempotent)
    print(load_faq_to_chroma())

    for q in EXAMPLES:
        run_example(q)

    # Interactive mode – optional
    while True:
        try:
            user_input = input("\nAsk a question (or 'quit'): ")
        except EOFError:
            break
        if user_input.lower() in {"quit", "exit"}:
            break
        run_example(user_input)
"""