import gradio as gr
from transformers import pipeline
import json
from py2neo import Graph
import time
import requests
import json
import os

# Hugging Face model endpoint for text-generation (GPT-2) and NER (BERT)
llm_endpoint = "https://api-inference.huggingface.co/models/gpt2"
ner_endpoint = "https://api-inference.huggingface.co/models/dbmdz/bert-large-cased-finetuned-conll03-english"
# Define headers for authorization
hf_token = os.getenv('HF_TOKEN')
headers = {
    "Authorization": f"Bearer {hf_token}"
}
# Connect to Neo4j database
uri = "neo4j://143.198.173.3:7687"  # Replace with your database URI
username = "neo4j"       # Replace with your Neo4j username
password = "pickumaterinu"    # Replace with your Neo4j password
graph = Graph(uri, auth=(username, password))
print("DB connected")

def query_neo4j(question):
    """Extract entities from a question and retrieve their relationships from Neo4j."""
    try:
        # Generate entities using NER
        ner_response = query_with_retry(ner_endpoint, {"inputs": question})
        if 'error' in ner_response:
            return f"Error in NER response: {ner_response.get('error')}"
        
        # If NER response is empty or does not contain entities, handle gracefully
        if not ner_response:
            return "No entities extracted from the question."
        
        # Extract the entities from the NER response
        entities = []
        for item in ner_response:
            if 'word' in item:
                entities.append(item['word'])  # Add the word to entities list
        
        # Query Neo4j for relationships of extracted entities
        result = []
        try:
            for entity in entities:
                query_response = graph.run(
                    """
                    CALL {
                        WITH $entity AS entity
                        MATCH (p:Place {id: entity})-[r]->(e)
                        RETURN p.id AS source_id, type(r) AS relationship, e.id AS target_id
                        LIMIT 50
                    }
                    RETURN source_id, relationship, target_id
                    """,
                    entity=entity
                )
                result.extend([f"{row['source_id']} - {row['relationship']} -> {row['target_id']}" for row in query_response])
            
        except Exception as e:
            return f"Error with querying Neo4j: {str(e)}"

        return "\n".join(result) if result else None

    except Exception as e:
        return f"An error occurred while processing the question: {str(e)}"


def query_llm(question):
    """Fetch data from Neo4j and provide an answer using GPT-2."""
    graph_data = query_neo4j(question)
    if graph_data is not None:
        if 'error' in graph_data:
            return graph_data  

        context = f"Graph data:\n{graph_data}"
        llm_response = query_with_retry(llm_endpoint, {"inputs": context})
        
        if 'error' in llm_response:
            return f"Error in LLM response: {llm_response.get('error')}"

        if not llm_response or 'generated_text' not in llm_response[0]:
            return "LLM generated an empty response."

        return llm_response[0]["generated_text"]
    else:
        return "No relevant data found"

# query hugging face api with retry for model load
def query_with_retry(endpoint, payload, max_retries=5):
    """Send a request to the Hugging Face API with retry logic in case of model loading delays."""
    for _ in range(max_retries):
        response = requests.post(endpoint, headers=headers, json=payload)
        response_data = response.json()

        # Check if the model is ready or if an error occurred
        if 'error' in response_data:
            if 'currently loading' in response_data['error']:
                # Wait for a few seconds before retrying
                estimated_time = response_data['error'].split('estimated_time":')[1].split('}')[0]
                print(f"Model loading... waiting for {estimated_time} seconds.")
                time.sleep(float(estimated_time))
            else:
                print(f"Error: {response_data['error']}")
                break
        else:
            return response_data
    return {"error": "Max retries exceeded, model still loading."}

# Gradio Interface
iface = gr.Interface(
    fn=query_llm,
    inputs="text",
    outputs="text",
    live=False,
    title="RAG - Neo4j & LLM Integration",
    description="Fetches data from Neo4j and generates a response using a Hugging Face GPT-2 model and NER."
)
iface.launch(share=True)