import os
import asyncio
import json
from pathlib import Path
from typing import List

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

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

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

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

# ---------------------------------------------------------------------------
# Chroma persistence
# ---------------------------------------------------------------------------
CHROMA_DIR = Path("./chroma_faq")
CHROMA_DIR.mkdir(parents=True, exist_ok=True)
vector_store = Chroma(
    collection_name="faq_collection",
    embedding_function=embeddings,
    persist_directory=str(CHROMA_DIR),
)

# ---------------------------------------------------------------------------
# Utility: load markdown files into Chroma
# ---------------------------------------------------------------------------
@tool
def load_faq_to_chroma() -> str:
    """Load all .md files from data/ into the Chroma vector store.
    This tool is idempotent – it will overwrite existing collection.
    """
    data_dir = Path("data")
    if not data_dir.exists():
        return "Data directory not found."
    docs: List[Document] = []
    for md_file in data_dir.glob("*.md"):
        text = md_file.read_text(encoding="utf-8")
        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
    if docs:
        vector_store.delete_collection()
        vector_store.add_documents(docs)
        vector_store.persist()
        return f"Loaded {len(docs)} documents into Chroma."
    return "No markdown files found."

# ---------------------------------------------------------------------------
# Tool: search knowledge base
# ---------------------------------------------------------------------------
@tool
def search_course_docs(query: str, k: int = 3) -> str:
    """Search the FAQ collection for relevant passages.
    Returns a string with the top k passages and a source tag.
    """
    docs = vector_store.similarity_search(query, k=k)
    if not docs:
        return "No relevant information found in the FAQ."
    results = "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
    return f"source: chroma\n{results}"

# ---------------------------------------------------------------------------
# MCP-style tool: fetch course metadata
# ---------------------------------------------------------------------------
# For simplicity we use a static JSON file in the repo. In production this
# would be an HTTP call to an MCP server.
METADATA_FILE = Path("course_meta.json")

@tool
def fetch_course_meta(query: str) -> str:
    """Return metadata that matches the query.
    The function performs a simple keyword search in the static JSON.
    """
    if not METADATA_FILE.exists():
        return "Metadata file not found."
    data = json.loads(METADATA_FILE.read_text(encoding="utf-8"))
    matches = [item for item in data if query.lower() in item.get("title", "").lower()]
    if not matches:
        return "No metadata matches the query."
    return f"source: mcp_meta\n" + json.dumps(matches, indent=2)

# ---------------------------------------------------------------------------
# Backend setup
# ---------------------------------------------------------------------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# ---------------------------------------------------------------------------
# Agent definition
# ---------------------------------------------------------------------------
agent = create_deep_agent(
    model=llm,
    tools=[load_faq_to_chroma, search_course_docs, fetch_course_meta],
    backend=backend,
    system_prompt=(
        "You are a helpful FAQ bot for a course. "
        "Use the search_course_docs tool for questions about lecture materials. "
        "Use fetch_course_meta for questions about schedule or metadata. "
        "Do not call both tools unless necessary. "
        "Always prefix your answer with the source tag (chroma or mcp_meta)."
    ),
)

# ---------------------------------------------------------------------------
# CLI helpers
# ---------------------------------------------------------------------------
PRESET_QUESTIONS = [
    "What is the deadline for the final project?",  # should hit metadata
    "Explain the concept of polymorphism in OOP.",  # should hit FAQ
    "How many lectures are there in the first module?",  # metadata
]

async def run_agent(question: str, thread_id: str = "session-1"):
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=question)]},
        {"configurable": {"thread_id": thread_id}},
    )
    return result["messages"][-1].content

async def main():
    # Ensure FAQ is loaded once
    await agent.ainvoke(
        {"messages": [HumanMessage(content="load_faq_to_chroma()")]},
        {"configurable": {"thread_id": "init"}},
    )
    print("\n--- Preset questions ---")
    for q in PRESET_QUESTIONS:
        ans = await run_agent(q)
        print(f"Q: {q}\nA: {ans}\n")
    print("\n--- Interactive mode (type 'exit' to quit) ---")
    while True:
        user_input = input("> ")
        if user_input.lower() in {"exit", "quit"}:
            break
        ans = await run_agent(user_input)
        print(ans)

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