"""main.py

Простой FAQ‑бот, использующий ChromaDB для локальных конспектов и один MCP‑подобный инструмент – HTTP‑запрос к статическому JSON.

Требования:
- Python 3.10+
- Ollama модели: `nomic-embed-text` и `llama3`
- Пакеты из requirements.txt

Запуск:
```bash
python main.py
```

В интерактивном режиме можно задавать вопросы. В примере уже есть три готовых вопроса – два из Chroma, один из MCP‑тул.
"""

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

import httpx
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_chroma import Chroma
from langchain_core.tools import BaseTool, tool
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import Runnable
from langchain.agents import create_agent, AgentExecutor, Tool

# ---------------------------------------------------------------------------
# 1. Векторная база
# ---------------------------------------------------------------------------
CHROMA_PATH = Path("./chroma_faq")
DATA_PATH = Path("./data")

# Создаём embeddings
embeddings = OllamaEmbeddings(model="nomic-embed-text")

# ---------------------------------------------------------------------------
# Чтение markdown‑файлов и загрузка в Chroma
# ---------------------------------------------------------------------------

def load_faq_to_chroma() -> Chroma:
    """Загружает все .md файлы из DATA_PATH в Chroma.
    Если коллекция уже существует – просто возвращаем её.
    """
    if CHROMA_PATH.exists():
        return Chroma(
            collection_name="faq",
            embedding_function=embeddings,
            persist_directory=str(CHROMA_PATH),
        )
    # Если нет, создаём новую коллекцию
    from langchain_text_splitters import MarkdownTextSplitter

    splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = []
    for md_file in DATA_PATH.glob("*.md"):
        text = md_file.read_text(encoding="utf-8")
        docs.extend(splitter.split_text(text))
    vector_store = Chroma.from_texts(
        texts=docs,
        embedding=embeddings,
        collection_name="faq",
        persist_directory=str(CHROMA_PATH),
    )
    return vector_store

# ---------------------------------------------------------------------------
# 2. Tool: поиск по Chroma
# ---------------------------------------------------------------------------
@tool
def search_course_docs(query: str, k: int = 3) -> str:
    """Возвращает кортеж из k наиболее релевантных фрагментов.
    Формат ответа: ``source: chroma`` + текст.
    """
    vector_store = load_faq_to_chroma()
    results = vector_store.similarity_search(query, k=k)
    snippets = "\n---\n".join([r.page_content for r in results])
    return f"source: chroma\n{snippets}"

# ---------------------------------------------------------------------------
# 3. MCP‑подобный инструмент – HTTP‑запрос к статическому JSON
# ---------------------------------------------------------------------------
MCP_JSON_PATH = Path("./course_meta.json")

@tool
def fetch_course_meta(query: str) -> str:
    """Имитирует вызов MCP‑сервера. Читает локальный JSON и возвращает
    информацию, содержащуюся в ключе, совпадающем с query.
    """
    if not MCP_JSON_PATH.exists():
        return "source: mcp_meta\nMeta file not found."
    data = json.loads(MCP_JSON_PATH.read_text(encoding="utf-8"))
    # простая логика: ищем ключ, содержащий query (case‑insensitive)
    for key, value in data.items():
        if query.lower() in key.lower():
            return f"source: mcp_meta\n{key}: {value}"
    return "source: mcp_meta\nNo matching metadata found."

# ---------------------------------------------------------------------------
# 4. Создание агента
# ---------------------------------------------------------------------------
# Системный промпт, который заставляет агент выбирать нужный инструмент
SYSTEM_PROMPT = (
    "You are a helpful assistant for a course FAQ. Use the provided tools to answer the user. "
    "If the question is about course materials, use search_course_docs. "
    "If it is about schedule or metadata, use fetch_course_meta. "
    "Do not use both tools unless necessary. "
    "Always prefix your answer with the source: chroma or source: mcp_meta."
)

# Создаём LLM
llm = ChatOllama(model="llama3", temperature=0.2)

# Prompt template
prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("human", "{input}"),
])

# Создаём Runnable, который будет использовать инструменты
agent = create_agent(
    llm=llm,
    tools=[search_course_docs, fetch_course_meta],
    system_message=SYSTEM_PROMPT,
    verbose=True,
)

# ---------------------------------------------------------------------------
# 5. CLI
# ---------------------------------------------------------------------------
PREDEFINED_QUESTIONS = [
    "What topics are covered in the first lecture?",  # Chroma
    "Explain the concept of polymorphism in the context of the course.",  # Chroma
    "What is the schedule for the next week?",  # MCP
]

def run_cli():
    print("--- FAQ Bot ---")
    print("Predefined questions:")
    for i, q in enumerate(PREDEFINED_QUESTIONS, 1):
        print(f"{i}. {q}")
    print("\nEnter 1-3 to ask a predefined question, or type your own question.")
    while True:
        user_input = input("\n> ")
        if user_input.lower() in {"exit", "quit"}:
            print("Goodbye!")
            break
        if user_input.isdigit() and 1 <= int(user_input) <= len(PREDEFINED_QUESTIONS):
            question = PREDEFINED_QUESTIONS[int(user_input) - 1]
        else:
            question = user_input
        try:
            response = agent.invoke({"input": question})
            print("\nAnswer:\n", response)
        except Exception as e:
            print("Error:", e)

if __name__ == "__main__":
    # Убедимся, что данные загружены
    load_faq_to_chroma()
    run_cli()