database_utils.py

import json
import os
import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
import together
from together import Together
import sqlite3
from collections import Counter, defaultdict
from dotenv import load_dotenv
from neo4j import GraphDatabase, NotificationDisabledCategory

load_dotenv()

# Suppress verbose HTTP logging from qdrant/httpx
import logging
for noisy in ["qdrant_client", "httpx", "urllib3", "httpcore"]:
    logging.getLogger(noisy).setLevel(logging.WARNING)

EMBEDDING_MAX_TOKENS = int(os.getenv('TOGETHER_EMBED_MAX_TOKENS', '480'))


def _truncate_for_embedding(text: str) -> str:
    from tools.utils import truncate_text
    normalized = truncate_text(text or '', max_tokens=EMBEDDING_MAX_TOKENS)
    return normalized[:EMBEDDING_MAX_TOKENS * 3]


def get_embedding(text: str) -> list:
    """Generate embedding for text/text_list using Together API"""
    if isinstance(text, str):
        text_list = [text]
    else:
        text_list = text
    try:
        client = Together(api_key=os.getenv("TOGETHER_API_KEY"))
        model = os.getenv("TOGETHER_MODEL", "intfloat/multilingual-e5-large-instruct")
        response = client.embeddings.create(
            input=[_truncate_for_embedding(text) for text in text_list],
            model=model
        )
        return [item.embedding for item in response.data]
    except Exception as e:
        print(f"Error generating embedding: {e}")
        return []


def get_graph_driver():
    """Get graph database driver (Neo4j/Memgraph compatible)"""
    uri = os.getenv("GRAPH_URI", "bolt://localhost:7687")
    user = os.getenv("GRAPH_USER", "neo4j")
    password = os.getenv("GRAPH_PASSWORD", "test123456")
    # Suppress unknown property key warnings
    return GraphDatabase.driver(
        uri, 
        auth=(user, password),
        notifications_disabled_categories=[NotificationDisabledCategory.UNRECOGNIZED]
    )


def _iter_levelled_entities(path: str):
    """Yield `(level, entity_dict)` pairs from newline-delimited or list-based JSON."""
    if not os.path.exists(path):
        return
    with open(path, 'r', encoding='utf-8') as stream:
        for level, line in enumerate(stream):
            line = line.strip()
            if not line:
                continue
            payload = json.loads(line)
            if isinstance(payload, list):
                for entity in payload:
                    yield level, entity
            elif isinstance(payload, dict):
                yield level, payload


def _iter_relations(path: str):
    """Yield relation dictionaries from JSONL or JSON-array files."""
    if not os.path.exists(path):
        return
    with open(path, 'r', encoding='utf-8') as stream:
        first_char = stream.read(1)
        while first_char and first_char.isspace():
            first_char = stream.read(1)
        stream.seek(0)
        if not first_char:
            return
        if first_char == '[':
            data = json.load(stream)
            items = data if isinstance(data, list) else [data]
            for relation in items:
                if isinstance(relation, list):
                    for item in relation:
                        yield item
                else:
                    yield relation
        else:
            for line in stream:
                line = line.strip()
                if not line:
                    continue
                payload = json.loads(line)
                if isinstance(payload, list):
                    for relation in payload:
                        yield relation
                else:
                    yield payload


def _iter_communities(path: str):
    """Yield community dictionaries from JSONL or JSON-array files."""
    if not os.path.exists(path):
        return
    with open(path, 'r', encoding='utf-8') as stream:
        for line in stream:
            line = line.strip()
            if not line:
                continue
            payload = json.loads(line)
            if isinstance(payload, list):
                for community in payload:
                    yield community
            else:
                yield payload


def build_vector_search(data, working_dir, recreate: bool = True):
    """Used for building up the vector search space consistently"""
    qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
    qdrant_api_key = os.getenv("QDRANT_API_KEY")
    qdrant_collection = os.getenv("QDRANT_COLLECTION", "entity_collection")

    client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
    
    collection_name = qdrant_collection
    exists = client.collection_exists(collection_name)
    if recreate and exists:
        client.delete_collection(collection_name)
        exists = False

    if not exists:
        client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
        )
    
    prepared = []
    total_levels = len(data)
    print(f"🔧 Preparing vectors for {total_levels} level groups...")
    next_id = 0
    for level, sublist in enumerate(data):
        level_count = 0
        nodes = sublist if isinstance(sublist, list) else [sublist]
        for node in nodes:
            vector = node.get('vector')
            if vector is None:
                continue
            vector = np.asarray(vector)
            if vector.ndim == 2 and vector.shape[0] == 1:
                vector = vector.squeeze(0)
            if vector.ndim != 1:
                logging.warning(
                    "Skipping vector for entity %s due to shape %s",
                    node.get('entity_name'),
                    getattr(vector, 'shape', None),
                )
                continue

            payload = {
                key: value
                for key, value in node.items()
                if key not in ('vector',)
            }
            payload['level'] = level
            prepared.append({
                'id': next_id,
                'vector': vector.astype(float).tolist(),
                'payload': payload,
            })
            next_id += 1
            level_count += 1
        print(f"  • Level {level+1}/{total_levels}: added {level_count} items (total {len(prepared)})", end='\r', flush=True)
    print()
    print(f"✅ Collected {len(prepared)} vectors. Uploading to Qdrant in batches...")

    piece = 32
    total = len(prepared)
    for batch_idx in range((total + piece - 1)//piece):
        start = batch_idx * piece
        end = min(start + piece, total)
        batch = prepared[start:end]
        client.upsert(collection_name=collection_name, points=batch)
        print(f"  • Uploaded {end}/{total} vectors", end='\r', flush=True)
    print(f"\n🚀 Vector collection '{collection_name}' ready (total {total})")

def search_vector_search(working_dir, query, topk=10, level_mode=2):
    """Search the Qdrant collection with optional filtering by hierarchy level.

    level_mode values:
        0 -> raw nodes only
        1 -> aggregate nodes only
        2 -> all nodes
    """
    from qdrant_client.models import Filter, FieldCondition, MatchValue
    
    qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
    qdrant_api_key = os.getenv("QDRANT_API_KEY")
    qdrant_collection = os.getenv("QDRANT_COLLECTION", "entity_collection")
    
    client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
    
    collection_name = qdrant_collection
    
    filter_conditions = None
    if level_mode == 0:
        filter_conditions = Filter(
            must=[FieldCondition(key="level", match=MatchValue(value=0))]
        )
    elif level_mode == 1:
        filter_conditions = Filter(
            must=[FieldCondition(key="level", range={"gt": 0})]
        )
    # else: all nodes, no filter
    
    # Extract the query vector (handle both single vector and batch)
    if hasattr(query, 'shape') and len(query.shape) > 1:
        query_vector = query[0].tolist()  # Take first vector if batched
    else:
        query_vector = query.tolist() if hasattr(query, 'tolist') else list(query)
    
    search_results = client.search(
        collection_name=collection_name,
        query_vector=query_vector,
        limit=topk,
        query_filter=filter_conditions,
        with_payload=["entity_name", "description", "parent", "level", "source_id"],
    )
    
    extract_results = [
        (
            hit.payload["entity_name"],
            hit.payload.get("parent"),
            hit.payload["description"],
            hit.payload.get("source_id"),
            hit.payload.get("level"),
        )
        for hit in search_results
    ]
    return extract_results

def get_db_path():
    """Return the canonical path to the LeanRAG relational store (dqlite)."""
    base_dir = "dqlite"
    os.makedirs(base_dir, exist_ok=True)
    return os.path.join(base_dir, "leanrag.db")


def create_db_table_sqlite(working_dir):
    # Always use unified dqlite path; working_dir ignored (kept for backward compatibility)
    # migrate_sqlite_to_dqlite()
    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    
    # Create tables
    cur.execute("""
        CREATE TABLE IF NOT EXISTS entities (
            entity_name TEXT,
            description TEXT,
            source_id TEXT,
            degree INTEGER DEFAULT 0,
            parent TEXT,
            level INTEGER,
            vector TEXT,
            is_new INTEGER DEFAULT 1,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS relations (
            src_tgt TEXT,
            tgt_src TEXT,
            description TEXT,
            weight INTEGER DEFAULT 1,
            level INTEGER,
            src_level INTEGER DEFAULT 0,
            tgt_level INTEGER DEFAULT 0,
            is_new INTEGER DEFAULT 1,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS communities (
            entity_name TEXT,
            entity_description TEXT,
            findings TEXT
        )
    """)
    
    cur.execute("""
        CREATE TABLE IF NOT EXISTS text_units (
            hash_code TEXT PRIMARY KEY,
            text TEXT,
            tokens INTEGER,
            quality_score REAL,
            strategy TEXT
        )
    """)
    
    # Backfill newly introduced columns for existing installations
    rel_columns = {row[1] for row in cur.execute("PRAGMA table_info(relations)").fetchall()}
    if 'src_level' not in rel_columns:
        cur.execute("ALTER TABLE relations ADD COLUMN src_level INTEGER DEFAULT 0")
    if 'tgt_level' not in rel_columns:
        cur.execute("ALTER TABLE relations ADD COLUMN tgt_level INTEGER DEFAULT 0")

    # Create indexes
    cur.execute("CREATE INDEX IF NOT EXISTS idx_entities_name ON entities(entity_name)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_entities_name_level ON entities(entity_name, level)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_relations_link ON relations(src_tgt, tgt_src)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_relations_nodes_levels ON relations(src_tgt, src_level, tgt_src, tgt_level)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_communities_name ON communities(entity_name)")
    cur.execute("CREATE INDEX IF NOT EXISTS idx_text_units_hash ON text_units(hash_code)")

    con.commit()
    con.close()

def store_text_units_in_db(working_dir, chunks_file):
    """Store text units from JSON file into SQLite database"""
    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    
    # Load chunks from file
    with open(chunks_file, 'r') as f:
        chunks_data = json.load(f)
    
    # Insert text units
    for chunk in chunks_data:
        cur.execute("""
            INSERT OR REPLACE INTO text_units (hash_code, text, tokens, quality_score, strategy)
            VALUES (?, ?, ?, ?, ?)
        """, (
            chunk['hash_code'],
            chunk['text'],
            chunk.get('tokens', 0),
            chunk.get('quality_score', 0.0),
            chunk.get('strategy', 'unknown')
        ))
    
    con.commit()
    con.close()
    print(f"Stored {len(chunks_data)} text units in database")

def store_text_units_records(records):
    """Store a list of chunk dicts (already loaded) into the text_units table.

    Each record must contain at minimum: hash_code, text. Optional: tokens, quality_score, strategy.
    """
    if not records:
        return 0
    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    inserted = 0
    for chunk in records:
        try:
            cur.execute(
                """
                INSERT OR REPLACE INTO text_units (hash_code, text, tokens, quality_score, strategy)
                VALUES (?, ?, ?, ?, ?)
                """,
                (
                    chunk.get('hash_code'),
                    chunk.get('text',''),
                    chunk.get('tokens', 0),
                    chunk.get('quality_score', 0.0),
                    chunk.get('strategy', 'unknown')
                )
            )
            inserted += 1
        except Exception:
            continue
    con.commit()
    con.close()
    return inserted


def insert_relationships_to_sqlite(working_dir, relationships):
    """Insert relationships into SQLite database"""
    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    
    for rel in relationships:
        src_tgt = f"{rel.get('source', '')} -> {rel.get('target', '')}"
        tgt_src = f"{rel.get('target', '')} -> {rel.get('source', '')}"
        
        cur.execute("""
            INSERT INTO relations (src_tgt, tgt_src, description, is_new)
            VALUES (?, ?, ?, 1)
        """, (
            src_tgt,
            tgt_src,
            rel.get('description', ''),
        ))
    
    con.commit()
    con.close()


def mark_entities_relations_processed(working_dir):
    """Mark all entities and relations as processed (is_new = 0)"""
    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    
    cur.execute("UPDATE entities SET is_new = 0 WHERE is_new = 1")
    cur.execute("UPDATE relations SET is_new = 0 WHERE is_new = 1")
    
    con.commit()
    con.close()

def persist_hierarchy_to_sqlite(
    entity_layers: list,
    base_relations: dict,
    aggregate_relations: dict,
    community_report: dict,
    replace_existing: bool = True,
) -> dict[str, int]:
    """Store hierarchical entities, relations, and communities in the shared SQLite store."""

    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()

    if replace_existing:
        cur.execute("DELETE FROM entities")
        cur.execute("DELETE FROM relations")
        cur.execute("DELETE FROM communities")

    entity_rows: list[tuple] = []
    entity_levels: defaultdict[str, set[int]] = defaultdict(set)
    for level, layer in enumerate(entity_layers):
        if isinstance(layer, dict):
            layer = [layer]
        if not isinstance(layer, list):
            continue
        for entity in layer:
            if not isinstance(entity, dict):
                continue
            name = entity.get('entity_name')
            if not name:
                continue
            entity_levels[name].add(level)
            vector = entity.get('vector')
            if isinstance(vector, np.ndarray):
                vector = vector.astype(float).tolist()
            elif vector is not None and not isinstance(vector, list):
                try:
                    vector = list(vector)
                except TypeError:
                    vector = None
            entity_rows.append((
                name,
                entity.get('description', ''),
                entity.get('source_id', ''),
                entity.get('degree', 0),
                entity.get('parent'),
                level,
                json.dumps(vector) if vector is not None else None,
            ))

    relation_rows: list[tuple] = []
    seen_relations: set[tuple[str, int, str, int, str]] = set()

    def resolve_level(name: str, preferred: int | None) -> int:
        levels = sorted(entity_levels.get(name, []))
        if not levels:
            logging.warning(
                "No entity levels recorded for %s; defaulting to level %s",
                name,
                preferred if preferred is not None else 0,
            )
            return preferred if preferred is not None else 0
        if preferred is not None and preferred in levels:
            return preferred
        if preferred is not None:
            for lvl in levels:
                if lvl >= preferred:
                    return lvl
            return levels[-1]
        return levels[0]

    def coerce_int(value, default=None):
        if value is None:
            return default
        try:
            return int(value)
        except (TypeError, ValueError):
            logging.warning("Unable to coerce value '%s' to int; using default %s", value, default)
            return default

    for rel in (base_relations or {}).values():
        src = rel.get('src_tgt')
        tgt = rel.get('tgt_src')
        if not src or not tgt:
            continue
        desc = rel.get('description', '')
        relation_level = coerce_int(rel.get('level'), default=0)
        preferred_level = relation_level if relation_level is not None else 0
        src_level = resolve_level(src, preferred_level)
        tgt_level = resolve_level(tgt, preferred_level)
        key = (src, src_level, tgt, tgt_level, desc)
        if key in seen_relations:
            continue
        seen_relations.add(key)
        relation_rows.append((
            src,
            tgt,
            desc,
            coerce_int(rel.get('weight', 1), default=1),
            preferred_level,
            src_level,
            tgt_level,
        ))

    for rel in (aggregate_relations or {}).values():
        if not isinstance(rel, dict):
            continue
        src = rel.get('src_tgt')
        tgt = rel.get('tgt_src')
        if not src or not tgt:
            continue
        desc = rel.get('description', '')
        relation_level = coerce_int(rel.get('level'))
        preferred_level = relation_level if relation_level is not None else None
        src_level = resolve_level(src, preferred_level)
        tgt_level = resolve_level(tgt, preferred_level)
        key = (src, src_level, tgt, tgt_level, desc)
        if key in seen_relations:
            continue
        seen_relations.add(key)
        relation_rows.append((
            src,
            tgt,
            desc,
            coerce_int(rel.get('weight', 1), default=1),
            preferred_level if preferred_level is not None else 0,
            src_level,
            tgt_level,
        ))

    community_rows: list[tuple] = []
    for community in (community_report or {}).values():
        if not isinstance(community, dict):
            continue
        name = community.get('entity_name')
        if not name:
            continue
        findings = community.get('findings', [])
        community_rows.append((
            name,
            community.get('entity_description', ''),
            json.dumps(findings, ensure_ascii=False),
        ))

    if entity_rows:
        cur.executemany(
            """
            INSERT OR REPLACE INTO entities(entity_name, description, source_id, degree, parent, level, vector, is_new)
            VALUES (?, ?, ?, ?, ?, ?, ?, 1)
            """,
            entity_rows,
        )

    if relation_rows:
        cur.executemany(
            """
            INSERT INTO relations(src_tgt, tgt_src, description, weight, level, src_level, tgt_level, is_new)
            VALUES (?, ?, ?, ?, ?, ?, ?, 1)
            """,
            relation_rows,
        )

    if community_rows:
        cur.executemany(
            """
            INSERT OR REPLACE INTO communities(entity_name, entity_description, findings)
            VALUES (?, ?, ?)
            """,
            community_rows,
        )

    con.commit()
    con.close()

    return {
        'entities': len(entity_rows),
        'relations': len(relation_rows),
        'communities': len(community_rows),
    }


def insert_data_to_memgraph(working_dir, clear_existing: bool = True) -> dict[str, int]:
    """Insert graph data from SQLite into the connected graph database."""
    driver = get_graph_driver()
    graph_db = os.getenv("GRAPH_DATABASE")

    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    entity_rows = cur.execute(
        """
        SELECT entity_name, description, source_id, degree, parent, level
        FROM entities
        WHERE entity_name IS NOT NULL
        """
    ).fetchall()
    relation_rows = cur.execute(
        """
        SELECT src_tgt, tgt_src, description, weight, level, src_level, tgt_level
        FROM relations
        WHERE src_tgt IS NOT NULL AND tgt_src IS NOT NULL
        """
    ).fetchall()
    community_rows = cur.execute(
        """
        SELECT entity_name, entity_description, findings
        FROM communities
        WHERE entity_name IS NOT NULL
        """
    ).fetchall()
    con.close()

    stats = {
        'cleared_nodes': 0,
        'entities': 0,
        'relationships': 0,
        'communities': 0,
    }

    session_kwargs = {}
    if graph_db:
        session_kwargs['database'] = graph_db

    try:
        with driver.session(**session_kwargs) as session:
            if clear_existing:
                result = session.run("MATCH (n) RETURN count(n) AS count")
                record = result.single()
                if record:
                    stats['cleared_nodes'] = record.get("count", 0)
                session.run("MATCH (n) DETACH DELETE n").consume()

            for name, description, source_id, degree, parent, level in entity_rows:
                summary = session.run(
                    """
                    MERGE (e:Entity {name: $name, level: $level})
                    SET e.description = $description,
                        e.source_id = $source_id,
                        e.degree = $degree,
                        e.parent = $parent,
                        e.level = $level,
                        e.updated_at = datetime()
                    """,
                    {
                        "name": name,
                        "description": description or '',
                        "source_id": "|".join((source_id or '').split("|")[:5]),
                        "degree": degree or 0,
                        "parent": parent,
                        "level": level or 0,
                    },
                ).consume()
                stats['entities'] += summary.counters.nodes_created or 0

            for src_name, tgt_name, description, weight, level, src_level, tgt_level in relation_rows:
                rel_src_level = src_level if src_level is not None else (level if level is not None else 0)
                rel_tgt_level = tgt_level if tgt_level is not None else (level if level is not None else 0)
                summary = session.run(
                    """
                    MATCH (src:Entity {name: $src_name, level: $src_level})
                    MATCH (tgt:Entity {name: $tgt_name, level: $tgt_level})
                    MERGE (src)-[rel:RELATES_TO]->(tgt)
                    SET rel.description = $description,
                        rel.weight = $weight,
                        rel.level = $level,
                        rel.updated_at = datetime()
                    """,
                    {
                        "src_name": src_name,
                        "src_level": rel_src_level,
                        "tgt_name": tgt_name,
                        "tgt_level": rel_tgt_level,
                        "description": description or '',
                        "weight": weight or 1,
                        "level": level or 0,
                    },
                ).consume()
                stats['relationships'] += summary.counters.relationships_created or 0

            for name, description, findings in community_rows:
                summary = session.run(
                    """
                    MERGE (c:Community {entity_name: $entity_name})
                    SET c.entity_description = $entity_description,
                        c.findings = $findings,
                        c.updated_at = datetime()
                    """,
                    {
                        "entity_name": name,
                        "entity_description": description or '',
                        "findings": findings or '[]',
                    },
                ).consume()
                stats['communities'] += summary.counters.nodes_created or 0
    finally:
        driver.close()

    return stats


def upsert_communities(working_dir: str, community_report: dict[str, dict]) -> int:
    """Insert or update community summaries in the SQLite store."""

    if not community_report:
        return 0

    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()

    rows = []
    for community in community_report.values():
        name = community.get('entity_name')
        if not name:
            continue
        description = community.get('entity_description', '')
        findings = community.get('findings', [])
        rows.append((
            name,
            description,
            json.dumps(findings, ensure_ascii=False),
        ))

    if rows:
        cur.executemany(
            """
            INSERT OR REPLACE INTO communities(entity_name, entity_description, findings)
            VALUES (?, ?, ?)
            """,
            rows,
        )
        con.commit()

    con.close()
    return len(rows)

def insert_data_to_sqlite(working_dir):
    db_path = os.path.join(working_dir, "leanrag.db")
    db = sqlite3.connect(db_path)
    cursor = db.cursor()
    
    entity_path = os.path.join(working_dir, "all_entities.json")
    entity_rows = []
    for level, entity in _iter_levelled_entities(entity_path):
        if not isinstance(entity, dict):
            continue
        entity_rows.append((
            entity.get('entity_name'),
            entity.get('description', ''),
            "|".join((entity.get('source_id') or '').split("|")[:5]),
            entity.get('degree', 0),
            entity.get('parent'),
            level,
        ))

    if entity_rows:
        cursor.executemany(
            "INSERT INTO entities(entity_name, description, source_id, degree, parent, level) VALUES (?, ?, ?, ?, ?, ?)",
            entity_rows,
        )
        db.commit()

    relation_path = os.path.join(working_dir, "generate_relations.json")
    relation_rows = []
    for relation in _iter_relations(relation_path):
        if not isinstance(relation, dict):
            continue
        relation_rows.append((
            relation.get('src_tgt'),
            relation.get('tgt_src'),
            relation.get('description', ''),
            relation.get('weight', 1),
            relation.get('level', 0),
        ))

    if relation_rows:
        cursor.executemany(
            "INSERT INTO relations(src_tgt, tgt_src, description, weight, level) VALUES (?, ?, ?, ?, ?)",
            relation_rows,
        )
        db.commit()

    community_path = os.path.join(working_dir, "community.json")
    community_rows = []
    for community in _iter_communities(community_path):
        if not isinstance(community, dict):
            continue
        community_rows.append((
            community.get('entity_name'),
            community.get('entity_description', ''),
            json.dumps(community.get('findings', []), ensure_ascii=False),
        ))

    if community_rows:
        cursor.executemany(
            "INSERT INTO communities(entity_name, entity_description, findings) VALUES (?, ?, ?)",
            community_rows,
        )
        db.commit()
    
    db.close()
def find_tree_root(working_dir, entity):
    """Find the root of the entity hierarchy in Memgraph"""
    driver = get_graph_driver()
    res = [entity]
    
    with driver.session() as session:
        # Since level property doesn't exist, use default depth
        depth = 5
        
        current_entity = entity
        for i in range(depth):
            # Since parent property doesn't exist, this will return no results
            query = "MATCH (e:Entity {name: $name}) WHERE e.parent IS NOT NULL RETURN e.parent AS parent"
            result = session.run(query, name=current_entity)
            record = result.single()
            
            if not record or not record["parent"]:
                break
            
            current_entity = record["parent"]
            res.append(current_entity)
    
    return res

def find_path(entity1, entity2, working_dir, level, depth=5):
    """Find path between two entities in Memgraph"""
    driver = get_graph_driver()
    
    with driver.session() as session:
        # Since level property doesn't exist on relationships, don't filter by it
        query = """
        MATCH path = (start:Entity {name: $entity1})-[*1..{depth}]->(end:Entity {name: $entity2})
        RETURN [node IN nodes(path) | node.name] AS path
        ORDER BY length(path) ASC
        LIMIT 1
        """
        result = session.run(query, entity1=entity1, entity2=entity2, depth=depth)
        record = result.single()
        
        if record:
            return record["path"]
        else:
            return None


def fetch_paths_between_entities(entity_pairs, max_depth: int = 4):
    """Retrieve shortest paths between entity pairs with relationship metadata."""

    if not entity_pairs:
        return []

    depth = max(1, int(max_depth or 1))
    query = f"""
        MATCH path = shortestPath((start:Entity {{name: $source}})-[:RELATES_TO*1..{depth}]-(end:Entity {{name: $target}}))
        RETURN [node IN nodes(path) | node.name] AS nodes,
               [rel IN relationships(path) | {{
                    start: startNode(rel).name,
                    end: endNode(rel).name,
                    description: rel.description,
                    weight: coalesce(rel.weight, 1.0),
                    level: coalesce(rel.level, 0)
               }}] AS relationships
        LIMIT 1
    """

    driver = get_graph_driver()
    results = []

    with driver.session() as session:
        for source, target in entity_pairs:
            if source == target:
                continue
            record = session.run(query, source=source, target=target).single()
            if record and record["nodes"]:
                results.append({
                    "pair": (source, target),
                    "nodes": record["nodes"],
                    "relationships": record["relationships"],
                })

    return results

def search_nodes_link(entity1, entity2, working_dir, level=0):
    """Search for direct link between two entities in Memgraph"""
    driver = get_graph_driver()
    
    with driver.session() as session:
        # Since level property doesn't exist on relationships, don't filter by it
        query = """
        MATCH (e1:Entity {name: $entity1})-[r:RELATES_TO]-(e2:Entity {name: $entity2})
        RETURN r, e1, e2
        """
        result = session.run(query, entity1=entity1, entity2=entity2)
        record = result.single()
        
        if record:
            return {
                "src_tgt": record["e1"]["name"],
                "tgt_src": record["e2"]["name"],
                "description": record["r"]["description"],
                "weight": record["r"]["weight"],
                "level": record["r"].get("level", 0)  # Use get() to handle missing level property
            }
        else:
            return None
def search_chunks(working_dir, entity_set):
    db_path = get_db_path()
    db = sqlite3.connect(db_path)
    res = []
    cursor = db.cursor()
    for entity in entity_set:
        if entity == 'root':
            continue
        sql = "SELECT source_id FROM entities WHERE entity_name = ?"
        cursor.execute(sql, (entity,))
        ret = cursor.fetchall()
        if ret:
            res.append(ret[0])
    db.close()
    return res
def search_nodes(entity_set, working_dir):
    """Search for entities in Memgraph"""
    driver = get_graph_driver()
    res = []
    
    with driver.session() as session:
        for entity in entity_set:
            query = """
            MATCH (e:Entity {name: $entity, level: 0})
            RETURN e.name AS entity_name, e.description, e.source_id, e.degree, e.parent, e.level
            """
            result = session.run(query, entity=entity)
            records = result.data()
            for record in records:
                res.append((
                    record["entity_name"],
                    record["description"],
                    record["source_id"],
                    record["degree"],
                    record["parent"],
                    record["level"]
                ))
    
    return res
def get_text_units(working_dir, chunks_set, chunks_file=None, k=5, max_length: int = 240):
    """Get trimmed text-unit previews from the database based on chunk hashes."""
    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()

    chunks_list = []
    for chunks in chunks_set:
        if chunks and "|" in chunks:
            temp_chunks = chunks.split("|")
        elif chunks:
            temp_chunks = [chunks]
        else:
            temp_chunks = []
        chunks_list += temp_chunks
    
    counter = Counter(chunks_list)
    
    # Select most frequent chunks, up to k
    duplicates = [item for item, _ in sorted(
        [(item, count) for item, count in counter.items() if count > 1],
        key=lambda x: x[1],
        reverse=True
    )[:k]]
    
    if len(duplicates) < k:
        used = set(duplicates)
        for item, _ in counter.items():
            if item not in used:
                duplicates.append(item)
                used.add(item)
            if len(duplicates) == k:
                break
    
    previews: list[str] = []
    if duplicates:
        # Query database for text units
        placeholders = ','.join('?' * len(duplicates))
        query = f"SELECT text FROM text_units WHERE hash_code IN ({placeholders})"
        cur.execute(query, duplicates)
        results = cur.fetchall()
        
        for (text,) in results:
            preview = (text or '').strip()
            if max_length and len(preview) > max_length:
                preview = preview[:max_length].rstrip() + "..."
            if preview:
                previews.append(preview)

    con.close()
    return previews

def get_entity_levels(names: list[str]) -> dict[str, int]:
    """Return mapping of entity name to hierarchy level from SQLite."""

    if not names:
        return {}

    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()

    placeholders = ','.join('?' for _ in names)
    query = f"SELECT entity_name, level FROM entities WHERE entity_name IN ({placeholders})"
    cur.execute(query, names)
    rows = cur.fetchall()

    con.close()
    return {name: level for name, level in rows}


def search_community(entity_name, working_dir):
    """Search for community in SQLite database"""
    db_path = get_db_path()
    con = sqlite3.connect(db_path)
    cur = con.cursor()
    
    cur.execute("""
        SELECT entity_name, entity_description, findings
        FROM communities
        WHERE entity_name = ?
    """, (entity_name,))
    
    result = cur.fetchone()
    con.close()
    
    if result:
        return result
    else:
        return ""
    

def insert_origin_relations(working_dir):
    """deprecated ???"""
    db_path = get_db_path()
    db = sqlite3.connect(db_path)
    cursor = db.cursor()
    
    relation_path = os.path.join("hi_ex", os.path.basename(working_dir), "relation.jsonl")
    with open(relation_path, "r") as f:
        val = []
        for relation_l in f:
            relation = json.loads(relation_l)
            src_tgt = relation['src_tgt']
            tgt_src = relation['tgt_src']
            if len(src_tgt) > 190 or len(tgt_src) > 190:
                print(f"src_tgt or tgt_src too long: {src_tgt} {tgt_src}")
                continue
            description = relation['description']
            weight = relation['weight']
            level = 0
            val.append((src_tgt, tgt_src, description, weight, level))
        
        cursor.executemany("INSERT INTO relations(src_tgt, tgt_src, description, weight, level) VALUES (?, ?, ?, ?, ?)", val)
        db.commit()
    
    db.close()