import os
import json
import logging
from typing import List, Dict, Any

from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_qdrant import QdrantStore
from langchain_core.tools import tool
from langchain.agents import AgentExecutor, ZeroShotAgent, Tool
from langchain.schema import HumanMessage, AIMessage

# ==========================
# Конфигурация и логирование
# ==========================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "knowledge_base")

# ==========================
# Векторный хранилище
# ==========================

try:
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
except Exception as e:
    logger.error(f"Не удалось инициализировать OllamaEmbeddings: {e}")
    raise

try:
    qdrant_store = QdrantStore(
        url=QDRANT_URL,
        collection_name=QDRANT_COLLECTION,
        embeddings=embeddings,
    )
except Exception as e:
    logger.error(f"Не удалось подключиться к Qdrant: {e}")
    raise


# ==========================
# Инструменты
# ==========================

@tool("search_knowledge_base")
def search_knowledge_base(query: str) -> List[Dict[str, Any]]:
    """
    Поиск релевантных документов в базе знаний.

    Args:
        query (str): поисковый запрос.

    Returns:
        list of dict: список найденных документов с полями 'id', 'text' и 'score'.
    """
    try:
        results = qdrant_store.similarity_search_with_score(query, k=5)
    except Exception as e:
        logger.exception(f"Ошибка при поиске в Qdrant: {e}")
        return []

    docs = []
    for doc, score in results:
        docs.append(
            {
                "id": getattr(doc.metadata, "doc_id", None),
                "text": doc.page_content,
                "score": score,
            }
        )
    logger.debug(f"search_knowledge_base: найдено {len(docs)} документов")
    return docs


@tool("add_to_knowledge_base")
def add_to_knowledge_base(text: str) -> str:
    """
    Добавление нового текста в базу знаний.

    Args:
        text (str): текст, который нужно добавить.

    Returns:
        str: сообщение об успехе или ошибке.
    """
    try:
        # Разбиваем на чанки с гибкой схемой
        chunk_size = 500  # символов
        overlap = 50      # символов перекрытия

        chunks = []
        start = 0
        while start < len(text):
            end = min(start + chunk_size, len(text))
            chunk_text = text[start:end]
            chunks.append(chunk_text)
            start += chunk_size - overlap

        for i, chunk in enumerate(chunks):
            doc_id = f"{hash(chunk)}_{i}"
            qdrant_store.add_documents(
                documents=[chunk],
                ids=[doc_id],
                metadata={"doc_id": doc_id},
            )
    except Exception as e:
        logger.exception(f"Ошибка при добавлении в Qdrant: {e}")
        return "Не удалось добавить документ."

    logger.info(f"add_to_knowledge_base: добавлено {len(chunks)} чанков")
    return f"Документ успешно добавлен, {len(chunks)} чанков."


# ==========================
# Агент
# ==========================

tools = [
    Tool.from_function(
        func=search_knowledge_base,
        name="search_knowledge_base",
        description="Используйте для поиска информации в базе знаний.",
    ),
    Tool.from_function(
        func=add_to_knowledge_base,
        name="add_to_knowledge_base",
        description="Добавьте новый текст в базу знаний.",
    ),
]

prompt = ZeroShotAgent.create_prompt(
    tools,
    prefix="""
Вы — интеллектуальный агент, помогающий пользователю находить и добавлять информацию в базу знаний.
Используйте инструменты search_knowledge_base и add_to_knowledge_base по мере необходимости.
""",
)

llm = ChatOllama(model="llama3")

agent_chain = ZeroShotAgent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor.from_agent_and_tools(
    agent=agent_chain,
    tools=tools,
    verbose=True,
)


# ==========================
# Тесты (можно запускать как `python agent.py`)
# ==========================

def _test_add_and_search():
    test_text = (
        "Python — это высокоуровневый язык программирования, "
        "который поддерживает несколько парадигм и широко используется в научных исследованиях."
    )
    add_result = add_to_knowledge_base(test_text)
    assert "успешно" in add_result.lower(), f"Добавление не удалось: {add_result}"

    search_query = "Python язык программирования"
    results = search_knowledge_base(search_query)
    assert isinstance(results, list), "Результат поиска должен быть списком."
    assert len(results) > 0, "Ни один документ не найден."

    logger.info("Тесты пройдены успешно.")


if __name__ == "__main__":
    # Пример использования агента
    user_query = input("Введите ваш запрос: ")
    try:
        response = agent_executor.invoke({"input": user_query})
        print("\nОтвет агента:")
        print(response["output"])
    except Exception as e:
        logger.exception(f"Ошибка при выполнении запроса: {e}")

    # Запуск тестов
    _test_add_and_search()