#!/usr/bin/env python3
"""Test AutoMergingRetriever against the markdown + mermaid pipeline.

Verifies that the md ingestion DAG produced a hierarchy in which:
  1. Mermaid atomic leaves are retrievable from Milvus
  2. Each mermaid leaf carries a PARENT relationship to the mid-level parent
     that also owns the placeholder leaf
  3. AutoMergingRetriever merges sibling hits up to that shared parent

Usage:
    # Inspect the hierarchy for a doc (no querying)
    python scripts/test_automerge_mermaid.py --doc-id <uuid> --inspect

    # Run the built-in mermaid-themed query battery
    python scripts/test_automerge_mermaid.py --doc-id <uuid>

    # Single custom query
    python scripts/test_automerge_mermaid.py --doc-id <uuid> --query "..."
"""

import argparse
import json
import os
import sys
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import load_dotenv

load_dotenv()

from llama_index.core import Settings, StorageContext, VectorStoreIndex
from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core.schema import NodeRelationship
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.storage.docstore.postgres import PostgresDocumentStore
from llama_index.storage.kvstore.postgres import PostgresKVStore
from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.vector_stores.milvus.utils import BM25BuiltInFunction
from sqlalchemy import create_engine, event as sa_event
from sqlalchemy.ext.asyncio import create_async_engine

from src.core.config import settings


DEFAULT_QUERIES = [
    "ネットキャッシング 概要図",
    "リボ変更処理 フロー",
    "TR_CREDIT_ADD テーブル",
    "サービス債権 照会画面",
    "振込予定年月日 算出",
]

# RLS tenant/node — match the test upload script defaults.
DEFAULT_TENANT_ID = "28aa3971-48af-4f2a-aa73-9e4f89fbfba3"
DEFAULT_NODE_ID = "22222222-2222-2222-2222-222222222222"


def get_embed_model():
    return AzureOpenAIEmbedding(
        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
        api_version=os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-15-preview"),
        azure_deployment=os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "text-embedding-3-large"),
        model="text-embedding-3-large",
        embed_batch_size=16,
    )


def get_docstore(tenant_id: str, node_id: str) -> PostgresDocumentStore:
    """Build a docstore whose connections set RLS tenant context on connect.

    The docstore_table rows are protected by row-level security keyed on
    ``app.tenant_id`` / ``app.node_id`` — without setting these the query
    returns zero rows.
    """
    # Force localhost — DAG-side URL points at the docker network host ("postgres").
    async_url = settings.database_url.replace("@postgres:", "@localhost:")
    sync_url = async_url.replace("+asyncpg", "+psycopg2")

    async_engine = create_async_engine(
        async_url,
        connect_args={"server_settings": {"app.tenant_id": tenant_id, "app.node_id": node_id}},
    )
    sync_engine = create_engine(sync_url)

    @sa_event.listens_for(sync_engine, "connect")
    def _set_tenant(dbapi_conn, _record):
        cur = dbapi_conn.cursor()
        cur.execute(f"SET app.tenant_id = '{tenant_id}'; SET app.node_id = '{node_id}'")
        cur.close()

    kvstore = PostgresKVStore(
        table_name=settings.docstore_table,
        engine=sync_engine,
        async_engine=async_engine,
        perform_setup=False,
        use_jsonb=True,
    )
    return PostgresDocumentStore(postgres_kvstore=kvstore, namespace=settings.docstore_namespace)


def get_vector_store() -> MilvusVectorStore:
    # Force localhost — DAG-side host is the docker network alias ("milvus").
    host = "localhost" if settings.milvus_host in ("milvus", "postgres") else settings.milvus_host
    return MilvusVectorStore(
        uri=f"http://{host}:{settings.milvus_port}",
        collection_name=settings.milvus_collection,
        dim=settings.embedding_dimensions,
        overwrite=False,
        enable_sparse=True,
        sparse_embedding_function=BM25BuiltInFunction(),
        hybrid_ranker="RRFRanker",
        hybrid_ranker_params={"k": 60},
    )


def inspect_hierarchy(doc_id: str, tenant_id: str, node_id: str) -> dict:
    """Walk the docstore for one document and report parent/child wiring."""
    docstore = get_docstore(tenant_id, node_id)
    all_docs = docstore.docs  # {node_id: node}

    nodes = [
        n for n in all_docs.values()
        if n.ref_doc_id == doc_id or n.metadata.get("document_id") == doc_id
    ]
    if not nodes:
        print(f"No nodes found for document_id={doc_id}")
        return {"document_id": doc_id, "node_count": 0}

    by_id = {n.node_id: n for n in nodes}
    parent_ids = set()
    for n in nodes:
        parent_ref = n.relationships.get(NodeRelationship.PARENT)
        if parent_ref:
            parent_ids.add(parent_ref.node_id)

    leaves = [n for n in nodes if n.node_id not in parent_ids]
    mermaid_leaves = [n for n in leaves if n.metadata.get("node_type") == "mermaid_diagram"]
    prose_leaves = [n for n in leaves if n.metadata.get("node_type") != "mermaid_diagram"]

    # Group mermaid leaves by parent — verifies they share the placeholder leaf's parent.
    by_parent = defaultdict(list)
    orphan_mermaid = []
    for m in mermaid_leaves:
        p = m.relationships.get(NodeRelationship.PARENT)
        if not p:
            orphan_mermaid.append(m)
        else:
            by_parent[p.node_id].append(m)

    print(f"\n{'=' * 72}")
    print(f"Hierarchy for document_id={doc_id}")
    print(f"{'=' * 72}")
    print(f"Total nodes      : {len(nodes)}")
    print(f"Leaf nodes       : {len(leaves)}")
    print(f"  prose leaves   : {len(prose_leaves)}")
    print(f"  mermaid leaves : {len(mermaid_leaves)} "
          f"({len(orphan_mermaid)} orphan, {len(mermaid_leaves) - len(orphan_mermaid)} wired)")
    print(f"Parents involved : {len(parent_ids)}")

    if orphan_mermaid:
        print("\nWARNING: mermaid leaves with NO parent (auto-merging will not work):")
        for m in orphan_mermaid:
            print(f"  - {m.node_id} diagram_index={m.metadata.get('diagram_index')}")

    print("\nMermaid → parent groupings:")
    for parent_id, members in by_parent.items():
        parent = by_id.get(parent_id)
        parent_token = parent.metadata.get("token_count", "?") if parent else "?"
        diagrams = sorted({m.metadata.get("diagram_index") for m in members})
        # Find sibling placeholder leaves under the same parent
        sibling_placeholders = []
        for n in nodes:
            p = n.relationships.get(NodeRelationship.PARENT)
            if p and p.node_id == parent_id and n.metadata.get("node_type") != "mermaid_diagram":
                if "[MERMAID_DIAGRAM_" in n.text:
                    sibling_placeholders.append(n.node_id)
        print(f"  parent {parent_id[:8]}.. tokens={parent_token} diagrams={diagrams} "
              f"placeholder_siblings={len(sibling_placeholders)}")

    return {
        "document_id": doc_id,
        "node_count": len(nodes),
        "leaf_count": len(leaves),
        "prose_leaf_count": len(prose_leaves),
        "mermaid_leaf_count": len(mermaid_leaves),
        "orphan_mermaid_count": len(orphan_mermaid),
        "parent_count": len(parent_ids),
    }


def _node_summary(node, score=None) -> dict:
    meta = dict(node.metadata or {})
    node_type = meta.get("node_type", "prose")
    summary = {
        "node_id": node.node_id,
        "node_type": node_type,
        "score": float(score) if score is not None else None,
        "char_count": len(node.text),
        "token_count": meta.get("token_count"),
        "diagram_index": meta.get("diagram_index"),
        "split_index": meta.get("split_index"),
        "split_total": meta.get("split_total"),
        "parent_id": (
            node.relationships[NodeRelationship.PARENT].node_id
            if NodeRelationship.PARENT in node.relationships else None
        ),
        "child_count": (
            len(node.relationships[NodeRelationship.CHILD])
            if isinstance(node.relationships.get(NodeRelationship.CHILD), list)
            else (1 if NodeRelationship.CHILD in node.relationships else 0)
        ),
        "preview": node.text[:240].replace("\n", " ⏎ "),
    }
    return summary


def _load_all_nodes_in_memory(tenant_id: str, node_id: str) -> SimpleDocumentStore:
    """Pull every node visible under this tenant's RLS context into an in-memory
    SimpleDocumentStore. Lets the retriever resolve any parent_id Milvus returns,
    regardless of which document the leaf came from. Result-level filtering on
    ``document_id`` happens in ``query_once``."""
    pg_docstore = get_docstore(tenant_id, node_id)
    nodes = list(pg_docstore.docs.values())
    if not nodes:
        raise RuntimeError("No nodes visible under current RLS context")
    mem = SimpleDocumentStore()
    mem.add_documents(nodes)
    return mem


def query_once(query: str, doc_id: str, tenant_id: str, node_id: str,
               top_k: int, merge_threshold: float, max_preview: int) -> dict:
    print(f"\n{'=' * 72}")
    print(f"Query: {query!r}")
    print(f"top_k={top_k}  merge_threshold={merge_threshold}")
    print(f"{'=' * 72}")

    Settings.embed_model = get_embed_model()
    vector_store = get_vector_store()
    docstore = _load_all_nodes_in_memory(tenant_id, node_id)
    storage_context = StorageContext.from_defaults(vector_store=vector_store, docstore=docstore)

    index = VectorStoreIndex.from_vector_store(
        vector_store=vector_store,
        storage_context=storage_context,
    )

    base_retriever = index.as_retriever(similarity_top_k=top_k)
    retriever = AutoMergingRetriever(base_retriever, storage_context, simple_ratio_thresh=merge_threshold, verbose=False)

    def _from_target_doc(h):
        return (h.node.ref_doc_id == doc_id
                or h.node.metadata.get("document_id") == doc_id)

    # Capture both pre- and post-merge, then keep only hits whose source is our target doc.
    base_hits = [h for h in base_retriever.retrieve(query) if _from_target_doc(h)]
    merged_hits = [h for h in retriever.retrieve(query) if _from_target_doc(h)]

    types_before = Counter(
        h.node.metadata.get("node_type", "prose") for h in base_hits
    )
    types_after = Counter(
        h.node.metadata.get("node_type", "prose") for h in merged_hits
    )
    merged_into_parent = max(0, len(base_hits) - len(merged_hits))

    print(f"  pre-merge hits  : {len(base_hits)}  types={dict(types_before)}")
    print(f"  post-merge hits : {len(merged_hits)} types={dict(types_after)}  (collapsed: {merged_into_parent})")

    base_summaries = [_node_summary(h.node, h.score) for h in base_hits]
    merged_summaries = [_node_summary(h.node, h.score) for h in merged_hits]

    print("\n  Top merged results:")
    for i, (h, s) in enumerate(zip(merged_hits, merged_summaries), 1):
        marker = "🔷" if s["node_type"] == "mermaid_diagram" else (
            "📚" if s["child_count"] else "📄"
        )
        preview = (h.node.text[:max_preview] + "...") if len(h.node.text) > max_preview else h.node.text
        score_str = f"{s['score']:.4f}" if s["score"] is not None else "NA"
        print(f"    {i:>2}. {marker} type={s['node_type']:<16} "
              f"chars={s['char_count']:>5} score={score_str}")
        print(f"        children={s['child_count']} parent={(s['parent_id'] or 'None')[:8]}..")
        if s["diagram_index"] is not None:
            print(f"        diagram_index={s['diagram_index']} split={s['split_index']}/{s['split_total']}")
        print(f"        preview: {preview[:max_preview]}")

    return {
        "query": query,
        "doc_id": doc_id,
        "top_k": top_k,
        "merge_threshold": merge_threshold,
        "pre_merge": {
            "count": len(base_hits),
            "type_counts": dict(types_before),
            "results": base_summaries,
        },
        "post_merge": {
            "count": len(merged_hits),
            "type_counts": dict(types_after),
            "collapsed": merged_into_parent,
            "results": merged_summaries,
        },
    }


def main():
    parser = argparse.ArgumentParser(description="Test AutoMergingRetriever for markdown+mermaid documents")
    parser.add_argument("--doc-id", required=True, help="Target document UUID")
    parser.add_argument("--query", default=None, help="Single query (default: built-in battery)")
    parser.add_argument("--top-k", type=int, default=8, help="Leaf nodes to retrieve before merging (default: 8)")
    parser.add_argument("--merge-threshold", type=float, default=0.5, help="AutoMerge ratio threshold (default: 0.5)")
    parser.add_argument("--max-preview", type=int, default=180, help="Preview chars per result (default: 180)")
    parser.add_argument("--inspect", action="store_true", help="Print hierarchy summary only; skip queries")
    parser.add_argument("--tenant-id", default=DEFAULT_TENANT_ID, help="RLS tenant UUID (default: test tenant)")
    parser.add_argument("--node-id", default=DEFAULT_NODE_ID, help="RLS node UUID (default: test node)")
    parser.add_argument("-o", "--output", type=Path, help="Optional JSON report path")
    args = parser.parse_args()

    if not os.getenv("AZURE_OPENAI_API_KEY") or not os.getenv("AZURE_OPENAI_ENDPOINT"):
        print("Error: AZURE_OPENAI_API_KEY / AZURE_OPENAI_ENDPOINT not set")
        sys.exit(1)

    report = {
        "timestamp": datetime.now().isoformat(),
        "collection": settings.milvus_collection,
        "docstore_table": settings.docstore_table,
        "docstore_namespace": settings.docstore_namespace,
    }

    report["hierarchy"] = inspect_hierarchy(args.doc_id, args.tenant_id, args.node_id)

    if args.inspect:
        print("\n(--inspect mode: skipping queries)")
    else:
        queries = [args.query] if args.query else DEFAULT_QUERIES
        report["queries"] = []
        for q in queries:
            try:
                report["queries"].append(query_once(
                    query=q,
                    doc_id=args.doc_id,
                    tenant_id=args.tenant_id,
                    node_id=args.node_id,
                    top_k=args.top_k,
                    merge_threshold=args.merge_threshold,
                    max_preview=args.max_preview,
                ))
            except Exception as e:
                print(f"\nQuery {q!r} failed: {e}")
                report["queries"].append({"query": q, "error": str(e)})

    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(report, f, indent=2, ensure_ascii=False, default=str)
        print(f"\nReport written to {args.output}")


if __name__ == "__main__":
    main()