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"

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)

            # Attempt to connect to the assigned Neo4j instance
            try:
                self.graph = Neo4jGraph(
                    url=self.profile.uri,
                    username=self.profile.username,
                    password=self.profile.password
                )
            except Exception as e:
                logger.error(f"Failed to initialize Neo4jGraph: {e}")
                self.graph = None

            # Initialize LLM if graph is available
            try:
                self.llm = ChatOpenAI(
                    api_key=self.profile.openai_api_key,
                    model_name=self.profile.model_name,
                    temperature=0
                )
            except Exception as e:
                logger.error(f"Failed to initialize ChatOpenAI: {e}")
                self.llm = None

            # Prepare graph transformer only if llm available
            if self.llm:
                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."
                )
                try:
                    self.graph_transformer = LLMGraphTransformer(
                        llm=self.llm,
                        prompt=english_prompt
                    )
                except Exception as e:
                    logger.error(f"Failed to initialize LLMGraphTransformer: {e}")
                    self.graph_transformer = None
            else:
                self.graph_transformer = None

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

        except Neo4jProfile.DoesNotExist:
            logger.error(f"Neo4j profile '{profile_name}' not found.")
            self.graph = None
            self.llm = None
            self.graph_transformer = None
        except Exception as e:
            logger.error(f"Unexpected error initializing Neo4jDatabase: {e}")
            self.graph = None
            self.llm = None
            self.graph_transformer = None

    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`.
        """
        if not self.graph_transformer:
            logger.warning("Graph transformer not available, skipping store_interaction.")
            return

        try:
            timestamp = datetime.utcnow().isoformat()
            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}),
            ]
            graph_docs = self.graph_transformer.convert_to_graph_documents(documents)
            if self.graph:
                self.graph.add_graph_documents(graph_docs, include_source=True)
                logger.info(f"Stored interaction in Neo4j (Profile: {self.profile.name})")
            else:
                logger.warning("No Neo4j graph available, skipping add_graph_documents.")
        except Exception as e:
            logger.error(f"Failed to store interaction in Neo4j: {e}")

    def query_graph(self, user_query):
        """
        Queries the graph using GraphCypherQAChain and returns a structured response.
        """
        if not self.graph or not self.llm:
            logger.warning("Graph or LLM not available, cannot query.")
            return None

        try:
            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 query result: {result}")
            return result.get('result')
        except Exception as e:
            logger.error(f"Graph query failed: {e}")
            return None
