"""
=============================================================================
第 7 讲 · 工具模块：一个迷你的 SUQL —— 在 SQL 里调用"读文本"的 LLM 函数，以及一个做优化的执行器
=============================================================================

对应讲义：lectures/lecture07_suql.md 第 2~4 节

-----------------------------------------------------------------------------
【SUQL 是什么？】
-----------------------------------------------------------------------------
SUQL（Structured and Unstructured Query Language）= SQL + 两个处理自由文本的函数：

    answer(文本列, '问题')   → 让 LLM 读这一行的这段文字，回答这个问题（例如 'Yes' / 'No'、一个日期）
    summary(文本列)          → 让 LLM 给这段文字写一个摘要（相当于 answer(文本列, '摘要是什么？')）

它们可以出现在 SQL 里任何能写"一个文本值"的地方：SELECT 列表、WHERE 条件、ORDER BY、聚合函数里。
例如"Palo Alto 适合带孩子的餐厅"：

    SELECT *, summary(reviews) FROM restaurants
    WHERE location = 'Palo Alto' AND answer(reviews, 'is it family-friendly?') = 'Yes' LIMIT 1

-----------------------------------------------------------------------------
【这个模块提供什么？】
-----------------------------------------------------------------------------
    make_suql_world(seed)  第 6 讲的 600 家餐厅，加一列 reviews（全部评论拼成一段文字），
                           外加一个本讲才有的隐藏特征"停车方便"（只出现在部分餐厅的评论里）
    SimLLM                 模拟的 LLM：实现 answer() 和 summary()，统计调用次数和"模拟耗时"
    similarity / row_score 模拟的稠密检索：问题和一句评论有多像（认识一张小同义词表）
    parse_suql / to_dnf    把一条 SUQL 拆成 SELECT / FROM / WHERE / ORDER BY / LIMIT，把 WHERE 化成"或的与"
    run_sqlite             不做任何优化：把 answer、summary 注册成 SQLite 的函数，直接执行
    run_naive              幻灯片上的"朴素实现"：先对每一行调用 LLM，再做结构化过滤
    SUQLCompiler           做优化的执行器：谓词排序、只把最相似的行交给 LLM、惰性求值

-----------------------------------------------------------------------------
【模拟的边界】
-----------------------------------------------------------------------------
    · SimLLM 读评论时认得各种说法（"brought the kids" 和 "family-friendly" 是一回事），也知道一些
      常识（cacio e pepe 是一种 pasta），但会出错：文字提到了相关的话题时，有 error 的概率判断错
      （默认 3%）；文字根本没提这个话题时，出错的概率只有 error / 10；
    · "模拟耗时"按每次 answer 1 秒、每次 summary 2 秒计算，只用来比较不同执行策略的相对快慢；
    · 检索的"相似度"只是一个打分，不是判断：一条说"停车很难"的评论和问题"停车方便吗？"很相似。
这些都是设定的参数，说明的是机制，不代表真实 LLM 和检索模型的数字。
"""

import hashlib
import os
import random
import re
import sys
from collections import Counter

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_HERE, "..", "lecture06"))
from restaurant_world import FEATURES, LOCATIONS, PRICES, make_world  # noqa: E402

# 本讲新增的隐藏特征："停车方便"。第 6 讲的评论里有一句闲话 "parking can be tricky"（停车不方便），
# 这里再给一部分**没有**那句闲话的餐厅加上"停车方便"的说法，用来演示"相似 ≠ 满足条件"。
PARKING_EASY = ["plenty of free parking out front", "easy to park in their lot", "parking was a breeze"]
PARKING_HARD = ["parking can be tricky"]

# answer() 能理解的"概念"：问题里出现这些词，就知道在问哪个特征；评论里出现这些说法，就知道答案是"有"。
CONCEPT_WORDS = {
    "family_friendly": ["family", "kid", "kids", "child", "children"],
    "romantic": ["romantic", "date"],
    "view": ["view"],
    "spicy": ["spicy"],
    "groups": ["group", "groups", "party", "parties"],
    "outdoor": ["outdoor", "patio", "outside"],
    "parking": ["parking", "park"],
}
CONCEPT_PHRASES = dict(FEATURES)
CONCEPT_PHRASES["parking"] = PARKING_EASY

# LLM 的"常识"：一道菜属于哪一类（SQL 的字符串匹配做不到这一点）
DISH_KNOWLEDGE = {
    "pasta": ["pasta", "cacio e pepe"],
    "burgers": ["burger", "cheeseburger"],
    "noodles": ["noodle", "ramen", "pho", "pad see ew"],
    "sushi": ["sushi", "nigiri", "omakase"],
    "tacos": ["taco"],
    "dumplings": ["dumpling", "wonton", "har gow", "siu mai"],
    "curry": ["curry", "butter chicken"],
    "pizza": ["pizza"],
    "coffee": ["latte", "cappuccino", "coffee"],
}

# 模拟检索用的小同义词表：把不同的词映射到同一个"语义标签"（稠密检索的效果）
SYNONYMS = {"kids": "child", "kid": "child", "children": "child", "family": "child", "family-friendly": "child",
            "date": "romantic", "romantic": "romantic", "candlelit": "romantic", "intimate": "romantic",
            "patio": "outdoor", "outdoor": "outdoor", "sidewalk": "outdoor", "outside": "outdoor",
            "parking": "park", "park": "park", "lot": "park",
            "spicy": "spicy", "chili": "spicy", "hot": "spicy", "numbing": "spicy",
            "groups": "group", "group": "group", "parties": "group", "party": "group", "people": "group",
            "view": "view", "bay": "view", "gate": "view"}
STOP = {"is", "it", "this", "a", "an", "the", "for", "to", "do", "does", "you", "find", "restaurant", "place",
        "good", "great", "be", "in", "of", "and", "there", "they", "their", "we", "our", "very", "with", "can",
        "was", "what", "are", "have", "has", "serve", "serves"}


# =============================================================================
# 第 1 部分：数据
# =============================================================================
def make_suql_world(seed=0, p_parking=0.15):
    """
    第 6 讲的餐厅数据库，加上 SUQL 要用的一列 reviews，以及"停车方便"这个隐藏特征。

    【思路】
        · 第 6 讲把评论放在另一张表里（一条评论一行）。SUQL 的幻灯片把评论当作餐厅的一列（reviews TEXT[]），
          这里照做：给 restaurants 表加一列 reviews，内容是这家餐厅全部评论拼成的一段文字。
        · 没有 "parking can be tricky" 这句闲话的餐厅，以 p_parking 的概率在 reviews 末尾加一句"停车方便"。
    【参数】seed：随机种子（同一个 seed 总是生成同一个世界）；p_parking：加"停车方便"的概率
    【返回】第 6 讲的 World 对象，多了两样东西：
        world.truth[rid]：这家餐厅真实具备的特征（集合），包括 "parking"
        world.text[rid]  ：reviews 列的内容
    【边界】只改这个进程里的内存数据库，不影响第 6 讲的代码和输出。
    """
    world = make_world(seed=seed)
    rng = random.Random(f"suql-parking|{seed}")
    world.db.execute("ALTER TABLE restaurants ADD COLUMN reviews TEXT")
    world.truth, world.text = {}, {}
    for r in world.restaurants:
        text = world.review_text(r.id)
        truth = set(r.features)
        if "parking can be tricky" not in text.lower() and rng.random() < p_parking:
            phrase = rng.choice(PARKING_EASY)
            text += " " + phrase[0].upper() + phrase[1:] + "."
            truth.add("parking")
        world.truth[r.id], world.text[r.id] = truth, text
        world.db.execute("UPDATE restaurants SET reviews = ? WHERE id = ?", (text, r.id))
    return world


# =============================================================================
# 第 2 部分：模拟的 LLM —— answer() 和 summary()
# =============================================================================
def _stable_random(*keys):
    """
    由若干个键算出一个 [0, 1) 之间的"随机数"：同样的键永远得到同样的数。

    【为什么】模拟 LLM 对"同一行、同一个问题"的判断要可复现：不管执行器以什么顺序、调用几次，
        同一个 (文字, 问题) 的判断都一样 —— 这样比较不同执行策略时，差别只来自策略本身。
    【Python 小知识】hashlib.md5(...).hexdigest() 把任意字符串变成固定的 32 位十六进制串；
        取前 8 位转成整数再除以 16^8，就得到一个 [0, 1) 之间的数。Python 自带的 hash() 每次运行
        结果都可能不同（PYTHONHASHSEED），所以不能用。
    """
    h = hashlib.md5("|".join(str(k) for k in keys).encode("utf-8")).hexdigest()
    return int(h[:8], 16) / 16 ** 8


def question_concept(question):
    """
    这个问题在问哪个特征？例如 'is it family-friendly?' → 'family_friendly'。

    【返回】CONCEPT_WORDS 的一个键；认不出来时返回 None
    【边界】按问题里的词匹配，第一个匹配到的概念获胜（"parking" 要排在 "park" 相关的判断之前——
        这里两个词都映射到同一个概念，所以没有冲突）。
    """
    words = re.findall(r"[a-z\-]+", question.lower())
    for concept, keys in CONCEPT_WORDS.items():
        for w in words:
            if w in keys or w.split("-")[0] in keys:
                return concept
    return None


class SimLLM:
    """
    模拟的 LLM，提供 SUQL 的两个自由文本函数。

    【用法】
        llm = SimLLM(error=0.03)
        llm.answer(reviews_text, 'is it family-friendly?')   → 'Yes' / 'No'
        llm.summary(reviews_text)                            → 一段摘要
        llm.calls        → Counter({'answer': 12, 'summary': 1})
        llm.seconds()    → 模拟耗时（秒）

    【参数】
        error   : 每次判断出错（Yes 和 No 颠倒、抽错了值）的概率
        seed    : 出错的"随机性"由 (seed, 文字, 问题) 决定，可复现
        extract : 可选的"抽取规则"列表，每项是 (问题里的关键词, 正则表达式)。
                  用在第 6 节 HybridQA 那样"从一段介绍里抽出一个日期 / 地点"的问题上。
    """
    SECONDS = {"answer": 1.0, "summary": 2.0}

    def __init__(self, error=0.03, seed=0, extract=None):
        self.error, self.seed = error, seed
        self.extract = extract or []
        self.calls = Counter()

    def reset(self):
        """把调用计数清零（比较不同执行策略之前调用）。"""
        self.calls = Counter()

    def seconds(self):
        """按每种调用的模拟耗时，算出总的"模拟秒数"。"""
        return sum(self.SECONDS.get(k, 1.0) * n for k, n in self.calls.items())

    # ---- answer ----
    def answer(self, text, question):
        """
        SUQL 的 answer(文本, 问题)：读这段文字，回答问题。

        【思路】按问题的种类分三种情况：
            1. "抽取"问题（第 6 节的 HybridQA）：用 self.extract 里的规则从文字里找答案；
            2. 菜品问题 'does this restaurant serve pasta?'：用 DISH_KNOWLEDGE 的常识判断；
            3. 特征问题 'is it family-friendly?'：看文字里有没有这个特征的任何一种说法。
        【参数】text：一行里某个文本列的值（可能是 None）；question：问题
        【返回】字符串：'Yes' / 'No'，或者抽取出的值；抽不出来时返回 'Unknown'
        【边界】出错的方式（不对称，模仿真实 LLM 的常见错误）：
            · 文字里**确实有**这个特征时，以 self.error 的概率漏判（没认出换了说法的表达）→ 'No'；
            · 文字**提到了这个话题**但并不满足时（例如 "parking can be tricky"），以 self.error 的概率误判 → 'Yes'；
            · 文字**根本没提这个话题**时，误判的概率只有 self.error / 10 —— LLM 很少凭空认定文字里说了什么；
            · 抽取题以 self.error 的概率返回另一个候选值（例如段落里的另一个日期）；
            · text 为 None（这一行这一列没有内容）时返回 'Unknown'。
        """
        self.calls["answer"] += 1
        if text is None:
            return "Unknown"
        r = _stable_random("answer", self.seed, text, question)
        q, low = question.lower(), text.lower()
        for key, pattern in self.extract:                          # 情况 1：抽取
            if key in q:
                found = re.findall(pattern, text)
                if not found:
                    return "Unknown"
                if r < self.error and len(found) > 1:
                    return found[-1]                                # 抽错：拿了另一个值
                return found[0]
        m = re.search(r"serves? ([a-z ]+?)\??$", q)                 # 情况 2：菜品
        if m and m.group(1).strip() in DISH_KNOWLEDGE:
            dish = m.group(1).strip()
            truth = any(d in low for d in DISH_KNOWLEDGE[dish])
            topic = truth or dish.rstrip("s") in low                # 话题是否被提到
        else:                                                       # 情况 3：特征
            concept = question_concept(question)
            if concept is None:
                return "Unknown"
            truth = any(p.lower() in low for p in CONCEPT_PHRASES[concept])
            topic = truth or any(re.search(rf"\b{w}", low) for w in CONCEPT_WORDS[concept])
        rate = self.error if topic else self.error / 10
        if r < rate:
            truth = not truth
        return "Yes" if truth else "No"

    # ---- summary ----
    def summary(self, text):
        """
        SUQL 的 summary(文本)：给一段评论写个摘要。

        【返回】一段很短的英文摘要，列出评论里提到的特征（模拟 LLM 的"概括"）
        【边界】只是演示用的字符串；不会出错（摘要的质量这里不评估）。
        """
        self.calls["summary"] += 1
        if not text:
            return ""
        found = [c.replace("_", "-") for c in CONCEPT_PHRASES
                 if any(p.lower() in text.lower() for p in CONCEPT_PHRASES[c])]
        return "Reviewers mention: " + (", ".join(found) if found else "nothing special") + "."


# =============================================================================
# 第 3 部分：模拟的检索 —— 问题和一句评论有多"像"
# =============================================================================
def _tags(text):
    """把一段文字变成"语义标签"的集合：去掉停用词，同义词映射到同一个标签。"""
    words = re.findall(r"[a-z\-]+", text.lower())
    return {SYNONYMS.get(w, w) for w in words if w not in STOP}


def similarity(question, sentence):
    """
    问题和一句话的相似度：问题的语义标签里，有多大比例出现在这句话里。

    【返回】0~1 之间的数
    【边界】这只是"像不像"，不是"是不是"：'is parking easy?' 和 'Parking can be tricky.' 都含有标签
        'park'，相似度很高，但后者的意思恰恰相反 —— 所以检索之后必须让 LLM 判断（演示 2）。
    """
    q = _tags(question)
    if not q:
        return 0.0
    return len(q & _tags(sentence)) / len(q)


def row_score(questions, text):
    """
    一行的"检索得分"（SUQL 论文第 5.1 节的公式）：对每个文本条件 c，取这一行里和它最像的那句话的相似度，
    再把各个条件加起来：  score(r) = Σ_c  max_t sim(c, t)

    【参数】questions：这一行要满足的全部 answer() 条件里的问题；text：这一行的文字
    【为什么取 max 再求和】一家餐厅有很多句评论，只要有一句提到"适合带孩子"就够了（max）；
        多个条件时，同时和几个条件都像的行排在前面（求和）。
    """
    sentences = [s for s in re.split(r"(?<=[.!?])\s+", text or "") if s]
    return sum(max((similarity(q, s) for s in sentences), default=0.0) for q in questions)


# =============================================================================
# 第 4 部分：解析 SUQL —— 拆出各个子句，把 WHERE 化成"或的与"（DNF）
# =============================================================================
_ANSWER_RE = re.compile(r"answer\(\s*(\w+)\s*,\s*'((?:[^']|'')*)'\s*\)", re.I)


def split_top(text, keyword):
    """
    在"最外层"按关键字（AND / OR）切开一段条件：括号里面和引号里面的不切。

    【例子】split_top("a = 1 AND (b = 2 OR c = 3)", "OR") → ["a = 1 AND (b = 2 OR c = 3)"]（OR 在括号里）
            split_top("a = 1 AND (b = 2 OR c = 3)", "AND") → ["a = 1", "(b = 2 OR c = 3)"]
    【思路】从左到右扫一遍，记住当前在几层括号里（depth）、是不是在单引号字符串里（in_quote）；
        只有 depth == 0 且不在引号里时，遇到 " AND " 才切。
    【边界】关键字不区分大小写；字符串里的 '' 表示一个单引号（SQL 的转义写法），这里会把它当作
        "出引号再进引号"，效果一样。
    """
    parts, depth, in_quote, start, i = [], 0, False, 0, 0
    pad = f" {keyword} "
    upper = text.upper()
    while i < len(text):
        ch = text[i]
        if ch == "'":
            in_quote = not in_quote
        elif not in_quote and ch == "(":
            depth += 1
        elif not in_quote and ch == ")":
            depth -= 1
        elif not in_quote and depth == 0 and upper.startswith(pad, i):
            parts.append(text[start:i].strip())
            i += len(pad)
            start = i
            continue
        i += 1
    parts.append(text[start:].strip())
    return [p for p in parts if p]


def _strip_parens(text):
    """如果整段条件被一对括号包住（例如 "(a OR b)"），去掉这对括号。"""
    text = text.strip()
    while text.startswith("(") and text.endswith(")"):
        depth = 0
        for i, ch in enumerate(text):
            depth += ch == "("
            depth -= ch == ")"
            if depth == 0 and i < len(text) - 1:
                return text                                        # 第一个括号在中途就闭合了，不是整体包住
        text = text[1:-1].strip()
    return text


def to_dnf(condition):
    """
    把 WHERE 条件化成**析取范式**（DNF，"或的与"）：[[a, b], [c]] 表示 (a AND b) OR c。

    【为什么】SUQL 编译器要在每个"与"子句里决定先算哪个条件（谓词排序）。化成 DNF 之后，
        每个子句就是一串必须同时满足的条件，可以自由地重新排序（论文第 5.3 节）。
    【思路】递归：
        · 按最外层的 OR 切开 → 各部分的 DNF 直接拼起来；
        · 只有 AND：每个部分先化成 DNF，再做"分配律"展开：(a OR b) AND c = (a AND c) OR (b AND c)；
        · 不能再切的就是一个原子条件（例如 location = 'Palo Alto'）。
    【返回】list[list[str]]
    【边界】不处理 NOT（这里的查询用不到）；条件很多、OR 嵌套很深时，展开后子句数会指数增长。
    """
    condition = _strip_parens(condition)
    ors = split_top(condition, "OR")
    if len(ors) > 1:
        return [clause for part in ors for clause in to_dnf(part)]
    ands = split_top(condition, "AND")
    if len(ands) == 1:
        inner = _strip_parens(ands[0])
        if inner != ands[0]:                                       # 原子外面包了括号，里面可能还有 OR
            return to_dnf(inner)
        return [[ands[0]]]
    result = [[]]
    for part in ands:                                              # 分配律
        result = [left + right for left in result for right in to_dnf(part)]
    return result


def parse_suql(suql):
    """
    把一条 SUQL 拆成几个部分。

    【返回】dict：{"select": str, "table": str, "where": str 或 None, "order_by": str 或 None, "limit": int 或 None}
    【边界】只支持本讲用到的形式：SELECT … FROM 一张表 [WHERE …] [ORDER BY …] [LIMIT n]；
        不支持 JOIN、GROUP BY 和子查询（这些情况用 run_sqlite 直接执行）。
    """
    s = suql.strip().rstrip(";")
    m = re.match(r"(?is)^SELECT\s+(.*?)\s+FROM\s+(\w+)\s*(.*)$", s)
    if not m:
        raise ValueError(f"无法解析：{suql}")
    select, table, rest = m.group(1), m.group(2), m.group(3)
    limit = order_by = where = None
    lm = re.search(r"(?is)\s*LIMIT\s+(\d+)\s*$", rest)
    if lm:
        limit, rest = int(lm.group(1)), rest[:lm.start()]
    om = re.search(r"(?is)\s*ORDER\s+BY\s+(.*)$", rest)
    if om:
        order_by, rest = om.group(1).strip(), rest[:om.start()]
    wm = re.match(r"(?is)^\s*WHERE\s+(.*)$", rest)
    if wm:
        where = wm.group(1).strip()
    return {"select": select, "table": table, "where": where, "order_by": order_by, "limit": limit}


def is_free_text(atom):
    """这个原子条件里有没有 answer()？有的话要调用 LLM（贵），没有的话 SQL 自己就能算（便宜）。"""
    return "answer(" in atom.lower()


def questions_in(atoms):
    """一组原子条件里，全部 answer() 的问题（用来算检索得分）。"""
    return [m.group(2).replace("''", "'") for a in atoms for m in _ANSWER_RE.finditer(a)]


def register(world, llm):
    """把 answer、summary 注册成 SQLite 函数（之后 SQL 里就能直接写 answer(reviews, '...')）。"""
    world.db.create_function("answer", 2, llm.answer)
    world.db.create_function("summary", 1, llm.summary)


# =============================================================================
# 第 5 部分：三种执行方式
# =============================================================================
def run_sqlite(world, llm, suql):
    """
    【方式 A】不做任何优化：SQLite 原样执行（answer、summary 只是两个外部函数）。

    【为什么能这样做】SUQL 是 SQL 的扩展，只多了两个函数 —— SQL 引擎不用改就能执行（幻灯片："SQL can
        run SUQL programs without modification"）。
    【代价】SQLite 不知道 answer() 很贵：它按条件的**书写顺序**逐行求值。如果 answer() 写在前面，
        每一行都会先调用一次 LLM。
    【返回】结果行（list[tuple]）
    """
    register(world, llm)
    return world.query(suql)


def run_naive(world, llm, suql):
    """
    【方式 B】幻灯片上的"朴素实现"：
        1. 取出每一行的文字，对每一行调用 LLM 回答 answer() 的问题；
        2. 再按结构化条件（地点等）过滤；
        3. 对全部通过的行生成摘要；
        4. 最后才取前几个（LIMIT）。
    【返回】结果行（list[tuple]）
    【边界】只支持 parse_suql 能解析的形式。WHERE 里的每个"与"子句都按这个顺序执行。
    """
    register(world, llm)
    p = parse_suql(suql)
    ids = [row[0] for row in world.query(f"SELECT id FROM {p['table']}")]
    keep = set()
    for clause in to_dnf(p["where"]) if p["where"] else [[]]:
        free = [a for a in clause if is_free_text(a)]
        struct = [a for a in clause if not is_free_text(a)]
        passed = [i for i in ids if all(_holds(world, p["table"], i, a) for a in free)]          # 先问 LLM
        passed = [i for i in passed if all(_holds(world, p["table"], i, a) for a in struct)]     # 再过滤
        keep.update(passed)
    return _project(world, p, [i for i in ids if i in keep], apply_limit_after=True)


def _holds(world, table, rid, atom):
    """一行是否满足一个原子条件：让 SQLite 只对这一行求值这个条件（answer() 会被调用一次）。"""
    return bool(world.query(f"SELECT 1 FROM {table} WHERE id = ? AND ({atom})", (rid,)))


def _project(world, p, ids, apply_limit_after=False):
    """
    对选中的行做最后的"投影"：计算 SELECT 列表（可能含 summary），按 ORDER BY 排序，取 LIMIT。

    【参数】ids：选中的行（按候选顺序）；apply_limit_after：朴素实现里，先对**全部**选中的行计算
        SELECT 列表（包括 summary），最后才截断 —— 白白生成了很多摘要。
    【Python 小知识】SQL 里的 "id IN (1, 2, 3)" 需要把列表拼成字符串；这里的 id 都是整数，拼接是安全的。
        如果是用户输入的文字，一定要用参数（?）而不是拼接，防止 SQL 注入（第 3 讲）。
    """
    if not ids:
        return []
    pos = " ".join(f"WHEN {i} THEN {k}" for k, i in enumerate(ids))
    order = p["order_by"] or f"CASE id {pos} END"
    sql = f"SELECT {p['select']} FROM {p['table']} WHERE id IN ({','.join(map(str, ids))}) ORDER BY {order}"
    if p["limit"] is not None and not apply_limit_after:
        sql += f" LIMIT {p['limit']}"
    rows = world.query(sql)
    return rows[:p["limit"]] if p["limit"] is not None else rows


class SUQLCompiler:
    """
    【方式 C】做优化的 SUQL 执行器（论文第 5 节的三个优化）。

    【用法】
        compiler = SUQLCompiler(world, llm)                       # 三个优化全开
        rows = compiler.execute("SELECT ... WHERE ... LIMIT 1")
        compiler.last_plan   → 一段文字，说明这次是怎么执行的

    【参数】
        order    : 谓词排序 —— 每个"与"子句里，先用 SQL 算结构化条件，只对剩下的行调用 LLM
        lazy     : 惰性求值 —— 有 LIMIT 时，凑够了就停
        retrieve : 检索排序 —— 没有 ORDER BY 时，按检索得分（row_score）从高到低把候选行交给 LLM，
                   最可能满足条件的先问
        max_candidates : 检索时最多把多少行交给 LLM（None 表示不限）。设上限会更快，但可能漏掉答案。

    【思路】
        1. 把 WHERE 化成"或的与"（DNF）；每个"与"子句分成结构化条件和 answer() 条件；
        2. （order）每个子句用 SQL 算出满足它全部结构化条件的候选行；一行只要是某个子句的候选，就是总的候选；
        3. 确定候选行的顺序：有 ORDER BY（只涉及结构化的列）时按它排；没有时，（retrieve）按检索得分排，
           最可能满足条件的先问 LLM；
        4. 按顺序检查每个候选行：依次看它是哪些子句的候选，对那个子句的 answer() 条件逐个求值（一个不满足
           就跳过后面的，短路）；有一个子句满足，这一行就选中（OR 也短路）；
        5. （lazy）有 LIMIT 时，凑够了就停。
       最后由 SQL 完成排序和投影（summary 只对最终输出的行计算）。
    【为什么候选行按 ORDER BY 的顺序检查，惰性求值仍然正确？】按"最终的顺序"检查，先选中的行一定排在
        后选中的行前面；凑够 LIMIT 行时，后面的行不可能挤进前 LIMIT 名。
    【边界】
        · 只支持 parse_suql 能解析的形式；ORDER BY 里含 answer() 时不能惰性求值（要先知道每一行的值才能排序），
          这时退回 run_sqlite；
        · 关掉 order 时，所有行都是候选，结构化条件和 answer() 条件按书写顺序逐个求值（和 SQLite 一样）；
        · 设了 max_candidates 时，检索得分排在后面的行不会被检查 —— 答案可能就在那里（召回率下降）。
    """

    def __init__(self, world, llm, order=True, lazy=True, retrieve=True, max_candidates=None):
        self.world, self.llm = world, llm
        self.order, self.lazy, self.retrieve = order, lazy, retrieve
        self.max_candidates = max_candidates
        self.last_plan = ""

    def execute(self, suql):
        """执行一条 SUQL，返回结果行（list[tuple]）。执行计划的说明存在 self.last_plan。"""
        register(self.world, self.llm)
        p = parse_suql(suql)
        if p["order_by"] and is_free_text(p["order_by"]):
            self.last_plan = "ORDER BY 里有 answer()：不能惰性求值，交给 SQLite 逐行计算"
            return run_sqlite(self.world, self.llm, suql)
        clauses = to_dnf(p["where"]) if p["where"] else [[]]
        plan = [f"WHERE 化成 {len(clauses)} 个“与”子句"]
        # 第 2 步：每个子句的候选行，以及这个子句里要逐行求值的条件
        work = []                                                  # [(候选 id 的集合, 要逐行求值的条件)]
        for clause in clauses:
            free = [a for a in clause if is_free_text(a)]
            struct = [a for a in clause if not is_free_text(a)]
            if self.order:
                ids = set(self._ids(p, struct))                    # 先做便宜的
                work.append((ids, free))
                plan.append(f"  子句 {len(work)}：结构化条件 → {len(ids)} 个候选，再问 {len(free)} 个 answer()")
            else:
                ids = set(self._ids(p, []))                        # 不排序：所有行都是候选
                work.append((ids, clause))                         # 按书写顺序逐个求值
                plan.append(f"  子句 {len(work)}：不排序，{len(ids)} 行逐行按书写顺序求值")
        # 第 3 步：候选行的顺序
        order = [i for i in self._ids(p, []) if any(i in ids for ids, _ in work)]
        free_all = [a for clause in clauses for a in clause if is_free_text(a)]
        if self.retrieve and free_all and not p["order_by"]:
            score = {i: max(row_score(questions_in(todo), self.world.text.get(i, ""))
                            for ids, todo in work if i in ids) for i in order}
            order = sorted(order, key=lambda i: -score[i])         # 最像的先问 LLM
            plan.append("  候选行按检索得分排序")
            if self.max_candidates is not None:
                order = order[:self.max_candidates]
                plan.append(f"  只检查得分最高的 {self.max_candidates} 行")
        # 第 4、5 步：逐行检查，凑够就停
        target = p["limit"] if self.lazy else None
        chosen = []
        for rid in order:
            if any(rid in ids and all(_holds(self.world, p["table"], rid, a) for a in todo)
                   for ids, todo in work):                         # any() 和 all() 都会短路
                chosen.append(rid)
                if target is not None and len(chosen) >= target:
                    plan.append(f"  惰性求值：凑够 {target} 行就停")
                    break
        self.last_plan = "\n".join(plan)
        return _project(self.world, p, chosen)

    def _ids(self, p, struct):
        """用 SQL 算出满足全部结构化条件 struct 的行（有 ORDER BY 时按它排好，否则按 id）。"""
        sql = f"SELECT id FROM {p['table']}"
        if struct:
            sql += " WHERE " + " AND ".join(f"({a})" for a in struct)
        sql += f" ORDER BY {p['order_by']}" if p["order_by"] else " ORDER BY id"
        return [row[0] for row in self.world.query(sql)]


def gold_ids(world, suql):
    """
    标准答案：一个**不会出错**的 LLM（error = 0）逐行执行这条查询的 WHERE，满足条件的全部行。

    【用途】评估执行结果：返回的行是不是都在这个集合里（精确率）？这个集合里的行找到了多少（召回率）？
    【返回】id 的集合
    【边界】调用之后，SQLite 里注册的 answer() 是这个"完美"的 LLM；之后执行别的查询前要重新注册
        （run_sqlite、run_naive、SUQLCompiler.execute 都会自动重新注册）。
    """
    p = parse_suql(suql)
    register(world, SimLLM(error=0.0))
    sql = f"SELECT id FROM {p['table']}" + (f" WHERE {p['where']}" if p["where"] else "")
    return {row[0] for row in world.query(sql)}


def _phrase(atom):
    """
    把一个条件读成一个英文短语（verbalize 用）。

    【例子】location = 'X' → "in X"；answer(reviews, 'is it romantic?') = 'Yes' → "where the reviews say it is romantic"
    【边界】被括号包住的 "A OR B" 会拆开，分别读出来，再用 "or" 连起来；认不出来的条件原样返回。
    """
    inner = _strip_parens(atom)
    ors = split_top(inner, "OR")
    if len(ors) > 1:
        return "(" + " or ".join(_phrase(a) for a in ors) + ")"
    rules = [(r"location = '(.+)'", "in {}"), (r"has\(cuisines, '(.+)'\)", "serving {} food"),
             (r"price = '(.+)'", "with {} prices"), (r"rating >= ([\d.]+)", "rated at least {}"),
             (r"opening_hours LIKE '%(\w+)'", "open until {}"),
             # 如实读出 SQL 的含义：文字之间的 > 是按字母顺序比较的
             (r"opening_hours > '(\w+)'", "whose opening hours come after '{}' in alphabetical order")]
    for pattern, template in rules:
        m = re.match(pattern, inner)
        if m:
            return template.format(m.group(1))
    m = _ANSWER_RE.search(inner)
    if not m:
        return inner
    col, q = m.group(1).replace("_", " "), m.group(2).rstrip("?")
    yes_no = re.match(r"(is|does) (?:it|this restaurant) (\w+)(.*)$", q)
    if yes_no and yes_no.group(1) == "is":                    # "is it X" → "it is X"
        return f"where the {col} say it is {yes_no.group(2)}{yes_no.group(3)}"
    if yes_no:                                                # "does it have X" → "it has X"
        verb = {"have": "has"}.get(yes_no.group(2), yes_no.group(2) + "s")
        return f"where the {col} say it {verb}{yes_no.group(3)}"
    return f"where the {col} answer '{q}' with yes"


def verbalize(suql):
    """
    把一条 SUQL "读"回成一句英文（智能体回复时先说"我查了什么"，第 5 节）。

    【思路】按最外层的 AND 切开，每个条件用 _phrase 读成一个短语，再拼起来。
        真实系统用 LLM 根据 SUQL 生成这句话；这里用规则，足够说明这个做法。
    【边界】只认识本讲用到的几种条件；认不出来的条件原样附在后面。
    """
    p = parse_suql(suql)
    parts = [_phrase(a) for a in (split_top(p["where"], "AND") if p["where"] else [])]
    order = ""
    if p["order_by"]:
        order = {"rating DESC": ", sorted by rating", "num_reviews DESC": ", sorted by popularity"}.get(p["order_by"], "")
    return "I searched for restaurants " + ", ".join(parts) + order + "."


if __name__ == "__main__":
    w = make_suql_world()
    llm = SimLLM()
    q = ("SELECT name, summary(reviews) FROM restaurants WHERE location = 'Palo Alto' "
         "AND answer(reviews, 'is it family-friendly?') = 'Yes' LIMIT 1")
    print(SUQLCompiler(w, llm).execute(q), dict(llm.calls))
    print(verbalize(q))
