import os
from pathlib import Path
from typing import List

from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from langchain_tavily import TavilySearchResults
from langchain.agents import Tool, AgentExecutor, initialize_agent
from dotenv import load_dotenv

# Загрузка переменных окружения из .env
load_dotenv()
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
if not TAVILY_API_KEY:
    raise ValueError("Требуется переменная окружения TAVILY_API_KEY")

def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
    """
    Создаёт и возвращает объект Chroma, подключённый к указанной директории.
    Если директория не существует, она будет создана.
    """
    Path(persist_directory).mkdir(parents=True, exist_ok=True)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    return Chroma(
        persist_directory=persist_directory,
        embedding_function=embeddings,
    )

def _load_text_files(directory: str) -> List[Document]:
    """
    Читает все .txt и .md файлы из указанной директории и возвращает список Document.
    """
    docs: List[Document] = []
    for file_path in Path(directory).rglob("*"):
        if file_path.suffix.lower() in {".txt", ".md"}:
            text = file_path.read_text(encoding="utf-8")
            docs.append(Document(page_content=text, metadata={"source": str(file_path)}))
    return docs

def load_documents(directory: str, vectorstore: Chroma) -> None:
    """
    Загружает документы из указанной директории в ChromaDB.
    Делает чанкинг с помощью RecursiveCharacterTextSplitter и добавляет в коллекцию.
    """
    # Читаем файлы
    raw_docs = _load_text_files(directory)

    # Разбиваем на чанки
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        separators=["\n\n", "\n", " ", ""],
    )
    chunks = splitter.split_documents(raw_docs)

    # Добавляем в коллекцию
    vectorstore.add_documents(chunks)
    # Сохраняем изменения
    vectorstore.persist()
    print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}")

# ------------------------------------------------------------------
# Инструмент web_search, использующий TavilySearchResults
def web_search(query: str) -> str:
    """
    Выполняет поиск по интернету через Tavily и возвращает строку с результатами.
    """
    tavily = TavilySearchResults(api_key=TAVILY_API_KEY)
    results = tavily.run(query)
    return "\n".join([f"{r['title']}: {r['url']}" for r in results])

# ------------------------------------------------------------------
# Инструмент search_local_kb, использующий локальный Chroma
def search_local_kb(query: str, vectorstore: Chroma) -> List[dict]:
    """
    Выполняет поиск по локальной базе знаний (Chroma) и возвращает список документов.
    """
    docs = vectorstore.similarity_search_with_score(query, k=5)
    return [{"content": doc.page_content, "score": score} for doc, score in docs]

# ------------------------------------------------------------------
# Создание агента, выбирающего между локальным поиском и web search
def create_agent(vectorstore: Chroma) -> AgentExecutor:
    """
    Инициализирует агент с двумя инструментами:
      - search_local_kb: поиск в локальной базе знаний
      - web_search: поиск через Tavily
    Агент выводит источник ответа.
    """
    tools = [
        Tool(
            name="search_local_kb",
            func=lambda q: search_local_kb(q, vectorstore),
            description="Search the local knowledge base stored in Chroma."
        ),
        Tool(
            name="web_search",
            func=web_search,
            description="Search the web using Tavily. Useful for up-to-date information."
        )
    ]

    llm = OllamaLLM(model="llama3")
    agent_executor = initialize_agent(
        tools,
        llm,
        agent_type="zero-shot-react-description",
        verbose=True,
        return_intermediate_steps=False
    )
    return agent_executor

# ------------------------------------------------------------------
# Пример использования агента (необязательно, можно удалить)
if __name__ == "__main__":
    vs = create_vectorstore()
    # Предположим, что документы уже загружены ранее
    agent = create_agent(vs)
    query = "What is the capital of France?"
    response = agent.run(query)
    print("\nОтвет агента:")
    print(response)