import os
import asyncio
from pathlib import Path
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# ===================== CONFIG =====================
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY is 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 = Chroma(
    collection_name="faq_knowledge",
    embedding_function=embeddings,
    persist_directory="./chroma_faq",
)

# ===================== TOOL: SEARCH IN CHROMA =====================
@tool
def search_course_docs(query: str, k: int = 3) -> str:
    """Search the local FAQ knowledge base for relevant passages."""
    docs: list[Document] = 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"**{d.metadata.get('title', 'Untitled')}**\n{d.page_content}" for d in docs)

# ===================== TOOL: FETCH METADATA (MCP‑STYLE) =====================
# For the purpose of this assignment we use a static JSON file as the mock MCP server response.
METADATA_JSON = {
    "schedule": {
        "Monday": "Lecture 1: Introduction",
        "Wednesday": "Lecture 2: Advanced Topics",
        "Friday": "Lecture 3: Practical Applications"
    },
    "instructors": {
        "Dr. Smith": "smith@example.com",
        "Prof. Doe": "doe@example.com"
    }
}

@tool
def fetch_course_meta(query: str) -> str:
    """Mock MCP tool that returns course metadata based on the query.
    In production this would be an HTTP GET to an MCP server.
    """
    query = query.lower()
    if "schedule" in query:
        return "\n".join(f"{day}: {info}" for day, info in METADATA_JSON["schedule"].items())
    if "instructor" in query or "email" in query:
        return "\n".join(f"{name}: {email}" for name, email in METADATA_JSON["instructors"].items())
    return "No metadata matches your query."

# ===================== LOAD FAQ TO CHROMA =====================
MD_DIR = Path("data")
if not MD_DIR.exists():
    MD_DIR.mkdir(parents=True, exist_ok=True)
    # Create example markdown files if none exist
    (MD_DIR / "faq1.md").write_text("# FAQ 1\nWhat is the course about?\nThe course covers advanced AI techniques.")
    (MD_DIR / "faq2.md").write_text("# FAQ 2\nHow to install dependencies?\nRun `pip install -r requirements.txt`.")
    (MD_DIR / "faq3.md").write_text("# FAQ 3\nWhere to find the schedule?\nCheck the course website.")

# Chunking and persisting
from langchain_text_splitters import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)

for md_file in MD_DIR.glob("*.md"):
    content = md_file.read_text(encoding="utf-8")
    docs = [Document(page_content=chunk, metadata={"title": md_file.stem}) for chunk in text_splitter.split_text(content)]
    vector_store.add_documents(docs)

vector_store.persist()

# ===================== 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,
)

backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

agent = create_deep_agent(
    model=llm,
    tools=[search_course_docs, fetch_course_meta],
    backend=backend,
    system_prompt="You are a helpful FAQ assistant. Use only the provided tools. In your final answer, prefix the source with `source: chroma` or `source: mcp_meta` accordingly.",
)

# ===================== CLI =====================
PRESET_QUESTIONS = [
    "What is the course about?",  # chroma
    "How to install dependencies?",  # chroma
    "What is the lecture schedule?",  # mcp_meta
]

async def run_cli():
    print("=== FAQ Assistant ===")
    print("Type your question or 'exit' to quit.")
    for i, q in enumerate(PRESET_QUESTIONS, 1):
        print(f"\nPreset {i}: {q}")
        await handle_question(q)
    while True:
        user_input = input("\nYour question: ")
        if user_input.lower() in {"exit", "quit"}:
            break
        await handle_question(user_input)

async def handle_question(question: str):
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=question)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    answer = result["messages"][-1].content
    print("\nAnswer:\n", answer)

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