import os
import logging
from typing import List, Dict, Any, Optional
from dotenv import load_dotenv

try:
    from neo4j import AsyncGraphDatabase, AsyncDriver
except ImportError:
    AsyncGraphDatabase = None
    AsyncDriver = None

load_dotenv()
logger = logging.getLogger(__name__)

class GraphRetriever:
    """
    Modul untuk mengambil konteks relasional dari Neo4j Knowledge Graph.
    Menggunakan driver asinkron dan praktik terbaik lifespan management.
    """

    def __init__(self, uri: Optional[str] = None, username: Optional[str] = None, password: Optional[str] = None):
        self.uri = uri or os.getenv("NEO4J_URI", "bolt://localhost:7687")
        self.username = username or os.getenv("NEO4J_USERNAME", "neo4j")
        self.password = password or os.getenv("NEO4J_PASSWORD", "")
        self._driver: Optional[AsyncDriver] = None

    async def connect(self):
        """Membangun koneksi ke Neo4j (Panggil di lifespan startup)."""
        if AsyncGraphDatabase is None:
            logger.error("neo4j library is not installed.")
            return

        try:
            self._driver = AsyncGraphDatabase.driver(
                self.uri,
                auth=(self.username, self.password),
                max_connection_pool_size=50,
                connection_timeout=10.0
            )
            # Verify connectivity to fail fast if config is wrong
            await self._driver.verify_connectivity()
            logger.info("Berhasil terhubung ke Neo4j Database.")
        except Exception as e:
            logger.error(f"Gagal terhubung ke Neo4j: {e}")
            self._driver = None

    async def close(self):
        """Menutup koneksi Neo4j (Panggil di lifespan shutdown)."""
        if self._driver:
            await self._driver.close()
            logger.info("Koneksi Neo4j ditutup.")

    async def get_node_title_by_hash(self, file_hash: str) -> Optional[str]:
        """Mengambil judul dokumen (label) dari Neo4j berdasarkan hash-nya."""
        if not self._driver:
            return None
        cypher_query = """
        MATCH (n:Literature {hash: $hash})
        RETURN n.title AS title
        LIMIT 1
        """
        try:
            async with self._driver.session() as session:
                result = await session.run(cypher_query, hash=file_hash)
                record = await result.single()
                if record and record["title"]:
                    return record["title"]
        except Exception as e:
            logger.error(f"Error get_node_title_by_hash untuk {file_hash}: {e}")
        return None

    async def get_relevant_neighbors(
        self,
        target_hash: str,
        project_id: str,
        user_id: str,
        max_nodes: int = 3,
    ) -> List[Dict[str, Any]]:
        """
        Mencari artikel tetangga yang terhubung lewat SIMILAR_* edge.
        Difilter dan diurutkan berdasarkan visual_pagerank_size tertinggi.
        
        Args:
            target_hash: file_hash dari node utama (Primary).
            project_id: project aktif untuk isolasi query graf.
            user_id: user aktif untuk isolasi query graf.
            max_nodes: jumlah maksimal tetangga yang diambil (Expansion layer).
            
        Returns:
            List of dict berisi 'hash' tetangga dan 'reasoning' hubungan.
        """
        if not self._driver:
            logger.warning("Driver Neo4j tidak diinisialisasi. Gagal mengambil tetangga.")
            return []

        # Cypher query untuk mendapatkan tetangga terdekat dengan bobot pagerank
        cypher_query = """
        MATCH (source:Literature {
            hash: $target_hash,
            project_id: $project_id,
            user_id: $user_id
        })-[r]->(neighbor:Literature {
            project_id: $project_id,
            user_id: $user_id
        })
        WHERE type(r) STARTS WITH 'SIMILAR_'
        WITH neighbor, r
        ORDER BY coalesce(neighbor.visual_pagerank_size, 0) DESC, r.score DESC
        LIMIT $limit
        RETURN neighbor.hash AS neighbor_hash,
               neighbor.title AS neighbor_title,
               r.reasoning AS reasoning,
               r.score AS similarity_score,
               type(r) AS relation_type
        """

        neighbors = []
        try:
            async with self._driver.session() as session:
                result = await session.run(
                    cypher_query,
                    target_hash=target_hash,
                    project_id=project_id,
                    user_id=user_id,
                    limit=max_nodes,
                )
                async for record in result:
                    neighbors.append({
                        "hash": record["neighbor_hash"],
                        "title": record["neighbor_title"],
                        "reasoning": record["reasoning"],
                        "similarity_score": record["similarity_score"],
                        "relation_type": record["relation_type"]
                    })
            return neighbors
        except Exception as e:
            logger.error(
                "Error saat query graf tetangga",
                extra={
                    "target_hash": target_hash,
                    "project_id": project_id,
                    "user_id": user_id,
                    "error": str(e),
                },
            )
            return []
