"""DuckDB-backed queries over aggregates.jsonl (fast path) or runs.jsonl (when grouping)."""

import json
import sys
from pathlib import Path
from typing import Any

import duckdb

from abench_speckz import aggregate, prettynames
from abench_speckz import render as _render_mod
from abench_speckz.filters import check_ident, parse_kv

# Re-exports kept for backwards compatibility with tests that import the
# private helpers from this module. New code should use ``render`` directly.
from abench_speckz.render import _fmt, _json_default  # noqa: F401

_AGGREGATE_COLUMNS = ("n", "n_failed", "mean", "stddev", "p50", "p95", "p99", "ci95_low", "ci95_high")


def var_names(results_dir: Path) -> list[str]:
    """Return variable names present in results, in their original order.

    Reads the first row of aggregates.jsonl (or runs.jsonl as fallback).
    Returns an empty list when no data exists yet.
    """
    for fname in ("aggregates.jsonl", "runs.jsonl"):
        path = results_dir / fname
        if path.exists() and path.stat().st_size > 0:
            with path.open() as f:
                for line in f:
                    row = json.loads(line)
                    if isinstance(row.get("vars"), dict):
                        return list(row["vars"].keys())
    return []


def compute_rows(
    results_dir: Path,
    *,
    group_by: list[str] | None = None,
    where: list[str] | None = None,
    require_tags: list[str] | None = None,
    exclude_tags: list[str] | None = None,
    metrics: list[str] | None = None,
    from_raw: bool = False,
) -> list[dict[str, Any]]:
    """Return aggregate rows for the results dir.

    Reads pre-aggregated stats from ``aggregates.jsonl`` by default, but
    routes through ``runs.jsonl`` whenever the caller groups across
    config_hashes — stddev and percentiles aren't additive across
    pre-aggregated rows, so they have to be recomputed from raw values.
    """
    group_by = group_by or []
    where = where or []
    require_tags = require_tags or []
    exclude_tags = exclude_tags or []
    metrics = metrics or []

    for g in group_by:
        check_ident(g, context="group-by key")

    must_recompute = bool(group_by)
    if from_raw or must_recompute:
        return _from_raw(results_dir, group_by, where, require_tags, exclude_tags, metrics)
    return _from_aggregates(results_dir, group_by, where, require_tags, exclude_tags, metrics)


def stats(
    results_dir: Path,
    *,
    group_by: list[str] | None = None,
    where: list[str] | None = None,
    require_tags: list[str] | None = None,
    exclude_tags: list[str] | None = None,
    metrics: list[str] | None = None,
    from_raw: bool = False,
    fmt: str = "table",
    pretty: bool = False,
    report: Path | None = None,
    plots: Path | None = None,
    run_reports_mode: str = "embedded",
    out=sys.stdout,
) -> None:
    group_by = group_by or []
    rows = compute_rows(
        results_dir,
        group_by=group_by,
        where=where,
        require_tags=require_tags,
        exclude_tags=exclude_tags,
        metrics=metrics,
        from_raw=from_raw,
    )

    if not rows:
        print("# 0 rows", file=out)
        return

    pretty_map = prettynames.load(results_dir / "pretty_names.json") if pretty else {}
    _render_mod.render(rows, fmt, out, pretty_map=pretty_map)

    if report is not None:
        from abench_speckz import plots as _plots
        from abench_speckz import report as _report

        pm = prettynames.load(results_dir / "pretty_names.json")
        if plots is not None:
            doc = _plots.load_plots_file(plots)
        else:
            doc = _plots.discover_from_tool_snapshots(results_dir)
        if doc.plots or doc.sections:
            _report.generate_plots_html(results_dir, doc, report, pm, run_reports_mode=run_reports_mode)
        else:
            _report.generate_html(
                rows, report, group_by, pm, results_dir=results_dir, run_reports_mode=run_reports_mode
            )
        print(f"report written to {report}", file=out)


def _from_aggregates(
    results_dir: Path,
    group_by: list[str],
    where_exprs: list[str],
    require_tags: list[str],
    exclude_tags: list[str],
    metrics: list[str],
) -> list[dict[str, Any]]:
    path = results_dir / "aggregates.jsonl"
    if not path.exists() or path.stat().st_size == 0:
        return []

    con = duckdb.connect(":memory:")
    src = f"read_json_auto('{path}', format='newline_delimited')"
    where_clause, params = _build_where(where_exprs, require_tags, exclude_tags, metrics)
    group_cols = [f"(vars->>'{g}')" for g in group_by]

    if group_by:
        select_keys = ", ".join(f"{c} AS {g}" for c, g in zip(group_cols, group_by, strict=True))
        agg_picks = ", ".join(f"any_value({c}) AS {c}" for c in _AGGREGATE_COLUMNS)
        sql = (
            f"SELECT {select_keys}, metric, {agg_picks} "
            f"FROM {src} "
            f"WHERE {where_clause} "
            f"GROUP BY {', '.join(group_cols)}, metric "
            f"ORDER BY {', '.join(group_cols)}, metric"
        )
    else:
        cols = ", ".join(_AGGREGATE_COLUMNS)
        sql = f"SELECT metric, {cols} FROM {src} WHERE {where_clause} ORDER BY metric"

    res = con.execute(sql, params).fetchall()
    headers = [d[0] for d in con.description]
    return [dict(zip(headers, r, strict=True)) for r in res]


def _from_raw(
    results_dir: Path,
    group_by: list[str],
    where_exprs: list[str],
    require_tags: list[str],
    exclude_tags: list[str],
    metrics: list[str],
) -> list[dict[str, Any]]:
    from collections import defaultdict

    path = results_dir / "runs.jsonl"
    if not path.exists() or path.stat().st_size == 0:
        return []

    con = duckdb.connect(":memory:")
    src = f"read_json_auto('{path}', format='newline_delimited')"

    where_clause, params = _build_where(
        where_exprs,
        require_tags,
        exclude_tags,
        metrics=[],
        extra_clauses=("exit_code = 0", "results IS NOT NULL", "NOT list_contains(tags, 'warmup')"),
    )
    rows_sql = f"SELECT * FROM {src} WHERE {where_clause}"
    raw_rows = con.execute(rows_sql, params).fetchall()
    headers = [d[0] for d in con.description]
    runs = [dict(zip(headers, r, strict=True)) for r in raw_rows]

    n_failed_by_key: dict[tuple, int] = defaultdict(int)
    failed_sql = f"SELECT vars, tags FROM {src} WHERE exit_code <> 0 AND NOT list_contains(tags, 'warmup')"
    for r in con.execute(failed_sql).fetchall():
        d = dict(zip([h[0] for h in con.description], r, strict=True))
        key = tuple((d.get("vars") or {}).get(g, "") for g in group_by)
        n_failed_by_key[key] += 1

    by_group: dict[tuple, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list))
    by_group_meta: dict[tuple, dict[str, str]] = {}
    for run in runs:
        vars_ = run.get("vars") or {}
        key = tuple(vars_.get(g, "") for g in group_by)
        if key not in by_group_meta:
            by_group_meta[key] = {g: vars_.get(g, "") for g in group_by}
        for metric, value in (run.get("results") or {}).items():
            if metrics and metric not in metrics:
                continue
            if isinstance(value, (int, float)) and not isinstance(value, bool):
                by_group[key][metric].append(float(value))

    output_rows: list[dict[str, Any]] = []
    for key, metric_values in by_group.items():
        meta = by_group_meta[key]
        for metric, values in sorted(metric_values.items()):
            if not values:
                continue
            s = aggregate.stats_for(values)
            row: dict[str, Any] = dict(meta)
            row["metric"] = metric
            row["n"] = s["n"]
            row["n_failed"] = n_failed_by_key.get(key, 0)
            for c in ("mean", "stddev", "p50", "p95", "p99", "ci95_low", "ci95_high"):
                row[c] = s[c]
            output_rows.append(row)

    output_rows.sort(key=lambda r: tuple(r.get(g, "") for g in group_by) + (r["metric"],))
    return output_rows


def _build_where(
    where_exprs: list[str],
    require_tags: list[str],
    exclude_tags: list[str],
    metrics: list[str],
    *,
    extra_clauses: tuple[str, ...] = (),
) -> tuple[str, list[Any]]:
    """Compose a DuckDB WHERE clause from CLI-style filter args.

    ``extra_clauses`` are AND-joined verbatim before the dynamic predicates —
    used by callers that need fixed-shape filters (e.g. ``exit_code = 0``).
    """
    clauses: list[str] = list(extra_clauses) or ["1=1"]
    params: list[Any] = []
    for pred in parse_kv(where_exprs):
        clauses.append(f"(vars->>'{pred.key}') = ?")
        params.append(pred.value)
    for t in require_tags:
        clauses.append("list_contains(tags, ?)")
        params.append(t)
    for t in exclude_tags:
        clauses.append("NOT list_contains(tags, ?)")
        params.append(t)
    if metrics:
        placeholders = ", ".join("?" for _ in metrics)
        clauses.append(f"metric IN ({placeholders})")
        params.extend(metrics)
    return " AND ".join(clauses), params
