#!/usr/bin/env python3
"""
lead_multi_criteria_v2.py — Run the slow multi-criteria KGQuery against lead data.

Executes the multi-criteria SPARQL query (MQL=true + MQLRating>=65 + CompanyStateCode='CA')
via the v2 SPARQL-to-SQL pipeline, outputs the generated SQL and EXPLAIN ANALYZE.

Usage:
    python vitalgraph_sparql_sql/scripts/lead_multi_criteria_v2.py
    python -m vitalgraph_sparql_sql.scripts.lead_multi_criteria_v2
"""

import asyncio
import logging
import os
import re
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..'))

from vitalgraph.db.jena_sparql.jena_sidecar_client import SidecarClient
from vitalgraph.db.jena_sparql.jena_ast_mapper import map_compile_response
from vitalgraph.db.sparql_sql.generator import generate_sql as v2_generate, warm_stats_cache
from vitalgraph.db.sparql_sql import db_provider
from vitalgraph_sparql_sql import db
from vitalgraph_sparql_sql.db import DevDbImpl

SPACE_ID = "sp_sql_lead_dataset"

logger = logging.getLogger(__name__)


def _format_sql(sql: str, indent: int = 0) -> str:
    """Format SQL for display: collapse whitespace, indent keywords."""
    sql = re.sub(r'\s+', ' ', sql.strip())
    pad = ' ' * indent
    for kw in ['FROM', 'JOIN', 'LEFT JOIN', 'WHERE', 'GROUP BY',
               'ORDER BY', 'LIMIT', 'ON ', 'AND ']:
        sql = sql.replace(f' {kw}', f'\n{pad}  {kw}')
    sql = sql.replace('WITH ', f'WITH\n{pad}  ')
    sql = sql.replace(') SELECT ', f')\n{pad}SELECT ')
    return sql


async def _run_query(label: str, sparql: str, sidecar: SidecarClient) -> dict:
    """Compile, generate, execute, and EXPLAIN a single SPARQL query via v2."""
    print(f"\n{'─' * 70}")
    print(f"  {label}")
    print(f"{'─' * 70}")

    # Compile via sidecar
    t0 = time.monotonic()
    raw_json = sidecar.compile(sparql)
    sidecar_ms = (time.monotonic() - t0) * 1000

    cr = map_compile_response(raw_json)
    if not cr.ok:
        print(f"  ERROR: {cr.error}")
        return {}

    async with db.get_connection() as conn:
        # Generate SQL via v2
        t0 = time.monotonic()
        gen = await v2_generate(cr, SPACE_ID, conn=conn)
        gen_ms = (time.monotonic() - t0) * 1000

        sql = gen.sql
        var_map = gen.var_map or {}

        print(f"\n  SQL ({len(sql)} chars, {sql.upper().count(' JOIN ')} JOINs)")
        print(_format_sql(sql, indent=2))
        print()

        # Execute
        t0 = time.monotonic()
        rows = await db.execute_query(sql, conn=conn)
        exec_ms = (time.monotonic() - t0) * 1000

        # EXPLAIN ANALYZE
        explain_rows = []
        try:
            raw = await db.execute_query(f"EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) {sql}", conn=conn)
            explain_rows = [list(r.values())[0] for r in raw]
        except Exception as e:
            explain_rows = [f"EXPLAIN failed: {e}"]

    # ── Timing Summary ────────────────────────────────────────────────
    print()
    print(f"  {'Metric':<20} {'Value':>12}")
    print(f"  {'─' * 34}")
    print(f"  {'Rows':<20} {len(rows):>12}")
    print(f"  {'Sidecar':<20} {sidecar_ms:>11.0f}ms")
    print(f"  {'Generate (v2)':<20} {gen_ms:>11.0f}ms")
    print(f"  {'Execute':<20} {exec_ms:>11.0f}ms")
    print(f"  {'SQL chars':<20} {len(sql):>12}")
    print(f"  {'─' * 34}")

    # ── EXPLAIN ANALYZE ───────────────────────────────────────────────
    print()
    print(f"  EXPLAIN ANALYZE (BUFFERS):")
    print(f"  {'─' * 66}")
    for line in explain_rows:
        print(f"  {line}")
    print()

    return {
        "rows": rows,
        "row_count": len(rows),
        "sql": sql,
        "explain": explain_rows,
        "sidecar_ms": sidecar_ms,
        "gen_ms": gen_ms,
        "exec_ms": exec_ms,
    }


async def run():
    """Run the multi-criteria query against lead dataset."""

    # The exact SPARQL generated by KGQueryCriteriaBuilder for the multi-criteria test
    multi_criteria_sparql = """
        PREFIX haley: <http://vital.ai/ontology/haley-ai-kg#>
        PREFIX vital-core: <http://vital.ai/ontology/vital-core#>
        PREFIX vg-direct: <http://vital.ai/vitalgraph/direct#>
        PREFIX rdf: <http://www.w3.org/1999/02/22-rdf-syntax-ns#>
        PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
        PREFIX xsd: <http://www.w3.org/2001/XMLSchema#>

            SELECT DISTINCT ?entity WHERE {
                GRAPH <urn:sql_lead_dataset> {

        {
            ?entity vital-core:vitaltype haley:KGEntity .
        } UNION {
            ?entity vital-core:vitaltype haley:KGNewsEntity .
        } UNION {
            ?entity vital-core:vitaltype haley:KGProductEntity .
        } UNION {
            ?entity vital-core:vitaltype haley:KGWebEntity .
        }
        ?frame_edge_0 vital-core:vitaltype <http://vital.ai/ontology/haley-ai-kg#Edge_hasEntityKGFrame> .
        ?frame_edge_0 vital-core:hasEdgeSource ?entity .
        ?frame_edge_0 vital-core:hasEdgeDestination ?frame_0 .
        ?frame_0 haley:hasKGFrameType <urn:cardiff:kg:frame:LeadStatusFrame> .
        ?frame_edge_0_0 vital-core:vitaltype <http://vital.ai/ontology/haley-ai-kg#Edge_hasKGFrame> .
        ?frame_edge_0_0 vital-core:hasEdgeSource ?frame_0 .
        ?frame_edge_0_0 vital-core:hasEdgeDestination ?frame_0_0 .
        ?frame_0_0 haley:hasKGFrameType <urn:cardiff:kg:frame:LeadStatusQualificationFrame> .
        ?slot_edge_0_0_0 vital-core:vitaltype <http://vital.ai/ontology/haley-ai-kg#Edge_hasKGSlot> .
        ?slot_edge_0_0_0 vital-core:hasEdgeSource ?frame_0_0 .
        ?slot_edge_0_0_0 vital-core:hasEdgeDestination ?slot_0_0_0 .
        ?slot_0_0_0 haley:hasKGSlotType <urn:cardiff:kg:slot:MQLv2> .
        ?slot_0_0_0 haley:hasBooleanSlotValue true .
        ?slot_edge_0_0_1 vital-core:vitaltype <http://vital.ai/ontology/haley-ai-kg#Edge_hasKGSlot> .
        ?slot_edge_0_0_1 vital-core:hasEdgeSource ?frame_0_0 .
        ?slot_edge_0_0_1 vital-core:hasEdgeDestination ?slot_0_0_1 .
        ?slot_0_0_1 haley:hasKGSlotType <urn:cardiff:kg:slot:MQLRating> .
        ?slot_0_0_1 haley:hasDoubleSlotValue ?val_0_0_1 .
        FILTER(?val_0_0_1 >= 65.0)
        ?frame_edge_1 vital-core:vitaltype <http://vital.ai/ontology/haley-ai-kg#Edge_hasEntityKGFrame> .
        ?frame_edge_1 vital-core:hasEdgeSource ?entity .
        ?frame_edge_1 vital-core:hasEdgeDestination ?frame_1 .
        ?frame_1 haley:hasKGFrameType <urn:cardiff:kg:frame:CompanyFrame> .
        ?frame_edge_1_0 vital-core:vitaltype <http://vital.ai/ontology/haley-ai-kg#Edge_hasKGFrame> .
        ?frame_edge_1_0 vital-core:hasEdgeSource ?frame_1 .
        ?frame_edge_1_0 vital-core:hasEdgeDestination ?frame_1_0 .
        ?frame_1_0 haley:hasKGFrameType <urn:cardiff:kg:frame:CompanyAddressFrame> .
        ?slot_edge_1_0_0 vital-core:vitaltype <http://vital.ai/ontology/haley-ai-kg#Edge_hasKGSlot> .
        ?slot_edge_1_0_0 vital-core:hasEdgeSource ?frame_1_0 .
        ?slot_edge_1_0_0 vital-core:hasEdgeDestination ?slot_1_0_0 .
        ?slot_1_0_0 haley:hasKGSlotType <urn:cardiff:kg:slot:CompanyStateCode> .
        ?slot_1_0_0 haley:hasTextSlotValue "CA" .
                }
            }
            ORDER BY ?entity
            LIMIT 100
            OFFSET 0
    """

    print(f"lead_multi_criteria_v2 — SPARQL-to-SQL v2 pipeline")
    print(f"  space={SPACE_ID}")
    print(f"  query: MQL=true + MQLRating>=65 + CompanyStateCode='CA'")

    # Configure db_provider
    dev_impl = DevDbImpl()
    await dev_impl.connect()
    db_provider.configure(dev_impl)

    # Warm stats cache
    t0 = time.monotonic()
    async with db.get_connection() as conn:
        await warm_stats_cache(SPACE_ID, conn=conn)
    warmup_ms = (time.monotonic() - t0) * 1000
    print(f"  stats cache warmed in {warmup_ms:.0f}ms")

    sidecar = SidecarClient()

    print()
    print("=" * 70)
    print("Multi-criteria: MQL=true + MQLRating>=65 + CompanyStateCode='CA'")
    print("=" * 70)

    info = await _run_query("Multi-criteria query", multi_criteria_sparql, sidecar)

    if info.get("rows"):
        print(f"  Result entity URIs:")
        for row in info["rows"]:
            vals = list(row.values())
            print(f"    {vals[0] if vals else '?'}")

    sidecar.close()
    return 0


def main():
    logging.basicConfig(
        level=logging.WARNING,
        format="%(levelname)s:%(name)s:%(message)s",
    )
    sys.exit(asyncio.run(run()))


if __name__ == "__main__":
    main()
