"""pg_stat_statements deltas and EXPLAIN helpers, per phase-0a-test-suite-spec.md §7."""
from contextlib import contextmanager

import psycopg


class _Counter:
    """Populated with `.calls` (and `.total_time_ms`) once the `with` block exits."""

    def __init__(self):
        self.calls = None
        self.total_time_ms = None

    def __repr__(self):
        return f"Counter(calls={self.calls}, total_time_ms={self.total_time_ms})"

    def __str__(self):
        return repr(self)


@contextmanager
def statement_count(dsn):
    """Reset pg_stat_statements, run the caller's block, yield the call-count delta.

    pg_stat_statements is reset globally, so this is only meaningful when
    nothing else is hitting the database concurrently during the `with` block
    -- true for the single-request pathology assertions this is built for.
    """
    counter = _Counter()
    with psycopg.connect(dsn, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT pg_stat_statements_reset()")
        try:
            yield counter
        finally:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT COALESCE(SUM(calls), 0), COALESCE(SUM(total_exec_time), 0) "
                    "FROM pg_stat_statements"
                )
                calls, total_time_ms = cur.fetchone()
            counter.calls = calls
            counter.total_time_ms = total_time_ms


def explain(dsn, sql, params=None):
    """Run EXPLAIN (ANALYZE, FORMAT JSON) and return the plan dict."""
    with psycopg.connect(dsn, autocommit=True) as conn:
        with conn.cursor() as cur:
            cur.execute(f"EXPLAIN (ANALYZE, FORMAT JSON) {sql}", params or ())
            (plan,) = cur.fetchone()
            return plan[0]


def plan_node_types(plan: dict) -> set:
    """Flatten an EXPLAIN JSON plan (or its top-level {"Plan": ...} wrapper)
    into the set of "Node Type" values it contains, at any depth."""
    types: set = set()
    root = plan.get("Plan", plan) if isinstance(plan, dict) else plan

    def walk(node):
        if isinstance(node, dict):
            if "Node Type" in node:
                types.add(node["Node Type"])
            for key, value in node.items():
                if key == "Plans" and isinstance(value, list):
                    for child in value:
                        walk(child)
        elif isinstance(node, list):
            for item in node:
                walk(item)

    walk(root)
    return types