from langchain_community.llms import Ollama
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain.memory import ConversationBufferWindowMemory
import os
from dotenv import load_dotenv
from tempfile import NamedTemporaryFile

# Load environment variables
load_dotenv()

# Global dictionary to store memory instances per session
session_memories = {}

# Initialize Ollama LLM and Embeddings
llm = Ollama(model="tinyllama", temperature=0.7)
embeddings = OllamaEmbeddings(model="tinyllama")

# Initialize global Chroma vector store (in-memory)
vector_store = Chroma.from_texts([""], embeddings)  # Initialize empty store

# Function to index uploaded file
def index_file(file_content: bytes, file_name: str):
    with NamedTemporaryFile(delete=False, suffix=os.path.splitext(file_name)[1]) as temp_file:
        temp_file.write(file_content)
        temp_file_path = temp_file.name

    loader = TextLoader(temp_file_path)
    documents = loader.load()

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = text_splitter.split_documents(documents)

    vector_store.add_documents(chunks)

    os.unlink(temp_file_path)

# Define prompt templates
def get_prompt_with_history(session_id):
    memory = session_memories.get(session_id)
    if not memory:
        memory = ConversationBufferWindowMemory(k=3)  # Keep last 3 turns
        session_memories[session_id] = memory
    return PromptTemplate(
        input_variables=["question", "history"],
        template=f"Previous conversation:\n{{history}}\n\nResponda à seguinte pergunta: {{question}}"
    )

def get_prompt_with_history_and_docs(session_id, docs):
    memory = session_memories.get(session_id)
    if not memory:
        memory = ConversationBufferWindowMemory(k=3)
        session_memories[session_id] = memory
    history_text = memory.buffer if memory.buffer else "No previous conversation."
    docs_text = "\n".join([f"Source: {doc.page_content}" for doc in docs]) if docs else "No relevant documents found."
    return PromptTemplate(
        input_variables=["question", "history"],
        template=f"Previous conversation:\n{history_text}\n\nRelevant documents:\n{docs_text}\n\nResponda à seguinte pergunta usando as fontes relevantes e citando trechos como fontes: {{question}}"
    )

def get_answer(session_id: str, question: str) -> str:
    memory = session_memories.get(session_id, ConversationBufferWindowMemory(k=3))
    prompt = get_prompt_with_history(session_id)
    chain = LLMChain(llm=llm, prompt=prompt, memory=memory)
    response = chain.run(question=question)
    response = response[:100] if len(response) > 100 else response
    return response

def ask_rag(session_id: str, question: str, file_content: bytes = None, file_name: str = None) -> dict:
    if file_content and file_name:
        index_file(file_content, file_name)
    
    memory = session_memories.get(session_id, ConversationBufferWindowMemory(k=3))
    docs = vector_store.similarity_search(question, k=3)
    prompt = get_prompt_with_history_and_docs(session_id, docs)
    chain = LLMChain(llm=llm, prompt=prompt, memory=memory)
    
    response = chain.run(question=question)
    response = response[:100] if len(response) > 100 else response
    
    sources = [doc.page_content for doc in docs]
    return {"answer": response, "sources": sources}

if __name__ == "__main__":
    session_id = "test_session"
    print(get_answer(session_id, "Qual a capital da França?"))
    print(get_answer(session_id, "E a capital da Espanha?"))