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

# ---------------------------------------------------------------------------
# 1. Настройка LLM и Embeddings (OpenRouter)
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

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

# ---------------------------------------------------------------------------
# 2. ChromaDB: загрузка .md файлов и поиск
# ---------------------------------------------------------------------------
CHROMA_PATH = Path("./chroma_faq")
CHROMA_COLLECTION = "faq_collection"

vector_store = Chroma(
    collection_name=CHROMA_COLLECTION,
    embedding_function=embeddings,
    persist_directory=str(CHROMA_PATH),
)

# Если коллекция пуста, загрузим данные из data/*.md
if not vector_store.get_collection().list_documents():
    md_files = list(Path("data").glob("*.md"))
    docs: List[Document] = []
    for f in md_files:
        text = f.read_text(encoding="utf-8")
        docs.append(Document(page_content=text, metadata={"source": f.name}))
    vector_store.add_documents(docs)
    vector_store.persist()

@tool
def search_course_docs(query: str, k: int = 3) -> str:
    """Search the local FAQ collection for relevant passages."""
    results = vector_store.similarity_search(query, k=k)
    if not results:
        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 results)

# ---------------------------------------------------------------------------
# 3. MCP‑style tool (mocked via local JSON file)
# ---------------------------------------------------------------------------
META_JSON = Path("course_meta.json")
if not META_JSON.exists():
    # Создаём простую статическую мета‑информацию
    META_JSON.write_text(json.dumps({
        "schedule": {
            "Monday": "Lecture 1",
            "Wednesday": "Lecture 2",
            "Friday": "Lab"
        },
        "instructor": "Dr. Example"
    }, indent=2))

@tool
def fetch_course_meta(query: str) -> str:
    """Return course metadata that matches the query.
    The function simply looks for the query string in the keys of the JSON.
    """
    data = json.loads(META_JSON.read_text())
    for key, value in data.items():
        if query.lower() in key.lower():
            return json.dumps({key: value}, indent=2)
    return "No metadata found for the given query."

# ---------------------------------------------------------------------------
# 4. DeepAgent с маршрутизацией
# ---------------------------------------------------------------------------
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"
        "If the question is about course content, use the search_course_docs tool.\n"
        "If the question is about schedule, instructor, or other metadata, use fetch_course_meta.\n"
        "Do not call both tools unless absolutely necessary.\n"
        "In your final answer, prepend 'source: chroma' or 'source: mcp_meta' to indicate where the answer came from."
    ),
)

# ---------------------------------------------------------------------------
# 5. CLI
# ---------------------------------------------------------------------------
PRESET_QUESTIONS = [
    "What topics are covered in Lecture 1?",  # should hit chroma
    "Who is the instructor for this course?",  # should hit mcp_meta
    "Explain the concept of polymorphism."
]

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

async def main():
    print("\n--- FAQ Bot Demo ---\n")
    for i, q in enumerate(PRESET_QUESTIONS, 1):
        print(f"Q{i}: {q}")
        ans = await run_agent(q)
        print(f"A{i}: {ans}\n")

    print("Enter your own question (or press Ctrl+C to exit):")
    while True:
        try:
            user_q = input("> ")
            if not user_q.strip():
                continue
            ans = await run_agent(user_q)
            print(ans)
        except KeyboardInterrupt:
            print("\nExiting.")
            break

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