# Установка: pip install langgraph langchain langchain-community chromadb
from typing import TypedDict, List, Literal
from langgraph.graph import StateGraph, END
from langchain_community.vectorstores import Chroma
from langchain_community.llms import Ollama

# Определяем состояние графа
class AgenticRAGState(TypedDict):
    question: str              # Исходный вопрос
    documents: List[str]      # Извлечённые документы
    generation: str           # Сгенерированный ответ
    web_results: str          # Результаты web-поиска (fallback)
    hallucination_score: float # Оценка галлюцинаций
    retry_count: int           # Счётчик повторов

# Инициализация узлов StateGraph
workflow = StateGraph(AgenticRAGState)

# Узлы будут добавлены: route → retrieve → grade → transform → generate → check → [web | END]
print("LangGraph: построение графа Agentic RAG...")

# Логический роутер: vectorstore vs web search vs SQL database
def route_question(state: AgenticRAGState) -> Literal["vectorstore", "web_search", "sql_database"]:
    router_prompt = f"""Ты — маршрутизатор запросов. Определи источник данных:
Вопрос: {state['question']}

Варианты:
- vectorstore: вопрос о внутренних документах, отчётах, политиках компании
- web_search: вопрос о текущих событиях, новостях, внешней информации
- sql_database: вопрос, требующий агрегации (сумма, среднее, количество, группировка)

Ответь одним словом: vectorstore, web_search или sql_database."""

    decision = llm.invoke(router_prompt).strip().lower()
    if "web" in decision:
        return "web_search"
    elif "sql" in decision:
        return "sql_database"
    return "vectorstore"

# Добавляем узел роутера в граф
workflow.add_node("route", route_question)
workflow.set_entry_point("route")

# Условные рёбра: в зависимости от решения идём в разные ветки
workflow.add_conditional_edges(
    "route",
    lambda x: x,  # возвращает имя следующего узла
    {"vectorstore": "retrieve", "web_search": "web_search", "sql_database": "sql_query"}
)

# LLM-as-Judge: оцениваем релевантность каждого документа
def retrieve_and_grade(state: AgenticRAGState) -> AgenticRAGState:
    # Шаг 1: Извлечение документов
    raw_docs = vectorstore.similarity_search(state["question"], k=5)

    # Шаг 2: LLM оценивает каждый документ
    relevant_docs = []
    for doc in raw_docs:
        grade_prompt = f"""Оцени релевантность документа вопросу (yes/no).
Вопрос: {state['question']}
Документ: {doc.page_content[:500]}
Отвечает ли документ на вопрос? Ответь ТОЛЬКО yes или no:"""
        grade = llm.invoke(grade_prompt).strip().lower()
        if "yes" in grade:
            relevant_docs.append(doc.page_content)

    # Шаг 3: Определяем следующий шаг
    if len(relevant_docs) == 0:
        state["next_step"] = "transform_query"  # Нет релевантных → переформулируем
    elif len(relevant_docs) >= 2:
        state["next_step"] = "generate"       # Достаточно → генерируем
    else:
        state["next_step"] = "web_search"      # Мало → ищем в web

    state["documents"] = [d.page_content for d in relevant_docs]
    return state

workflow.add_node("retrieve_grade", retrieve_and_grade)

# Переформулировка запроса если документы нерелевантны
def transform_query(state: AgenticRAGState) -> AgenticRAGState:
    # Три стратегии переформулировки
    transform_prompt = f"""Исходный запрос не дал релевантных результатов.
Переформулируй запрос ТРЕМЯ разными способами (по одному на строку):
1. Синонимы и перефраз: замени ключевые слова синонимами
2. Более общий запрос: убери специфичные детали
3. Более конкретный запрос: добавь уточняющие детали

Исходный запрос: {state['question']}

Новые запросы:"""

    new_queries = llm.invoke(transform_prompt).strip().split("\n")

    # Пробуем каждый вариант
    for new_q in new_queries:
        docs = vectorstore.similarity_search(new_q, k=3)
        if docs:
            state["documents"] = [d.page_content for d in docs]
            state["retry_count"] = state.get("retry_count", 0) + 1
            break
    return state

workflow.add_node("transform_query", transform_query)
# После трансформации возвращаемся на retrieve_grade
workflow.add_edge("transform_query", "retrieve_grade")

# Проверка: соответствует ли ответ извлечённым документам?
def check_hallucination(state: AgenticRAGState) -> AgenticRAGState:
    # LLM проверяет каждый факт из ответа
    hallucination_prompt = f"""Проверь, все ли факты в ответе подтверждены документами.
Оцени по шкале 0.0-1.0, где 1.0 = все факты подтверждены.

Документы: {state['documents']}
Сгенерированный ответ: {state['generation']}

Дай ТОЛЬКО число от 0.0 до 1.0:"""

    score_str = llm.invoke(hallucination_prompt).strip()
    try:
        state["hallucination_score"] = float(score_str)
    except:
        state["hallucination_score"] = 0.5  # default если парсинг не удался

    # Если галлюцинации — пробуем web search
    if state["hallucination_score"] < 0.7:
        state["next_step"] = "web_search"
    else:
        state["next_step"] = "end"
    return state

workflow.add_node("hallucination_check", check_hallucination)

# Сборка и компиляция полного графа
def generate_answer(state: AgenticRAGState) -> AgenticRAGState:
    prompt = f"Ответь на вопрос, используя ТОЛЬКО документы ниже.\nДокументы: {state['documents']}\nВопрос: {state['question']}"
    state["generation"] = llm.invoke(prompt)
    return state

def web_search_fallback(state: AgenticRAGState) -> AgenticRAGState:
    # В реальном коде — вызов Tavily/Brave/SerpAPI
    state["web_results"] = f"Web search results for: {state['question']}"
    state["documents"].append(state["web_results"])
    return state

# Добавляем все узлы
workflow.add_node("generate", generate_answer)
workflow.add_node("web_search", web_search_fallback)

# Определяем маршруты графа
workflow.add_edge("retrieve_grade", "generate")
workflow.add_edge("generate", "hallucination_check")
workflow.add_edge("web_search", "generate")
workflow.add_edge("hallucination_check", END)

# Компилируем и запускаем
app = workflow.compile()

# Запуск с реальным вопросом
result = app.invoke({
    "question": "Какая стратегия выхода на рынок Азии описана в документах?",
    "documents": [],
    "retry_count": 0,
})
print(f"Ответ: {result['generation']}")
print(f"Hallucination score: {result['hallucination_score']}")
print(f"Retries: {result['retry_count']}")