"""评论树：递归 CTE

一次 SQL 查完一篇文章的所有评论 + 层级 + 路径，无需 N+1。

运行：python code/comments_tree.py [post_id]
"""
from __future__ import annotations

import sys

from db import fetch_all, close_pool


TREE_SQL = """
WITH RECURSIVE tree AS (
    SELECT c.id,
           c.post_id,
           c.parent_id,
           c.author_id,
           c.content,
           c.created_at,
           1                                     AS depth,
           ARRAY[c.id]                           AS path
    FROM   comments c
    WHERE  c.post_id = %(p)s
      AND  c.parent_id IS NULL

    UNION ALL

    SELECT c.id, c.post_id, c.parent_id, c.author_id,
           c.content, c.created_at,
           t.depth + 1,
           t.path || c.id
    FROM   comments c
    JOIN   tree t ON c.parent_id = t.id
)
SELECT id, parent_id, author_id, depth, path, content
FROM   tree
ORDER  BY path;                                  -- 深度优先、父在前
"""


def build_tree(post_id: int) -> list[dict]:
    return fetch_all(TREE_SQL, {"p": post_id})


def pretty_print(post_id: int) -> None:
    rows = build_tree(post_id)
    if not rows:
        print(f"文章 {post_id} 没有评论")
        return

    print(f"=== 文章 {post_id} 的评论树（{len(rows)} 条） ===")
    for r in rows:
        indent = "    " * (r["depth"] - 1)
        print(f"{indent}└─ [#{r['id']} by u{r['author_id']}] "
              f"{r['content'][:50]}")


def explain(post_id: int) -> list[dict]:
    sql = "EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) " + TREE_SQL
    return fetch_all(sql, {"p": post_id})


if __name__ == "__main__":
    pid = int(sys.argv[1]) if len(sys.argv) > 1 else 1
    try:
        pretty_print(pid)
        print("\n-- EXPLAIN --")
        for r in explain(pid):
            print(" ", r["QUERY PLAN"])
    finally:
        close_pool()