"""
query_similarity.py
====================
PoC Task 2 - Neo4j Vector Migration
Alat verifikasi & eksplorasi Knowledge Graph.

Menampilkan seluruh node :Literature dan relasi kemiripan
(edge SIMILAR_*) yang tersimpan di Neo4j AuraDB dalam format
terstruktur dan ramah-baca di terminal.

Cara menjalankan:
  python query_similarity.py
  python query_similarity.py --component methodology
  python query_similarity.py --min-score 0.80
  python query_similarity.py --node <article_id_prefix>
"""

import os
import sys
import argparse
from dotenv import load_dotenv
from neo4j import GraphDatabase

# --- Load credentials ---
load_dotenv()

NEO4J_URI      = os.getenv("NEO4J_URI", "")
NEO4J_USERNAME = os.getenv("NEO4J_USERNAME", "neo4j")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "")

# Komponen yang tersedia
COMPONENTS = ["background", "methodology", "gap", "objective", "futurework"]


# ============================================================
# QUERY HELPERS
# ============================================================

def get_all_nodes(driver) -> list[dict]:
    """Mengambil semua node :Literature beserta metadata ringkas."""
    cypher = """
    MATCH (n:Literature)
    RETURN elementId(n) AS id,
           n.title AS title
    ORDER BY n.title
    """
    with driver.session() as session:
        records = list(session.run(cypher))
        return [{"id": r["id"], "title": r["title"]} for r in records]


def get_similarity_edges(driver, component: str = None, min_score: float = 0.0, node_id_prefix: str = None) -> list[dict]:
    """
    Mengambil edge kemiripan dari Neo4j dengan filter opsional.

    Args:
        driver: Neo4j driver.
        component: Filter komponen tertentu (e.g., "methodology"). None = semua.
        min_score: Skor minimum cosine similarity.
        node_id_prefix: Prefix node identifier (elementId/hash/article_id) untuk filter node tertentu.
    """
    # Build WHERE clauses
    where_clauses = ["r.score >= $min_score"]

    if component:
        rel_type = f"SIMILAR_{component.upper()}"
        where_clauses.append(f"type(r) = '{rel_type}'")

    if node_id_prefix:
        where_clauses.append(
            "(" \
            "elementId(a) STARTS WITH $node_prefix OR elementId(b) STARTS WITH $node_prefix OR " \
            "coalesce(a.hash, '') STARTS WITH $node_prefix OR coalesce(b.hash, '') STARTS WITH $node_prefix OR " \
            "coalesce(a.article_id, '') STARTS WITH $node_prefix OR coalesce(b.article_id, '') STARTS WITH $node_prefix" \
            ")"
        )
    else:
        # Jika tidak ada filter tunggal, cegah duplikasi hasil (A-B dan B-A) karena undirected edge (RD-11)
        where_clauses.append("elementId(a) < elementId(b)")

    where_str = " AND ".join(where_clauses)

    cypher = f"""
    MATCH (a:Literature)-[r]-(b:Literature)
    WHERE type(r) STARTS WITH 'SIMILAR_' AND {where_str}
    RETURN elementId(a) AS source_id,
           a.title AS source_title,
           type(r) AS rel_type,
           r.score AS score,
            coalesce(r.reasoning, r.summary, '') AS reasoning,
            CASE
             WHEN r.reasoning IS NOT NULL AND trim(r.reasoning) <> '' THEN 'reasoning'
             WHEN r.summary IS NOT NULL AND trim(r.summary) <> '' THEN 'summary_legacy'
             ELSE 'none'
            END AS reasoning_source,
            coalesce(r.matched_pillars, []) AS matched_pillars,
            coalesce(r.gate, 'legacy') AS gate,
           elementId(b) AS target_id,
           b.title AS target_title
    ORDER BY rel_type, r.score DESC
    """
    

    params = {
        "min_score": min_score,
        "node_prefix": node_id_prefix or "",
    }

    with driver.session() as session:
        records = list(session.run(cypher, params))
        return [{
            "source_id":    r["source_id"],
            "source_title": r["source_title"],
            "rel_type":     r["rel_type"],
            "score":        r["score"],
            "reasoning":    r["reasoning"] or "",
            "reasoning_source": r["reasoning_source"],
            "matched_pillars": r["matched_pillars"] or [],
            "gate": r["gate"],
            "target_id":    r["target_id"],
            "target_title": r["target_title"],
        } for r in records]


def get_graph_stats(driver) -> dict:
    """Statistik ringkas tentang Knowledge Graph."""
    stats = {}
    with driver.session() as session:
        # Total node
        result = session.run("MATCH (n:Literature) RETURN count(n) AS total")
        stats["total_nodes"] = result.single()["total"]

        # Edge per tipe
        result = session.run("""
            MATCH ()-[r]->()
            WHERE type(r) STARTS WITH 'SIMILAR_'
            RETURN type(r) AS rel, count(r) AS cnt, avg(r.score) AS avg_score
            ORDER BY rel
        """)
        stats["edges"] = [{
            "type": r["rel"],
            "count": r["cnt"],
            "avg_score": round(r["avg_score"], 4) if r["avg_score"] else 0,
        } for r in result]

        stats["total_edges"] = sum(e["count"] for e in stats["edges"])

    return stats


# ============================================================
# DISPLAY FORMATTERS
# ============================================================

def print_header(title: str):
    """Mencetak header dengan border."""
    print(f"\n{'═' * 70}")
    print(f"  {title}")
    print(f"{'═' * 70}")


def print_section(title: str):
    """Mencetak sub-section."""
    print(f"\n{'─' * 70}")
    print(f"  {title}")
    print(f"{'─' * 70}")


def display_nodes(nodes: list[dict]):
    """Menampilkan daftar node."""
    print_section(f"📚 Node :Literature ({len(nodes)} total)")
    for i, node in enumerate(nodes, 1):
        print(f"  {i:>3}. [{node['id'][:12]}...] {node['title'][:55]}")


def display_stats(stats: dict):
    """Menampilkan statistik Knowledge Graph."""
    print_section("📊 Statistik Knowledge Graph")
    print(f"  🔷 Total node       : {stats['total_nodes']}")
    print(f"  🔗 Total edge       : {stats['total_edges']}")
    print()
    if stats["edges"]:
        print(f"  {'Tipe Relasi':<28} {'Jumlah':>8} {'Avg Score':>10}")
        print(f"  {'─' * 48}")
        for e in stats["edges"]:
            print(f"  {e['type']:<28} {e['count']:>8} {e['avg_score']:>10.4f}")


def display_edges(edges: list[dict], show_reasoning: bool = True):
    """Menampilkan daftar edge kemiripan."""
    if not edges:
        print("\n  ℹ️  Tidak ada edge yang sesuai filter.\n")
        return

    # Group by rel_type
    grouped = {}
    for e in edges:
        rt = e["rel_type"]
        if rt not in grouped:
            grouped[rt] = []
        grouped[rt].append(e)

    for rel_type, group in grouped.items():
        comp_name = rel_type.replace("SIMILAR_", "").lower()
        print_section(f"🔗 {rel_type} ({len(group)} edge)")

        for i, e in enumerate(group, 1):
            src_short = e["source_id"][:12]
            tgt_short = e["target_id"][:12]
            score_bar = "█" * int(e["score"] * 20) + "░" * (20 - int(e["score"] * 20))

            print(f"\n  [{i}] {src_short}... → {tgt_short}...")
            print(f"      📄 \"{e['source_title'][:50]}...\"")
            print(f"      📄 \"{e['target_title'][:50]}...\"")
            print(f"      📈 Score: {e['score']:.4f}  [{score_bar}]")
            print(f"      🧩 Evidence: source={e['reasoning_source']} | gate={e['gate']} | matched={','.join(e['matched_pillars']) if e['matched_pillars'] else '-'}")

            if show_reasoning and e["reasoning"]:
                # Wrap reasoning jika panjang
                reasoning = e["reasoning"]
                if len(reasoning) > 80:
                    print(f"      💬 \"{reasoning[:77]}...\"")
                else:
                    print(f"      💬 \"{reasoning}\"")


# ============================================================
# MAIN
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        description="📊 Knowledge Graph Explorer - PoC Task 2 Neo4j Vector Migration",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Contoh penggunaan:
  python query_similarity.py                          # Tampilkan semua
  python query_similarity.py --component methodology  # Filter komponen
  python query_similarity.py --min-score 0.80         # Filter skor minimum
    python query_similarity.py --node facb52d3          # Filter node (elementId/hash/article_id)
  python query_similarity.py --stats-only             # Statistik saja
        """,
    )
    parser.add_argument(
        "--component", "-c",
        choices=COMPONENTS,
        help="Filter berdasarkan komponen tertentu.",
    )
    parser.add_argument(
        "--min-score", "-s",
        type=float,
        default=0.0,
        help="Skor cosine similarity minimum (default: 0.0).",
    )
    parser.add_argument(
        "--node", "-n",
        type=str,
        default=None,
        help="Prefix node identifier (elementId/hash/article_id) untuk filter node tertentu.",
    )
    parser.add_argument(
        "--stats-only",
        action="store_true",
        help="Hanya tampilkan statistik, tanpa detail edge.",
    )
    parser.add_argument(
        "--no-reasoning", "--no-summary",
        action="store_true",
        help="Sembunyikan teks reasoning di output edge.",
    )

    args = parser.parse_args()

    # Validasi kredensial
    if not all([NEO4J_URI, NEO4J_USERNAME, NEO4J_PASSWORD]):
        print("❌ ERROR: Kredensial Neo4j belum lengkap di file .env")
        print("   Pastikan NEO4J_URI, NEO4J_USERNAME, NEO4J_PASSWORD terisi.")
        sys.exit(1)

    # Koneksi
    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
    try:
        driver.verify_connectivity()
    except Exception as e:
        print(f"❌ Gagal koneksi ke Neo4j: {e}")
        sys.exit(1)

    print_header("📊 Knowledge Graph Explorer — My-SRE PoC Task 2")

    # Tampilkan daftar node
    nodes = get_all_nodes(driver)
    display_nodes(nodes)

    # Tampilkan statistik
    stats = get_graph_stats(driver)
    display_stats(stats)

    # Jika stats-only, berhenti di sini
    if args.stats_only:
        driver.close()
        print(f"\n{'═' * 70}\n")
        return

    # Tampilkan filter yang aktif
    active_filters = []
    if args.component:
        active_filters.append(f"komponen={args.component}")
    if args.min_score > 0:
        active_filters.append(f"min_score={args.min_score}")
    if args.node:
        active_filters.append(f"node={args.node}")

    if active_filters:
        print(f"\n  🔍 Filter aktif: {', '.join(active_filters)}")

    # Query dan tampilkan edge
    edges = get_similarity_edges(
        driver,
        component=args.component,
        min_score=args.min_score,
        node_id_prefix=args.node,
    )
    display_edges(edges, show_reasoning=not args.no_reasoning)

    driver.close()

    print(f"\n{'═' * 70}")
    print(f"  ✅ Selesai. {len(nodes)} node, {len(edges)} edge ditampilkan.")
    print(f"{'═' * 70}\n")


if __name__ == "__main__":
    main()
