"""
=============================================================================
第 8 讲 · 工具模块：一个迷你的 SPARQL 引擎（只支持本讲用到的子集）
=============================================================================

对应讲义：lectures/lecture08_spinach.md 第 1.3 节

-----------------------------------------------------------------------------
【SPARQL 是什么？】
-----------------------------------------------------------------------------
知识图谱由**三元组**组成：(主语, 谓语, 宾语)，例如 (wd:Q41506 斯坦福大学, wdt:P112 创办者, wd:Q... 某人)。
SPARQL 是查询三元组的语言。一条查询描述一个"图案"，引擎找出图谱里所有能"套上"这个图案的方式：

    SELECT ?x WHERE { wd:Q41506 wdt:P112 ?x . }          # 斯坦福大学的创办者是谁？

    · ?x 是变量；wd:Q41506 是一个实体（QID）；wdt:P112 是一个属性（PID）；
    · 花括号里的每一行是一个"三元组模式"，用 " . " 隔开；多个模式共用变量，就是 SQL 里的连接（JOIN）。

-----------------------------------------------------------------------------
【这个引擎支持什么？】
-----------------------------------------------------------------------------
    SELECT [DISTINCT] ?a ?b (COUNT(DISTINCT ?c) AS ?n) WHERE { … } [GROUP BY ?a] [ORDER BY DESC(?n)] [LIMIT 5]
    ASK WHERE { … }                                                   # 有没有匹配？返回 True / False

    三元组模式里的谓语可以是"属性路径"：
        wdt:P31/wdt:P279*      先走一步 P31，再走零步或多步 P279（"是某一类，或某一类的子类"）
        wdt:P131+              一步或多步
        ^wdt:P184              反方向走（"谁的导师是 X" = X 的学生）
    FILTER(?d < "1950-01-01")、FILTER(YEAR(?d) >= 1990)、FILTER(?e > 3000)、FILTER(?x != wd:Q5)
    rdfs:label：实体的名字（字符串），例如 ?x rdfs:label "Northgate University"

不支持：OPTIONAL、UNION、VALUES、子查询、语言标签……（真实的 Wikidata 都支持）。

-----------------------------------------------------------------------------
【四种反馈】（SPINACH 论文第 4.1 节：执行一条查询，智能体会得到四种反馈之一）
-----------------------------------------------------------------------------
    1. 有结果  2. 没有结果  3. 语法错误（SPARQLSyntaxError，带错误信息）  4. 超时（SPARQLTimeout）
Wikidata 把每条查询的执行时间限制在 60 秒；这里用"中间结果的行数"模拟：超过 max_rows 就算超时。
"""

import re


class SPARQLSyntaxError(Exception):
    """查询写错了（括号不配对、不认识的写法……）。错误信息会原样返回给智能体。"""


class SPARQLTimeout(Exception):
    """查询太"宽"了：中间结果超过了上限（模拟 Wikidata 的 60 秒超时）。"""


# =============================================================================
# 第 1 部分：切词
# =============================================================================
_TOKEN = re.compile(r'"[^"]*"|\(|\)|\{|\}|,|<=|>=|!=|[<>=]|\.(?=\s|$|\})|[^\s(){},]+')


def tokenize(text):
    """
    把查询切成"词"：带引号的字符串算一个词，括号和比较符号单独算一个词。

    【例子】tokenize('SELECT ?x WHERE { ?x rdfs:label "New York" . }')
            → ['SELECT', '?x', 'WHERE', '{', '?x', 'rdfs:label', '"New York"', '.', '}']
    【边界】"." 只有后面是空格、"}" 或结尾时才算分隔符，这样 "1.5" 这样的数字不会被切开。
    【Python 小知识】re.findall 按正则表达式从左到右找出所有匹配；"|" 表示"或"，排在前面的写法优先。
    """
    return _TOKEN.findall(text)


# =============================================================================
# 第 2 部分：解析
# =============================================================================
def _parse_term(tok):
    """
    把一个词变成"项"：变量、实体、字面值。

    【返回】("var", "x") / ("iri", "wd:Q5") / ("lit", "New York" 或 1984 或 3.5)
    """
    if tok.startswith("?"):
        return ("var", tok[1:])
    if tok.startswith('"'):
        return ("lit", tok[1:-1])
    if re.fullmatch(r"-?\d+", tok):
        return ("lit", int(tok))
    if re.fullmatch(r"-?\d+\.\d+", tok):
        return ("lit", float(tok))
    if re.fullmatch(r"wd:Q\d+", tok):
        return ("iri", tok)
    raise SPARQLSyntaxError(f"Unrecognized term: {tok}")


def _parse_path(tok):
    """
    把谓语（可能是属性路径）解析成一串"步"。

    【返回】list[(谓语, 修饰, 是否反向)]，例如 "wdt:P31/wdt:P279*" → [("wdt:P31", "", False), ("wdt:P279", "*", False)]
    【边界】谓语必须是 wdt: / p: / ps: / pq: 加 PID，或者 rdfs:label。
    """
    steps = []
    for part in tok.split("/"):
        inverse = part.startswith("^")
        part = part.lstrip("^")
        mod = part[-1] if part[-1] in "*+" else ""
        pred = part.rstrip("*+")
        if not re.fullmatch(r"(wdt|p|ps|pq):P\d+|rdfs:label", pred):
            raise SPARQLSyntaxError(f"Unrecognized predicate: {part}")
        steps.append((pred, mod, inverse))
    return steps


def parse(query):
    """
    把一条 SPARQL 查询解析成一个字典。

    【返回】{"ask": bool, "distinct": bool, "proj": [...], "patterns": [...], "filters": [...],
            "group": [...], "order": [(变量, 是否降序)], "limit": int 或 None}
        proj 里每一项是 ("var", 名字) 或 ("count", 被数的变量, 是否 DISTINCT, 结果变量名)
        patterns 里每一项是 (主语项, 路径, 宾语项)
        filters 里每一项是 (左边, 比较符, 右边)，左边是 ("var", x) 或 ("year", x)
    【抛出】SPARQLSyntaxError：任何不认识的写法
    """
    toks = tokenize(query)
    pos = 0

    def peek():
        return toks[pos].upper() if pos < len(toks) else ""

    def take(expected=None):
        nonlocal pos
        if pos >= len(toks):
            raise SPARQLSyntaxError(f"Unexpected end of query (expected {expected or 'more tokens'})")
        tok = toks[pos]
        if expected and tok.upper() != expected:
            raise SPARQLSyntaxError(f"Expected '{expected}' but found '{tok}'")
        pos += 1
        return tok

    q = {"ask": False, "distinct": False, "proj": [], "patterns": [], "filters": [], "group": [],
         "order": [], "limit": None}
    head = take().upper()
    if head == "ASK":
        q["ask"] = True
    elif head == "SELECT":
        if peek() == "DISTINCT":
            take()
            q["distinct"] = True
        while peek() not in ("WHERE", ""):
            if peek() == "(":                                      # (COUNT(DISTINCT ?x) AS ?n)
                take("(")
                take("COUNT")
                take("(")
                dist = False
                if peek() == "DISTINCT":
                    take()
                    dist = True
                var = _parse_term(take())
                take(")")
                take("AS")
                out = _parse_term(take())
                take(")")
                if var[0] != "var" or out[0] != "var":
                    raise SPARQLSyntaxError("COUNT needs variables")
                q["proj"].append(("count", var[1], dist, out[1]))
            else:
                term = _parse_term(take())
                if term[0] != "var":
                    raise SPARQLSyntaxError(f"Can only project variables, found {term}")
                q["proj"].append(("var", term[1]))
        if not q["proj"]:
            raise SPARQLSyntaxError("SELECT needs at least one variable")
    else:
        raise SPARQLSyntaxError(f"Query must start with SELECT or ASK, found '{head}'")
    take("WHERE")
    take("{")
    while peek() != "}":
        if peek() == "":
            raise SPARQLSyntaxError("Missing '}'")
        if peek() == ".":
            take()
            continue
        if peek() == "FILTER":
            take()
            take("(")
            if peek() == "YEAR":
                take()
                take("(")
                left = ("year", _parse_term(take())[1])
                take(")")
            else:
                t = _parse_term(take())
                if t[0] != "var":
                    raise SPARQLSyntaxError("FILTER must start with a variable")
                left = ("var", t[1])
            op = take()
            if op not in ("=", "!=", "<", ">", "<=", ">="):
                raise SPARQLSyntaxError(f"Unknown comparison '{op}'")
            right = _parse_term(take())
            take(")")
            q["filters"].append((left, op, right))
            continue
        s = _parse_term(take())
        path = _parse_path(take())
        o = _parse_term(take())
        q["patterns"].append((s, path, o))
    take("}")
    while pos < len(toks):
        kw = take().upper()
        if kw == "GROUP":
            take("BY")
            while pos < len(toks) and toks[pos].startswith("?"):
                q["group"].append(take()[1:])
        elif kw == "ORDER":
            take("BY")
            while pos < len(toks) and (peek() in ("DESC", "ASC") or toks[pos].startswith("?")):
                if peek() in ("DESC", "ASC"):
                    desc = take().upper() == "DESC"
                    take("(")
                    var = take()[1:]
                    take(")")
                else:
                    desc, var = False, take()[1:]
                q["order"].append((var, desc))
        elif kw == "LIMIT":
            n = take()
            if not n.isdigit():
                raise SPARQLSyntaxError(f"LIMIT needs a number, found '{n}'")
            q["limit"] = int(n)
        else:
            raise SPARQLSyntaxError(f"Unexpected '{kw}' after WHERE clause")
    if not q["patterns"]:
        raise SPARQLSyntaxError("WHERE clause has no triple patterns")
    return q


# =============================================================================
# 第 3 部分：执行
# =============================================================================
def sort_key(v):
    """
    给图谱里的值一个固定的排序：实体按 QID 的数字，字符串按字母，数字按大小。

    【为什么】Python 的集合（set）遍历顺序可能每次运行都不一样（受 PYTHONHASHSEED 影响），
        结果行的顺序（以及 LIMIT 截到哪几行）就会变。凡是从集合变成列表的地方都按这个键排序，
        保证每次运行的结果完全一样。
    """
    if isinstance(v, str) and v.startswith("wd:Q"):
        return (0, int(v[4:]), "")
    if isinstance(v, str):
        return (1, 0, v)
    return (2, v, "")


class Engine:
    """
    在一个三元组的集合上执行 SPARQL。

    【用法】
        engine = Engine(triples)                   # triples：(主语, 谓语, 宾语) 的列表
        rows = engine.run("SELECT ?x WHERE { … }")  # 返回 list[tuple]；ASK 查询返回 [(True,)] 或 [(False,)]
    【参数】max_rows：中间结果的上限，超过就抛出 SPARQLTimeout（模拟 Wikidata 的 60 秒超时）
    【思路】
        · 建两个索引：sp[主语][谓语] → 宾语列表，po[谓语][宾语] → 主语列表；
        · 从"一个空的绑定"开始，逐个三元组模式扩展：每个模式把每个已有的绑定扩展成零个或多个新绑定
          （这就是 JOIN）；
        · 最后做 FILTER、分组计数、排序、投影、去重、截断。
    """

    def __init__(self, triples, max_rows=20000):
        self.max_rows = max_rows
        self.sp, self.po, self.by_pred = {}, {}, {}
        for s, p, o in triples:
            self.sp.setdefault(s, {}).setdefault(p, []).append(o)
            self.po.setdefault(p, {}).setdefault(o, []).append(s)
            self.by_pred.setdefault(p, []).append((s, o))

    # ---- 属性路径 ----
    def _step(self, nodes, step):
        """
        从一组节点出发，沿一"步"走：返回能到达的节点集合。

        【思路】反向的步（^wdt:P184）用 po 索引；"*"（零步或多步）和"+"（一步或多步）用广度优先搜索：
            不断沿这个谓语往外走，直到没有新节点。"*" 还包括出发点本身。
        """
        pred, mod, inverse = step
        index = self.po.get(pred, {}) if inverse else None

        def one(n):
            if inverse:
                return index.get(n, [])
            return self.sp.get(n, {}).get(pred, [])

        if mod == "":
            return {m for n in nodes for m in one(n)}
        reached = set(nodes) if mod == "*" else set()
        frontier = set(nodes)
        while frontier:
            nxt = {m for n in frontier for m in one(n)} - reached
            reached |= nxt
            frontier = nxt
        return reached

    def _reverse(self, path):
        """把一条路径倒过来（主语未知、宾语已知时，从宾语往回走）。"""
        return [(p, mod, not inv) for p, mod, inv in reversed(path)]

    def _pairs(self, s, path, o):
        """
        给定主语（可能未知，None）、路径、宾语（可能未知，None），返回所有能连起来的 (主语, 宾语)。

        【边界】主语和宾语都未知、路径又以 "*" 开头时，"零步"会让每个节点都和自己配对 —— 这会非常大，
            这里只从"有这个谓语的主语"出发，避免把整个图谱都当成起点。
        """
        if s is not None:
            nodes = {s}
            for step in path:
                nodes = self._step(nodes, step)
            return [(s, x) for x in sorted(nodes, key=sort_key) if o is None or x == o]
        if o is not None:
            nodes = {o}
            for step in self._reverse(path):
                nodes = self._step(nodes, step)
            return [(x, o) for x in sorted(nodes, key=sort_key)]
        first, mod, inv = path[0]
        starts = {a for a, b in self.by_pred.get(first, [])} if not inv else {b for a, b in self.by_pred.get(first, [])}
        out = []
        for st in sorted(starts, key=sort_key):
            out += self._pairs(st, path, None)
            if len(out) > self.max_rows:
                raise SPARQLTimeout("Query timed out")
        return out

    # ---- 主循环 ----
    def run(self, query):
        """执行一条查询。【返回】list[tuple]【抛出】SPARQLSyntaxError、SPARQLTimeout"""
        q = parse(query)
        solutions = [{}]
        for s, path, o in q["patterns"]:
            new = []
            for b in solutions:
                sv = b.get(s[1]) if s[0] == "var" else s[1]
                ov = b.get(o[1]) if o[0] == "var" else o[1]
                for a, c in self._pairs(sv, path, ov):
                    nb = dict(b)
                    if s[0] == "var":
                        nb[s[1]] = a
                    if o[0] == "var":
                        if o[1] in nb and nb[o[1]] != c:
                            continue
                        nb[o[1]] = c
                    new.append(nb)
                    if len(new) > self.max_rows:
                        raise SPARQLTimeout("Query timed out")
            solutions = new
        solutions = [b for b in solutions if all(self._check(f, b) for f in q["filters"])]
        if q["ask"]:
            return [(bool(solutions),)]
        return self._finish(q, solutions)

    @staticmethod
    def _check(f, b):
        """一个 FILTER 条件是否成立。变量没有绑定、类型不能比较时，算作不成立（和 SPARQL 一样）。"""
        (kind, var), op, right = f
        if var not in b:
            return False
        left = b[var]
        if kind == "year":
            m = re.match(r"(-?\d{4})", str(left))
            if not m:
                return False
            left = int(m.group(1))
        rv = right[1]
        try:
            return {"=": left == rv, "!=": left != rv, "<": left < rv, ">": left > rv,
                    "<=": left <= rv, ">=": left >= rv}[op]
        except TypeError:                                          # 例如字符串和数字比较
            return False

    @staticmethod
    def _finish(q, solutions):
        """分组计数、排序、投影、去重、截断。"""
        counts = [p for p in q["proj"] if p[0] == "count"]
        if counts:
            groups = {}
            for b in solutions:
                key = tuple(b.get(g) for g in q["group"])
                groups.setdefault(key, []).append(b)
            if not q["group"] and not groups:
                groups[()] = []
            rows = []
            for key, bs in groups.items():
                row = dict(zip(q["group"], key))
                for _k, var, dist, out in counts:
                    vals = [b[var] for b in bs if var in b]
                    row[out] = len(set(vals)) if dist else len(vals)
                rows.append(row)
            solutions = rows
        for var, desc in reversed(q["order"]):                     # 从最后一个排序键往前排（稳定排序）
            solutions.sort(key=lambda b: (b.get(var) is None, b.get(var) if b.get(var) is not None else 0),
                           reverse=desc)
        names = [p[1] if p[0] == "var" else p[3] for p in q["proj"]]
        rows = [tuple(b.get(n) for n in names) for b in solutions]
        if q["distinct"]:
            seen, uniq = set(), []
            for r in rows:
                if r not in seen:
                    seen.add(r)
                    uniq.append(r)
            rows = uniq
        if q["limit"] is not None:
            rows = rows[:q["limit"]]
        return rows


if __name__ == "__main__":
    triples = [("wd:Q1", "wdt:P31", "wd:Q5"), ("wd:Q1", "rdfs:label", "Ada"), ("wd:Q2", "wdt:P31", "wd:Q5"),
               ("wd:Q1", "wdt:P1303", "wd:Q10"), ("wd:Q10", "wdt:P279", "wd:Q11"), ("wd:Q11", "wdt:P279", "wd:Q12")]
    e = Engine(triples)
    print(e.run('SELECT ?x WHERE { ?x wdt:P31 wd:Q5 . }'))
    print(e.run('SELECT ?p ?i WHERE { ?p wdt:P1303 ?i . ?i wdt:P279* wd:Q12 . }'))
    print(e.run('SELECT (COUNT(?x) AS ?n) WHERE { ?x wdt:P31 wd:Q5 . }'))
    print(e.run('ASK WHERE { wd:Q1 rdfs:label "Ada" . }'))
    for bad in ['SELECT ?x WHERE { ?x wdt:P31 wd:Q5 ', 'SELECT ?x WHERE { ?x foo:P31 wd:Q5 . }']:
        try:
            e.run(bad)
        except SPARQLSyntaxError as err:
            print("语法错误：", err)