import logging
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_openai import ChatOpenAI
from django.conf import settings
from langchain_community.graphs import Neo4jGraph
from langchain.schema import Document
from .models import Neo4jProfile
from datetime import datetime
from langchain.prompts import ChatPromptTemplate
from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain


logger = logging.getLogger(__name__)

DEFAULT_PROFILE_NAME = "DefaultNeo4jProfile"  # Define a default profile

class Neo4jDatabase:
    """
    Handles connection to Neo4j using different profiles for various use cases.
    """

    def __init__(self, profile_name=None):
        try:
            # Load the specified Neo4j profile, or fallback to the default
            if profile_name:
                self.profile = Neo4jProfile.objects.get(name=profile_name)
            else:
                logger.warning("No profile specified. Using default Neo4j profile.")
                self.profile = Neo4jProfile.objects.get(name=DEFAULT_PROFILE_NAME)

            # Connect to the assigned Neo4j instance
            self.graph = Neo4jGraph(
                url=self.profile.uri,
                username=self.profile.username,
                password=self.profile.password
            )

            # Load the correct AI model
            self.llm = ChatOpenAI(
                api_key=self.profile.openai_api_key,
                model_name=self.profile.model_name,
                temperature=0
            )

            # Add the custom prompt to enforce English
            english_prompt = ChatPromptTemplate.from_template(
                "Transform the following text into a graph structure. All nodes, relationships, and properties should be in English, regardless of the original language."
            )

            # Apply the custom English prompt here
            self.graph_transformer = LLMGraphTransformer(
                llm=self.llm,
                prompt=english_prompt
            )

            logger.info(f"Neo4jDatabase initialized with profile: {self.profile.name}")

        except Neo4jProfile.DoesNotExist:
            logger.error(f"Neo4j profile '{profile_name}' not found.")
            raise
        except Exception as e:
            logger.error(f"Failed to initialize Neo4jDatabase: {str(e)}")
            raise

    def store_interaction(self, user_id, bot_id, user_message, bot_response, platform):
        """
        Stores a chatbot interaction as a structured graph in Neo4j.
        Converts messages into `Document` objects to work with `LLMGraphTransformer`.
        """
        try:

            timestamp = datetime.utcnow().isoformat()

            # Convert messages into `Document` objects
            documents = [
                Document(page_content=user_message, metadata={"role": "user", "user_id": user_id, "platform": platform, "created_at": timestamp}),
                Document(page_content=bot_response, metadata={"role": "bot", "bot_id": bot_id, "platform": platform, "created_at": timestamp}),
            ]

            # Convert text into structured graph documents
            graph_docs = self.graph_transformer.convert_to_graph_documents(documents)

            # Store the structured graph data into Neo4j
            self.graph.add_graph_documents(graph_docs, include_source=True)

            logger.info(f"Stored interaction in Neo4j (Profile: {self.profile.name})")
        except Exception as e:
            logger.error(f"Failed to store interaction in Neo4j: {str(e)}")
            raise

    def query_graph(self, user_query):
        """
        Queries the graph using GraphCypherQAChain and returns a structured response.
        """
        logger.info(f"Calling Neo4j {user_query}")
        try:
            # Use AI model to generate Cypher query
            qa_chain = GraphCypherQAChain.from_llm(
                llm=self.llm,
                graph=self.graph,
                verbose=True,
                allow_dangerous_requests=True
            )

            result = qa_chain.invoke({"query": user_query})
            logger.info(f"Resulting Neo4j {result}")
            return result['result']

        except Exception as e:
            logger.error(f"Graph query failed: {e}")
            return None