# rag_tool.py
from langchain.agents import Tool
from langchain_groq import ChatGroq
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from neo4j import GraphDatabase
from typing import List, Dict

def neo4j_retriever(query: str, driver) -> List[Dict[str, str]]:
    with driver.session() as session:
        documents = []
        
        # Query for databases
        if "database" in query.lower():
            result = session.run("MATCH (d:Database) RETURN d.name AS database_name")
            databases = [record["database_name"] for record in result]
            documents.append({"content": f"Databases: {', '.join(databases)}"})
        
        # Query for tables with their respective databases
        result = session.run("""
        MATCH (d:Database)-[:HAS_TABLE]->(t:Table)
        RETURN d.name AS database_name, t.name AS table_name
        """)
        tables = [f"{record['database_name']}.{record['table_name']}" for record in result]
        documents.append({"content": f"Tables: {', '.join(tables)}"})
        
        # Query for columns
        result = session.run("""
        MATCH (d:Database)-[:HAS_TABLE]->(t:Table)-[:HAS_COLUMN]->(c:Column)
        RETURN d.name AS database_name, t.name AS table_name, 
               collect(c.name + ' (' + c.data_type + ')') AS column_info
        """)
        for record in result:
            content = f"Table: {record['database_name']}.{record['table_name']}, Columns: {', '.join(record['column_info'])}"
            documents.append({"content": content})
    
    return documents

def create_neo4j_db_rag_tool(api_key, neo4j_uri, neo4j_user, neo4j_password):
    llm = ChatGroq(groq_api_key=api_key, model_name="llama-3.3-70b-versatile", streaming=True)

    # Neo4j connection
    neo4j_driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))

    # Create a custom QA chain
    template = """Use the following pieces of context to answer the question at the end. 
    If you don't know the answer, just say that you don't know, don't try to make up an answer.
    When referring to tables, always use the format db_name.table_name.

    Context:
    {context}

    Question: {question}
    Answer:"""
    
    prompt = PromptTemplate(
        input_variables=["context", "question"],
        template=template,
    )

    chain = LLMChain(llm=llm, prompt=prompt)

    def qa_chain(input_str: str) -> str:
        # Split input into context and question if provided in format "context|||question"
        if "|||" in input_str:
            context_str, question = input_str.split("|||")
        else:
            # If no context provided, retrieve it using the entire input as the question
            question = input_str
            context = neo4j_retriever(question, neo4j_driver)
            context_str = "\n".join([doc["content"] for doc in context])
        
        return chain.run(context=context_str, question=question)

    return Tool(
        name="Neo4j RAG",
        func=qa_chain,
        description="Retrieves SQL database structure (Database, Table, Column) from Neo4j. Assists in SQL query preparation. Input: questions about database structure. Output: uses db_name.table_name format. Labels: Database, Table, Column."
    )