import os
import json
from pathlib import Path

import httpx
from dotenv import load_dotenv
from langchain_ollama import ChatOllama
from langchain_chroma import Chroma
from langchain_core.prompts import PromptTemplate
from langchain.agents import initialize_agent, Tool, AgentType
from langchain.schema import Document

# --------------------------------------------------------------------------- #
# 1. Загрузка FAQ в ChromaDB
# --------------------------------------------------------------------------- #

DATA_DIR = Path("data")
CHROMA_PATH = Path("./chroma_faq")

def load_faq_to_chroma() -> Chroma:
    """
    Читает все *.md файлы из DATA_DIR, разбивает их на чанки и сохраняет
    в локальный ChromaDB. Возвращает объект Chroma.
    """
    # Создаём список документов
    docs = []
    for md_file in DATA_DIR.glob("*.md"):
        text = md_file.read_text(encoding="utf-8")
        docs.append(Document(page_content=text, metadata={"source": str(md_file)}))

    # Инициализируем Chroma с Ollama embeddings
    chroma = Chroma.from_documents(
        documents=docs,
        embedding=ChatOllama(model="nomic-embed-text"),
        persist_directory=str(CHROMA_PATH),
    )
    return chroma

def search_course_docs(query: str, k: int = 3) -> list[Document]:
    """
    Поиск в ChromaDB. Возвращает список документов с наиболее релевантными
    ответами на запрос.
    """
    chroma = load_faq_to_chroma()
    return chroma.similarity_search_with_score(query, k=k)

# --------------------------------------------------------------------------- #
# 2. MCP‑style tool (mock)
# --------------------------------------------------------------------------- #

def fetch_course_meta(query: str) -> dict:
    """
    Имитация вызова внешнего MCP сервера. Читает статический JSON
    из файла mock_meta.json и возвращает часть, соответствующую query.
    """
    meta_file = Path("mock_meta.json")
    if not meta_file.exists():
        raise FileNotFoundError(f"{meta_file} не найден")

    data = json.loads(meta_file.read_text(encoding="utf-8"))
    # простая фильтрация по ключу
    return {k: v for k, v in data.items() if query.lower() in k.lower()}

# --------------------------------------------------------------------------- #
# 3. Agent + Tools
# --------------------------------------------------------------------------- #

def create_agent():
    """
    Создаёт агента LangChain с двумя инструментами:
      - search_course_docs (Chroma)
      - fetch_course_meta (MCP‑style)
    Агент использует Ollama LLM.
    """
    llm = ChatOllama(model="nomic-embed-text")

    tools = [
        Tool(
            name="search_course_docs",
            func=lambda q: "\n".join([f"{d.page_content[:200]}..." for d in search_course_docs(q)]),
            description=(
                "Используется, если вопрос относится к материалам курса. "
                "Возвращает краткие отрывки из FAQ."
            ),
        ),
        Tool(
            name="fetch_course_meta",
            func=lambda q: json.dumps(fetch_course_meta(q), indent=2),
            description=(
                "Вызывается, если нужно получить расписание или метаданные курса. "
                "Возвращает JSON‑объект."
            ),
        ),
    ]

    system_prompt = (
        "Ты FAQ‑бот для локального курса. "
        "Если вопрос относится к материалам занятия — используй search_course_docs. "
        "Если нужен график, даты или другая метаинформация — fetch_course_meta. "
        "Не вызывай оба инструмента одновременно. В ответе указывай источник: "
        "`source: chroma` или `source: mcp_meta`."
    )

    prompt = PromptTemplate(
        input_variables=["input"],
        template=system_prompt + "\n\nВопрос: {input}\nОтвет:",
    )

    agent = initialize_agent(
        tools,
        llm,
        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
        verbose=True,
        max_iterations=5,
        early_stopping_method="generate",
        system_message=prompt.format(input="{input}"),
    )
    return agent

# --------------------------------------------------------------------------- #
# 4. CLI
# --------------------------------------------------------------------------- #

def main():
    """
    Простая консольная утилита с тремя готовыми вопросами и интерактивным режимом.
    """
    print("=== FAQ‑бот для курса ===")
    sample_questions = [
        "Какой материал будет на следующем занятии?",
        "Когда начинается следующий модуль?",
        "Какие темы охватывает первый раздел?",
    ]

    for idx, q in enumerate(sample_questions, 1):
        print(f"\n{idx}. {q}")

    print("\nВведите номер вопроса (1‑3) или '0' для интерактивного режима.")
    choice = input("Ваш выбор: ").strip()

    if choice == "0":
        while True:
            q = input("\nВопрос (или 'exit'): ")
            if q.lower() in ("exit", "quit"):
                break
            agent = create_agent()
            response = agent.run(q)
            print(f"\nОтвет:\n{response}\n")
    else:
        try:
            idx = int(choice) - 1
            question = sample_questions[idx]
        except (ValueError, IndexError):
            print("Неверный выбор.")
            return

        agent = create_agent()
        response = agent.run(question)
        print(f"\nОтвет:\n{response}\n")

if __name__ == "__main__":
    load_dotenv()  # если понадобится переменные окружения
    main()