from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.chains import LLMChain

# Define the system prompt that teaches the LLM how to construct graph queries
SYSTEM_TEMPLATE = """You are a graph query expert that converts natural language questions into Cypher queries.
Follow these rules when constructing queries:
1. Always use lowercase for relationship types and property names
2. Use descriptive variable names (e.g., 'person' instead of 'p')
3. Include relevant properties in the RETURN clause
4. Add appropriate WHERE clauses for filtering
5. Use OPTIONAL MATCH when the relationship might not exist
6. Limit results when appropriate

The graph schema is:
- Nodes:
  * Person(id, name, age, title)
  * Company(id, name, industry)
  * Project(id, name, status)
- Relationships:
  * WORKS_AT(since)
  * MANAGES(role)
  * WORKS_ON(role, hours)

Convert the user's question into a Cypher query."""

# Define the human prompt template
HUMAN_TEMPLATE = """Question: {question}

Create a Cypher query that answers this question. 
Return only the Cypher query with no explanations."""

# Create the full prompt template
prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_TEMPLATE),
    ("human", HUMAN_TEMPLATE)
])

def construct_graph_query(question: str, llm: ChatOpenAI) -> str:
    """
    Constructs a Cypher query from a natural language question.
    
    Args:
        question: Natural language question
        llm: Language model instance
        
    Returns:
        str: Generated Cypher query
    """
    # Create the chain
    chain = LLMChain(llm=llm, prompt=prompt)
    
    # Generate the query
    result = chain.invoke({"question": question})
    
    return result["text"].strip()

# Example usage
if __name__ == "__main__":
    example_questions = [
        "Who are all the senior managers at tech companies?",
        "Which projects have no one working on them?",
        "Find people who work at multiple companies"
    ]
    
    llm = ChatOpenAI(temperature=0, model="gpt-4")
    
    for question in example_questions:
        query = construct_graph_query(question, llm)
        print(f"\nQuestion: {question}")
        print(f"Generated Query: {query}")
        
# Example outputs:
"""
Question: Who are all the senior managers at tech companies?
Generated Query: 
MATCH (person:Person)-[m:MANAGES]->(company:Company)
WHERE company.industry = 'tech' AND m.role CONTAINS 'senior'
RETURN person.name, company.name, m.role
ORDER BY company.name

Question: Which projects have no one working on them?
Generated Query:
MATCH (project:Project)
WHERE NOT EXISTS ((person:Person)-[:WORKS_ON]->(project))
RETURN project.name, project.status

Question: Find people who work at multiple companies
Generated Query:
MATCH (person:Person)-[w:WORKS_AT]->(company:Company)
WITH person, COUNT(company) as company_count
WHERE company_count > 1
MATCH (person)-[w:WORKS_AT]->(company:Company)
RETURN person.name, COLLECT(company.name) as companies, COUNT(company) as num_companies
ORDER BY num_companies DESC
"""