#!/usr/bin/env python3
"""Original LockedIn Labs teaching lab. Synthetic data; in-memory SQLite only.

Python 3.11+, standard library. No account, model, network, or persistent files.
An edge A DEPENDS_ON B says A depends on B. Walk incoming edges from the target.
This is graph retrieval, not the complete Microsoft GraphRAG pipeline.
"""
import argparse
import json
import sqlite3
import unittest

NODES = [
    ("desk", "Service desk", "demo"),
    ("assistant", "Draft assistant", "demo"),
    ("search", "Procedure search", "demo"),
    ("catalog", "Source catalog", "demo"),
    ("model", "Model endpoint", "demo"),
    ("billing", "Billing export", "demo"),
    ("outside", "Other workspace", "other"),
]
SOURCES = [
    ("S1", "demo", "Desk integration", "v3", "Service desk calls Draft assistant to prepare response drafts.", 1, "available"),
    ("S2", "demo", "Assistant design", "v2", "Draft assistant retrieves approved procedures using Procedure search.", 1, "available"),
    ("S3", "demo", "Search contract", "v5", "Procedure search reads procedure metadata from Source catalog.", 1, "available"),
    ("S4", "demo", "Model connection", "v1", "Draft assistant sends the selected prompt to Model endpoint.", 1, "available"),
    ("S5", "demo", "Unreviewed note", "v1", "Billing export may read Source catalog; this claim has not been verified.", 1, "available"),
    ("S6", "demo", "Counterfactual cycle", "v1", "In this changed fixture, Source catalog delegates draft creation to Draft assistant.", 1, "available"),
    ("S7", "demo", "Retired integration", "v0", "A retired design connected Billing export to Procedure search.", 0, "available"),
    ("S8", "other", "Other workspace design", "v1", "Other workspace depends on Procedure search.", 1, "available"),
]
EDGES = [
    ("E1", "desk", "assistant", "DEPENDS_ON", "S1", 1),
    ("E2", "assistant", "search", "DEPENDS_ON", "S2", 1),
    ("E3", "search", "catalog", "DEPENDS_ON", "S3", 1),
    ("E4", "assistant", "model", "DEPENDS_ON", "S4", 1),
    ("E5", "billing", "catalog", "DEPENDS_ON", "S5", 0),
    ("E7", "billing", "search", "DEPENDS_ON", "S7", 1),
    ("E8", "outside", "search", "DEPENDS_ON", "S8", 1),
]
CYCLE = ("E6", "catalog", "assistant", "DEPENDS_ON", "S6", 1)
CASES = [
    ("baseline", "Follow supported relationships", 3, False),
    ("one-hop", "Stop after one relationship", 1, False),
    ("withdrawn", "Withdraw the search contract", 3, False),
    ("unverified", "Compare a weakened evidence rule", 3, True),
    ("cycle", "Add a circular dependency", 6, False),
]

ELIGIBLE = """
SELECT e.* FROM edges e
JOIN nodes a ON a.id = e.origin
JOIN nodes b ON b.id = e.target
JOIN sources s ON s.id = e.source
WHERE a.tenant = :tenant AND b.tenant = :tenant AND s.tenant = :tenant
  AND s.is_current = 1 AND s.state = 'available'
  AND e.relation = 'DEPENDS_ON'
  AND (e.verified = 1 OR :trust_unverified = 1)
"""
QUERY = """
WITH RECURSIVE eligible AS (""" + ELIGIBLE + """),
walk(node, node_path, edge_path, depth) AS (
  SELECT id, '|' || id || '|', '', 0 FROM nodes
    WHERE id = :target AND tenant = :tenant
  UNION ALL
  SELECT e.origin, '|' || e.origin || w.node_path,
         e.id || '|' || w.edge_path, w.depth + 1
  FROM walk w JOIN eligible e ON e.target = w.node
  WHERE w.depth < :max_hops
    AND instr(w.node_path, '|' || e.origin || '|') = 0
)
SELECT node_path, edge_path, depth FROM walk WHERE depth > 0
ORDER BY depth, node_path, edge_path
"""


def database(case="baseline"):
    db = sqlite3.connect(":memory:")
    db.row_factory = sqlite3.Row
    db.executescript("""
      PRAGMA foreign_keys = ON;
      CREATE TABLE nodes(id TEXT PRIMARY KEY, label TEXT, tenant TEXT);
      CREATE TABLE sources(id TEXT PRIMARY KEY, tenant TEXT, title TEXT,
        version TEXT, statement TEXT, is_current INTEGER, state TEXT);
      CREATE TABLE edges(id TEXT PRIMARY KEY, origin TEXT REFERENCES nodes(id),
        target TEXT REFERENCES nodes(id), relation TEXT,
        source TEXT REFERENCES sources(id), verified INTEGER);
    """)
    db.executemany("INSERT INTO nodes VALUES (?,?,?)", NODES)
    db.executemany("INSERT INTO sources VALUES (?,?,?,?,?,?,?)", SOURCES)
    db.executemany("INSERT INTO edges VALUES (?,?,?,?,?,?)", EDGES)
    if case == "withdrawn":
        db.execute("UPDATE sources SET state = 'withdrawn' WHERE id = 'S3'")
    if case == "cycle":
        db.execute("INSERT INTO edges VALUES (?,?,?,?,?,?)", CYCLE)
    # A finite instruction budget is a backstop, separate from the hop limit.
    calls = 0
    def limit_work():
        nonlocal calls
        calls += 1
        return int(calls > 1000)
    db.set_progress_handler(limit_work, 100)
    return db


def retrieve(db, max_hops=3, trust_unverified=False, target="catalog", tenant="demo"):
    if not isinstance(max_hops, int) or isinstance(max_hops, bool) or not 1 <= max_hops <= 6:
        raise ValueError("Choose a hop limit from 1 to 6.")
    params = dict(tenant=tenant, target=target, max_hops=max_hops,
                  trust_unverified=int(trust_unverified))
    rows = db.execute(QUERY, params).fetchall()
    return {
        "maxHops": max_hops,
        "trustUnverified": trust_unverified,
        "eligibleEdges": [row["id"] for row in db.execute(ELIGIBLE + " ORDER BY e.id", params)],
        "paths": [{"nodes": row["node_path"].strip("|").split("|"),
                   "edges": row["edge_path"].strip("|").split("|"),
                   "depth": row["depth"]} for row in rows],
    }


def run_case(case_id, max_hops=None):
    config = next((row for row in CASES if row[0] == case_id), None)
    if config is None:
        raise ValueError("Unknown case.")
    db = database(case_id)
    try:
        return retrieve(db, config[2] if max_hops is None else max_hops, config[3])
    finally:
        db.close()


def export():
    return {
        "engine": "SQLite recursive CTE · Python standard library",
        "question": "Which services have a documented dependency path to Source catalog?",
        "nodes": [{"id": row[0], "label": row[1]} for row in NODES if row[2] == "demo"],
        "sources": [{"id": s[0], "title": s[2], "version": s[3], "statement": s[4]} for s in SOURCES if s[1] == "demo" and s[5]],
        "edges": [{"id": e[0], "origin": e[1], "target": e[2], "source": e[4], "verified": bool(e[5])} for e in EDGES[:5] + [CYCLE]],
        "cases": [{"id": row[0], "label": row[1], **run_case(row[0])} for row in CASES],
    }


class ProofTests(unittest.TestCase):
    def test_direction_and_support(self):
        result = run_case("baseline")
        self.assertEqual([p["nodes"] for p in result["paths"]], [
            ["search", "catalog"], ["assistant", "search", "catalog"],
            ["desk", "assistant", "search", "catalog"]])
        self.assertEqual(result["paths"][-1]["edges"], ["E1", "E2", "E3"])
        self.assertEqual(result["eligibleEdges"], ["E1", "E2", "E3", "E4"])

    def test_missing_source_breaks_paths(self):
        self.assertEqual(run_case("withdrawn")["paths"], [])

    def test_checks_every_hop(self):
        with database() as db:
            db.execute("UPDATE sources SET state = 'withdrawn' WHERE id = 'S2'")
            self.assertEqual([p["nodes"][0] for p in retrieve(db)["paths"]], ["search"])

    def test_scope_and_retirement_even_when_verification_is_weakened(self):
        result = run_case("unverified")
        self.assertIn("E5", result["eligibleEdges"])
        self.assertNotIn("E7", result["eligibleEdges"])
        self.assertNotIn("E8", result["eligibleEdges"])
        with database() as db:
            db.execute("UPDATE sources SET tenant = 'other' WHERE id = 'S2'")
            self.assertEqual(len(retrieve(db)["paths"]), 1)

    def test_cycle_terminates_without_repeating_a_node(self):
        result = run_case("cycle")
        self.assertEqual(result["paths"], run_case("baseline")["paths"])
        self.assertIn("E6", result["eligibleEdges"])
        for path in result["paths"]:
            self.assertEqual(len(path["nodes"]), len(set(path["nodes"])))

    def test_depth_is_not_completeness(self):
        self.assertEqual(len(run_case("one-hop")["paths"]), 1)
        self.assertEqual(len(run_case("baseline", 2)["paths"]), 2)
        with database() as db:
            for invalid in [0, 7, -1, 1.5, True]:
                with self.assertRaises(ValueError):
                    retrieve(db, invalid)

    def test_parameterized_target_and_relationship_type(self):
        with database() as db:
            self.assertEqual(retrieve(db, target="catalog' OR 1=1 --")["paths"], [])
            db.execute("UPDATE edges SET relation = 'MENTIONS' WHERE id = 'E3'")
            self.assertEqual(retrieve(db)["paths"], [])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["run", "export", "test"])
    parser.add_argument("--case", default="baseline", choices=[row[0] for row in CASES])
    parser.add_argument("--max-hops", type=int, choices=range(1, 7))
    args = parser.parse_args()
    if args.command == "test":
        result = unittest.TextTestRunner().run(unittest.defaultTestLoader.loadTestsFromTestCase(ProofTests))
        raise SystemExit(0 if result.wasSuccessful() else 1)
    print(json.dumps(export() if args.command == "export" else run_case(args.case, args.max_hops), indent=2))


if __name__ == "__main__":
    main()