import os
import json
import httpx
from pathlib import Path
from langchain_ollama import ChatOllama
from langchain_chroma import Chroma
from langchain.embeddings import OllamaEmbeddings
from langchain.schema import Document
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain

# 1. Load FAQ into Chroma

def load_faq_to_chroma(md_path: str, persist_dir: str = "./chroma_faq"):
    from langchain.document_loaders import TextLoader
    from langchain.text_splitter import RecursiveCharacterTextSplitter

    loader = TextLoader(md_path)
    docs = loader.load()
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    texts = splitter.split_documents(docs)
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    chroma = Chroma.from_documents(texts, embeddings, persist_directory=persist_dir)
    chroma.persist()
    return chroma

# 2. Search function

def search_course_docs(query: str, k: int = 3):
    chroma = Chroma(persist_directory="./chroma_faq", embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
    results = chroma.similarity_search(query, k=k)
    return [doc.page_content for doc in results]

# 3. MCP-style tool

def fetch_course_meta(query: str):
    # For demo, use static JSON file
    meta_path = Path("meta.json")
    if not meta_path.exists():
        return {"error": "Meta not found"}
    meta = json.loads(meta_path.read_text())
    # simple search by key
    return meta.get(query, {})

# 4. Agent logic

def answer_question(question: str):
    # Simple heuristic: if question contains "schedule" or "метаданные" use meta
    if any(word in question.lower() for word in ["schedule", "расписание", "метаданные"]):
        source = "mcp_meta"
        answer = fetch_course_meta(question)
    else:
        source = "chroma"
        answer = search_course_docs(question, k=1)[0]
    return {"answer": answer, "source": source}

if __name__ == "__main__":
    # CLI with preset questions
    preset = [
        "Что такое ChromaDB?",
        "Как подключить Ollama embeddings?",
        "Когда будет расписание следующего занятия?"
    ]
    for q in preset:
        res = answer_question(q)
        print(f"Q: {q}\nA: {res['answer']}\nSource: {res['source']}\n")
    # interactive
    while True:
        q = input("Ask a question (or 'exit'): ")
        if q.lower() == "exit":
            break
        res = answer_question(q)
        print(f"A: {res['answer']} (source: {res['source']})")