import os
import asyncio
import json
import httpx
from pathlib import Path
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain.tools import tool
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.agents import Tool
from langchain_community.utilities import RetrievalQA
from langchain_community.vectorstores import Chroma as ChromaStore

# ---------- Configuration ----------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY not set in environment")

# ---------- Embeddings & Vector Store ----------
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
)

vector_store = ChromaStore(
    collection_name="faq_collection",
    embedding_function=embeddings,
    persist_directory="./chroma_faq",
)

# ---------- Load FAQ into Chroma ----------

def load_faq_to_chroma(data_dir: str = "data"):
    """Load all .md files from data_dir into the Chroma vector store.
    The function clears the existing collection before loading.
    """
    vector_store.delete_collection()
    docs = []
    for md_file in Path(data_dir).glob("*.md"):
        text = md_file.read_text(encoding="utf-8")
        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
    vector_store.add_documents(docs)
    vector_store.persist()

# ---------- Tools ----------
@tool
def search_course_docs(query: str, k: int = 3) -> str:
    """Search the local FAQ collection for relevant passages."""
    docs = vector_store.similarity_search(query, k=k)
    if not docs:
        return "No relevant information found in the course materials."
    return "\n\n---\n\n".join([f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc in docs])

@tool
def fetch_course_meta(query: str) -> str:
    """MCP‑style tool that queries a local mock server for course metadata.
    The mock server should serve a JSON file at http://localhost:8000/meta.json.
    """
    url = "http://localhost:8000/meta.json"
    try:
        response = httpx.get(url, timeout=5.0)
        response.raise_for_status()
        data = response.json()
    except Exception as e:
        return f"Error fetching metadata: {e}"
    # Simple lookup: return value if query matches a key (case‑insensitive)
    key = query.strip().lower()
    value = data.get(key)
    if value is None:
        return f"No metadata entry found for '{query}'."
    return f"{key}: {value}"

# ---------- Agent Setup ----------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENAI_API_KEY,
    temperature=0.0,
)

# Define the system prompt with routing rule
system_prompt = (
    "You are a helpful assistant that answers questions about the course. "
    "If the question is about course content, use the search_course_docs tool. "
    "If the question is about schedule, metadata, or other non‑content info, "
    "use the fetch_course_meta tool. Do not call both tools unless the question "
    "explicitly requires both. In your final answer, prepend 'source: chroma' "
    "or 'source: mcp_meta' to indicate which tool provided the information."
)

# Create Tool objects
search_tool = Tool(name="search_course_docs", func=search_course_docs, description="Search local course documents.")
meta_tool = Tool(name="fetch_course_meta", func=fetch_course_meta, description="Fetch course metadata from MCP mock server.")

# Build the agent executor
agent = create_openai_tools_agent(llm=llm, tools=[search_tool, meta_tool], system_message=system_prompt)
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=[search_tool, meta_tool], verbose=True)

# ---------- CLI ----------
async def run_cli():
    # Preload data if not already present
    if not Path("./chroma_faq").exists():
        load_faq_to_chroma()

    # Predefined questions
    predefined = [
        "What is the main topic of the first lecture?",
        "Explain the concept of polymorphism in the course.",
        "What is the schedule for the next week?",
    ]
    print("--- Predefined questions ---")
    for i, q in enumerate(predefined, 1):
        print(f"{i}. {q}")
    print("\nEnter a number to ask a predefined question or type your own query.")
    user_input = input("> ")
    if user_input.isdigit() and 1 <= int(user_input) <= len(predefined):
        query = predefined[int(user_input)-1]
    else:
        query = user_input

    result = await agent_executor.ainvoke({"input": query})
    print("\n--- Answer ---")
    print(result["output"])

if __name__ == "__main__":
    asyncio.run(run_cli())