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

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# --- Embeddings and vector store (Qdrant + Ollama) ---
from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import Qdrant
from langchain_core.documents import Document

# Load OpenRouter key for LLM (required by Ollama embeddings as well)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# LLM configuration (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 Ollama (nomic-embed-text)
embeddings = OllamaEmbeddings(model="nomic-embed-text")

# Qdrant client (in‑memory for demo; replace with host/port for prod)
qdrant_client = Qdrant(
    collection_name="faq_collection",
    url="http://localhost:6333",  # default Qdrant local URL
    embedding_function=embeddings,
)

# --- Data loading and indexing ---
DATA_DIR = Path("data")
CHROMA_PERSIST = Path("./qdrant_faq")  # not used directly but kept for compatibility


def load_faq_to_qdrant() -> None:
    """Read all .md files from DATA_DIR, chunk them, embed and store in Qdrant."""
    if not DATA_DIR.exists():
        print("Data directory not found. Create 'data/' with .md files.")
        return
    docs: List[Document] = []
    for md_file in DATA_DIR.glob("*.md"):
        text = md_file.read_text(encoding="utf-8")
        # Simple split by double newlines as a naive chunker
        for i, chunk in enumerate(text.split("\n\n")):
            docs.append(Document(page_content=chunk, metadata={"source": md_file.name, "chunk": i}))
    # Add to Qdrant
    qdrant_client.add_documents(docs)
    print(f"Indexed {len(docs)} chunks into Qdrant.")

# --- Tools ---
@tool
def search_course_docs(query: str, k: int = 3) -> str:
    """Search the local FAQ collection for relevant passages."""
    results = qdrant_client.similarity_search(query, k=k)
    if not results:
        return "No relevant information found in the course materials."
    return "\n\n---\n\n".join(r.page_content for r in results)

# MCP-style tool: fetch metadata from a static JSON file
META_FILE = Path("meta.json")

@tool
def fetch_course_meta(query: str) -> str:
    """Return course metadata that matches the query.
    For demo purposes, we load a static JSON file and perform a simple keyword search.
    """
    if not META_FILE.exists():
        return "Metadata file not found."
    data = json.loads(META_FILE.read_text(encoding="utf-8"))
    # Very naive matching: return entries where query is a substring of any value
    matches = []
    for key, value in data.items():
        if isinstance(value, str) and query.lower() in value.lower():
            matches.append(f"{key}: {value}")
    if not matches:
        return "No metadata matches your query."
    return "\n".join(matches)

# --- Agent setup ---
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 bot for a course.\n"
        "When a user asks about course materials, use the search_course_docs tool.\n"
        "When a user asks about schedule, metadata, or other non‑material info, use fetch_course_meta.\n"
        "Do not call both tools unless absolutely necessary.\n"
        "In your final answer, prepend 'source: chroma' if you used search_course_docs,\n"
        "or 'source: mcp_meta' if you used fetch_course_meta."
    ),
)

# --- CLI ---
PRESET_QUESTIONS = [
    "What topics are covered in the first lecture?",
    "How can I access the lecture slides?",
    "What is the schedule for the next week?",
]

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 the vector store is populated
    load_faq_to_qdrant()

    print("\n--- FAQ Bot Demo ---\n")
    for i, q in enumerate(PRESET_QUESTIONS, 1):
        print(f"Q{i}: {q}")
        answer = await run_agent(q, thread_id=f"demo-{i}")
        print(f"A{i}: {answer}\n")

    # Interactive mode
    print("Enter your own questions (type 'exit' to quit):")
    while True:
        user_q = input("> ")
        if user_q.lower() in {"exit", "quit"}:
            break
        answer = await run_agent(user_q, thread_id="interactive")
        print(answer)

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