"""ChatViz analysis DSL: validate -> compile -> execute -> export.

A *plan* is a small JSON object describing one descriptive analysis over the
corpus. The LLM planner (planner.py) only ever produces a plan; SQL is
generated here from fixed templates with bound parameters, so every analysis
is deterministic, read-only, and correct by construction.

Plan shape (all keys optional except "metric"):

    {
      "filters": {
        "platform": "claude", "model": "gpt-4o", "language": "English",
        "topic": "coding", "role": "user", "fts": "homework help",
        "date_from": "2025-01-01", "date_to": "2025-06-30",
        "has_thinking": true, "has_code": false, "has_links": true,
        "complete": true
      },
      "metric": "conversation_count",
      "group_by": "platform",
      "sample": {"n": 5000, "seed": 42},
      "limit": 20,
      "sort": "value_desc",
      "output": "bar"
    }
"""

import base64
import json
import re
import sqlite3
import time

MAX_ROWS = 200
MAX_SECONDS = 10

# ------------------------------------------------------------------ catalog

# conversation-level metrics: aggregate over one row per conversation (alias t)
CONV_METRICS = {
    "conversation_count": ("COUNT(*)", "conversations"),
    "avg_messages":       ("AVG(t.n_messages)", "avg messages per conversation"),
    "avg_turns":          ("AVG(t.turns_count)", "avg turns per conversation"),
    "pct_with_thinking":  ("AVG(t.has_thinking) * 100.0",
                           "% of conversations with thinking traces"),
    "pct_with_code":      ("AVG(t.has_code) * 100.0",
                           "% of conversations with code blocks"),
    "pct_with_links":     ("AVG(t.has_links) * 100.0",
                           "% of conversations with links/citations"),
    "pct_complete":       ("AVG(t.is_complete) * 100.0",
                           "% of conversations ending with an assistant "
                           "reply"),
}
# message-level metrics: aggregate over one row per message (alias t)
MSG_METRICS = {
    "message_count":      ("COUNT(*)", "messages"),
    "avg_message_length": ("AVG(t.len)", "avg message length (characters)"),
}
SPECIAL_METRICS = {
    "avg_response_time": "avg assistant response time (seconds)",
}
ALL_METRICS = list(CONV_METRICS) + list(MSG_METRICS) + list(SPECIAL_METRICS)

TURN_BUCKET = """CASE
    WHEN c.n_messages <= 2  THEN '1-2'
    WHEN c.n_messages <= 5  THEN '3-5'
    WHEN c.n_messages <= 10 THEN '6-10'
    WHEN c.n_messages <= 20 THEN '11-20'
    ELSE '21+' END"""
TURN_BUCKET_ORDER = ("CASE grp WHEN '1-2' THEN 1 WHEN '3-5' THEN 2 "
                     "WHEN '6-10' THEN 3 WHEN '11-20' THEN 4 ELSE 5 END")

CONV_GROUPS = {
    "platform":    "c.platform",
    "model":       "c.model",
    "month":       "substr(cs.created_day, 1, 7)",
    "year":        "substr(cs.created_day, 1, 4)",
    "turn_bucket": TURN_BUCKET,
}
MSG_GROUPS = {
    "language": "m.language",
    "topic":    "m.topic",
    "role":     "m.role",
}
ALL_GROUPS = list(CONV_GROUPS) + list(MSG_GROUPS)

CONV_FILTERS = {
    "platform":     ("c.platform = ?", str),
    "model":        ("c.model = ?", str),
    "date_from":    ("cs.created_day >= ?", str),
    "date_to":      ("cs.created_day <= ?", str),
    "has_thinking": ("cs.has_thinking = ?", bool),
    "has_code":     ("cs.has_code = ?", bool),
    "has_links":    ("cs.has_links = ?", bool),
    "complete":     ("cs.is_complete = ?", bool),
}
MSG_FILTERS = {
    "language": ("language = ?", str),
    "topic":    ("topic = ?", str),
    "role":     ("role = ?", str),
}
ALL_FILTERS = list(CONV_FILTERS) + list(MSG_FILTERS) + ["fts"]

SORTS = ("value_desc", "value_asc", "key")
OUTPUTS = ("table", "bar", "timeline")
DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$")


def fts_quote(q):
    """Wrap each user term in double quotes so FTS5 operators can't break the
    query, while still allowing trailing * for prefix search."""
    terms = []
    for tok in q.split():
        prefix = tok.endswith("*")
        tok = tok.rstrip("*").replace('"', "")
        if tok:
            terms.append(f'"{tok}"' + ("*" if prefix else ""))
    return " ".join(terms)


# --------------------------------------------------------------- validation

def validate_plan(plan):
    """Normalize and check a plan. Returns (normalized_plan, errors, warnings).
    The plan is safe to compile iff errors == []."""
    errors, warnings = [], []
    if not isinstance(plan, dict):
        return {}, ["plan must be a JSON object"], []

    known_keys = {"filters", "metric", "group_by", "sample", "limit",
                  "sort", "output"}
    for k in plan:
        if k not in known_keys:
            errors.append(f"unknown key {k!r} (allowed: {sorted(known_keys)})")

    metric = plan.get("metric")
    if not isinstance(metric, str) or metric not in ALL_METRICS:
        errors.append(f"metric must be one of {ALL_METRICS}, got {metric!r}")
        metric = None   # avoid hashing junk in the dict lookups below

    group = plan.get("group_by")
    if group is not None and (not isinstance(group, str)
                              or group not in ALL_GROUPS):
        errors.append(f"group_by must be null or one of {ALL_GROUPS}, "
                      f"got {group!r}")
        group = None

    filters = plan.get("filters") or {}
    norm_filters = {}
    if not isinstance(filters, dict):
        errors.append("filters must be an object")
        filters = {}
    for k, v in filters.items():
        if v is None:
            continue
        if k == "fts":
            if not isinstance(v, str) or not v.strip():
                errors.append("filters.fts must be a non-empty string")
            elif not fts_quote(v).strip():
                errors.append("filters.fts contains no searchable characters")
            else:
                norm_filters["fts"] = v.strip()
        elif k in CONV_FILTERS or k in MSG_FILTERS:
            _, typ = (CONV_FILTERS.get(k) or MSG_FILTERS.get(k))
            if typ is bool:
                if not isinstance(v, bool):
                    errors.append(f"filters.{k} must be true or false")
                    continue
            elif not isinstance(v, str) or not v.strip():
                errors.append(f"filters.{k} must be a non-empty string")
                continue
            if k in ("date_from", "date_to") and not DATE_RE.match(v):
                errors.append(f"filters.{k} must be YYYY-MM-DD, got {v!r}")
                continue
            if k == "role" and v not in ("user", "llm"):
                warnings.append(f"filters.role is usually 'user' or 'llm', "
                                f"got {v!r}")
            norm_filters[k] = v.strip() if isinstance(v, str) else v
        else:
            errors.append(f"unknown filter {k!r} (allowed: {ALL_FILTERS})")

    # metric/group compatibility
    if metric in CONV_METRICS and group == "role" \
            and metric != "conversation_count":
        errors.append("group_by 'role' only works with message-level metrics "
                      "(message_count, avg_message_length) or "
                      "conversation_count")
    if metric == "avg_response_time":
        if group in MSG_GROUPS:
            errors.append("avg_response_time can only be grouped by "
                          f"{list(CONV_GROUPS)} or null")
        if "role" in norm_filters:
            warnings.append("filters.role is ignored for avg_response_time")
            norm_filters.pop("role")

    sample = plan.get("sample")
    if sample is not None:
        n = sample.get("n") if isinstance(sample, dict) else None
        seed = sample.get("seed", 0) if isinstance(sample, dict) else None
        # bool is an int subclass, so reject it explicitly; bound the seed so
        # it can never overflow SQLite's signed-64-bit integer column.
        if (not isinstance(sample, dict)
                or isinstance(n, bool) or not isinstance(n, int)
                or isinstance(seed, bool) or not isinstance(seed, int)
                or not 1 <= n <= 500_000
                or not 0 <= seed <= 2**31 - 1):
            errors.append('sample must look like {"n": 5000, "seed": 42} '
                          "with 1 <= n <= 500000 and 0 <= seed <= 2147483647")
        else:
            sample = {"n": n, "seed": seed}

    limit = plan.get("limit", 20)
    if isinstance(limit, bool) or not isinstance(limit, int) \
            or not 1 <= limit <= 100:
        errors.append("limit must be an integer between 1 and 100")
        limit = 20

    sort = plan.get("sort")
    if sort is None:
        sort = "key" if group in ("month", "year", "turn_bucket") \
            else "value_desc"
    elif sort not in SORTS:
        errors.append(f"sort must be one of {SORTS}, got {sort!r}")

    output = plan.get("output", "table")
    if output not in OUTPUTS:
        errors.append(f"output must be one of {OUTPUTS}, got {output!r}")
    if output == "timeline" and group not in ("month", "year"):
        warnings.append("output 'timeline' works best with group_by "
                        "'month' or 'year'")

    norm = {"metric": metric, "group_by": group, "limit": limit,
            "sort": sort, "output": output}
    if norm_filters:
        norm["filters"] = norm_filters
    if sample and not errors:
        norm["sample"] = sample
    return norm, errors, warnings


# -------------------------------------------------------------- compilation

def _conv_conditions(filters, sample):
    """Conversation-level WHERE conditions (aliases c, cs)."""
    where, params = [], []
    for key, (cond, _typ) in CONV_FILTERS.items():
        if key in filters:
            v = filters[key]
            where.append(cond)
            params.append(int(v) if isinstance(v, bool) else v)
    if sample:
        where.append(
            "c.id IN (SELECT id FROM conversations "
            "ORDER BY ((id + ?) * 2654435761) % 4294967296 LIMIT ?)")
        params += [sample["seed"], sample["n"]]
    return where, params


def _msg_conditions(filters, alias, direct):
    """Message-level conditions. If direct, apply to messages alias `alias`;
    otherwise wrap each in an EXISTS over the conversation."""
    where, params = [], []
    for key, (cond, _typ) in MSG_FILTERS.items():
        if key in filters:
            if direct:
                where.append(f"{alias}.{cond}")
            else:
                where.append(
                    f"EXISTS (SELECT 1 FROM messages mx WHERE "
                    f"mx.conversation_id = c.id AND mx.{cond})")
            params.append(filters[key])
    if "fts" in filters:
        if direct:
            where.append(f"{alias}.id IN (SELECT rowid FROM messages_fts "
                         "WHERE messages_fts MATCH ?)")
        else:
            where.append(
                "c.id IN (SELECT mf.conversation_id FROM messages_fts "
                "JOIN messages mf ON mf.id = messages_fts.rowid "
                "WHERE messages_fts MATCH ?)")
        params.append(fts_quote(filters["fts"]))
    return where, params


def _order_limit(plan, grouped):
    if not grouped:
        return ""
    if plan["sort"] == "key":
        order = TURN_BUCKET_ORDER if plan["group_by"] == "turn_bucket" \
            else "grp"
    elif plan["sort"] == "value_asc":
        order = "value ASC"
    else:
        order = "value DESC"
    return f"\nORDER BY {order}\nLIMIT {plan['limit']}"


def compile_plan(plan):
    """Compile a *validated* plan into (sql, params, meta)."""
    metric, group = plan["metric"], plan.get("group_by")
    filters = plan.get("filters", {})
    sample = plan.get("sample")
    grouped = group is not None

    join = ("JOIN conversations c ON c.id = m.conversation_id\n"
            "LEFT JOIN conv_signals cs ON cs.conversation_id = c.id")

    if metric == "avg_response_time":
        grp_sel = f"{CONV_GROUPS[group]} AS grp,\n           " if grouped else ""
        cw, cp = _conv_conditions(filters, sample)
        mw, mp = _msg_conditions(filters, "m", direct=False)
        where = " AND ".join(cw + mw) or "1=1"
        sql = f"""WITH base AS (
    SELECT {grp_sel}m.role, m.created_at,
           LAG(m.created_at) OVER (
             PARTITION BY m.conversation_id
             ORDER BY m.message_index) AS prev
    FROM messages m
    {join}
    WHERE {where}
), d AS (
    SELECT {'grp, ' if grouped else ''}(julianday(created_at) - julianday(prev)) * 86400 AS dt
    FROM base
    WHERE role = 'llm' AND created_at IS NOT NULL AND prev IS NOT NULL
)
SELECT {'grp, ' if grouped else ''}AVG(dt) AS value, COUNT(*) AS n
FROM d
WHERE dt > 0 AND dt < 3600{(' AND grp IS NOT NULL' if grouped else '')}"""
        if grouped:
            sql += "\nGROUP BY grp"
        sql += _order_limit(plan, grouped)
        params = cp + mp
        value_label = SPECIAL_METRICS[metric]

    elif metric in CONV_METRICS:
        expr, value_label = CONV_METRICS[metric]
        cw, cp = _conv_conditions(filters, sample)
        if group in MSG_GROUPS or (group == "role"):
            # one row per (group value, conversation)
            grp_expr = MSG_GROUPS[group]
            mw, mp = _msg_conditions(filters, "m", direct=True)
            where = " AND ".join(cw + mw + [f"{grp_expr} IS NOT NULL"])
            inner = f"""SELECT DISTINCT {grp_expr} AS grp, c.id AS cid,
        c.n_messages AS n_messages, c.turns_count AS turns_count,
        cs.has_thinking AS has_thinking, cs.has_code AS has_code,
        cs.has_links AS has_links, cs.is_complete AS is_complete
    FROM messages m
    {join}
    WHERE {where}"""
            params = cp + mp
        else:
            grp_sel = f"{CONV_GROUPS[group]} AS grp,\n        " if grouped else ""
            mw, mp = _msg_conditions(filters, "m", direct=False)
            where = " AND ".join(cw + mw) or "1=1"
            if grouped:
                where += f" AND {CONV_GROUPS[group]} IS NOT NULL"
            inner = f"""SELECT {grp_sel}c.id AS cid,
        c.n_messages AS n_messages, c.turns_count AS turns_count,
        cs.has_thinking AS has_thinking, cs.has_code AS has_code,
        cs.has_links AS has_links, cs.is_complete AS is_complete
    FROM conversations c
    LEFT JOIN conv_signals cs ON cs.conversation_id = c.id
    WHERE {where}"""
            params = cp + mp
        sel = f"{'grp, ' if grouped else ''}{expr} AS value, COUNT(*) AS n"
        sql = f"SELECT {sel}\nFROM (\n    {inner}\n) t"
        if grouped:
            sql += "\nGROUP BY grp"
        sql += _order_limit(plan, grouped)

    else:  # message-level metric
        expr, value_label = MSG_METRICS[metric]
        cw, cp = _conv_conditions(filters, sample)
        mw, mp = _msg_conditions(filters, "m", direct=True)
        if grouped:
            grp_expr = MSG_GROUPS.get(group) or CONV_GROUPS[group]
            grp_sel = f"{grp_expr} AS grp,\n        "
            extra = [f"{grp_expr} IS NOT NULL"]
        else:
            grp_sel, extra = "", []
        where = " AND ".join(cw + mw + extra) or "1=1"
        inner = f"""SELECT {grp_sel}LENGTH(m.plain_text) AS len
    FROM messages m
    {join}
    WHERE {where}"""
        sel = f"{'grp, ' if grouped else ''}{expr} AS value, COUNT(*) AS n"
        sql = f"SELECT {sel}\nFROM (\n    {inner}\n) t"
        if grouped:
            sql += "\nGROUP BY grp"
        sql += _order_limit(plan, grouped)
        params = cp + mp

    meta = {"value_label": value_label, "grouped": grouped,
            "group_label": group or "", "output": plan["output"]}
    return sql, params, meta


# ---------------------------------------------------------------- execution

def tune_connection(con):
    """Read-path pragmas for a corpus that never changes underneath us:
    memory-mapped I/O (random index probes skip syscalls), a bigger page
    cache, and in-memory temp stores for GROUP BY/ORDER BY spills."""
    for pragma in ("PRAGMA mmap_size=268435456",      # 256 MB window
                   "PRAGMA cache_size=-64000",        # 64 MB page cache
                   "PRAGMA temp_store=MEMORY"):
        con.execute(pragma)


def _connect_ro(db_path):
    con = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
    con.row_factory = sqlite3.Row
    tune_connection(con)
    return con

def run_plan(db_path, plan, max_seconds=MAX_SECONDS, max_rows=MAX_ROWS):
    """Execute a validated plan against db_path (read-only).
    Returns (rows, meta) where rows is a list of dicts."""
    sql, params, meta = compile_plan(plan)
    con = _connect_ro(db_path)
    deadline = time.time() + max_seconds
    con.set_progress_handler(
        lambda: 1 if time.time() > deadline else 0, 100_000)
    t0 = time.time()
    try:
        raw = con.execute(sql, params).fetchmany(max_rows)
    except sqlite3.OperationalError as e:
        if "interrupted" in str(e):
            raise TimeoutError(
                f"query exceeded {max_seconds}s on the full corpus; "
                "narrow the filters or export the script") from e
        raise
    finally:
        con.close()
    rows = []
    for r in raw:
        d = dict(r)
        if d.get("value") is not None:
            d["value"] = round(d["value"], 2)
        rows.append(d)
    meta["elapsed"] = round(time.time() - t0, 3)
    meta["sql"] = sql
    meta["params"] = params
    return rows, meta


def example_conversations(db_path, plan, k=5):
    """Representative conversations matching the plan's filters, so a number
    is never the only thing a researcher sees. Deterministic order (same hash
    as the sample op). When an fts filter is present the preview is a
    matching message; otherwise the first message."""
    filters = plan.get("filters", {})
    cw, cp = _conv_conditions(filters, plan.get("sample"))
    mw, mp = _msg_conditions(filters, "m", direct=False)
    where = " AND ".join(cw + mw) or "1=1"
    if "fts" in filters:
        preview = """(SELECT mm.plain_text FROM messages mm
                  WHERE mm.conversation_id = c.id AND mm.id IN
                        (SELECT rowid FROM messages_fts
                         WHERE messages_fts MATCH ?)
                  ORDER BY mm.message_index LIMIT 1)"""
        pre_params = [fts_quote(filters["fts"])]
    else:
        preview = """(SELECT mm.plain_text FROM messages mm
                  WHERE mm.conversation_id = c.id
                        AND mm.plain_text IS NOT NULL
                  ORDER BY mm.message_index LIMIT 1)"""
        pre_params = []
    sql = f"""
        SELECT c.id, c.platform, c.model, c.n_messages,
               {preview} AS preview
        FROM conversations c
        LEFT JOIN conv_signals cs ON cs.conversation_id = c.id
        WHERE {where}
        ORDER BY (c.id * 2654435761) % 4294967296
        LIMIT ?"""
    con = _connect_ro(db_path)
    deadline = time.time() + MAX_SECONDS
    con.set_progress_handler(
        lambda: 1 if time.time() > deadline else 0, 100_000)
    try:
        return [dict(r) for r in
                con.execute(sql, pre_params + cp + mp + [k]).fetchmany(k)]
    except sqlite3.OperationalError:
        return []   # never let the example panel break a successful analysis
    finally:
        con.close()


# ------------------------------------------------------------- explanation

def _fmt(v, pct=False):
    if v is None:
        return "–"
    s = "{:,}".format(int(v)) if float(v).is_integer() else "{:,.2f}".format(v)
    return (s + "%") if pct else s


def explain(rows, meta, plan):
    """A short deterministic summary of what the chart/table shows, so a
    reader never faces a bare number. Returns None when there is nothing
    to say (no rows)."""
    vals = [r for r in rows if r.get("value") is not None]
    if not vals:
        return None
    label = meta["value_label"]
    pct = label.startswith("%")
    filters = plan.get("filters") or {}
    scope = "; ".join(f"{k} = {str(v).lower() if isinstance(v, bool) else v}"
                      for k, v in filters.items())
    scope = f" Scope: {scope}." if scope else ""

    if not meta["grouped"]:
        r = vals[0]
        n = r.get("n")
        return (f"One overall figure: {label} is {_fmt(r['value'], pct)}"
                + (f", computed over {n:,} rows" if n else "") + f".{scope}")

    glabel = meta["group_label"]
    top = max(vals, key=lambda r: r["value"])
    low = min(vals, key=lambda r: r["value"])

    # timeline reading: rows arrive in key order for month/year (sort=key)
    if glabel in ("month", "year") and plan.get("sort") == "key" \
            and len(vals) > 1:
        first, last = vals[0], vals[-1]
        if first["value"]:
            change = (last["value"] - first["value"]) / first["value"] * 100
            verb = ("rose" if change > 5 else
                    "fell" if change < -5 else "stayed roughly flat")
        else:
            verb = "went"
        return (f"{label} over {len(vals)} {glabel}s: it {verb} from "
                f"{_fmt(first['value'], pct)} in {first['grp']} to "
                f"{_fmt(last['value'], pct)} in {last['grp']}, peaking at "
                f"{_fmt(top['value'], pct)} in {top['grp']}.{scope}")

    out = (f"Each bar is one {glabel}; the value is {label}. "
           f"{top['grp']} is highest at {_fmt(top['value'], pct)}")
    if low is not top:
        out += f"; {low['grp']} is lowest at {_fmt(low['value'], pct)}"
        if low["value"] and top["value"] / low["value"] >= 1.15:
            out += f" ({top['value'] / low['value']:.1f}× spread)"
    return out + f" across {len(vals)} groups.{scope}"


# ------------------------------------------------------------ shareability

def to_permalink(plan):
    raw = json.dumps(plan, separators=(",", ":"), sort_keys=True)
    return base64.urlsafe_b64encode(raw.encode()).decode()


def from_permalink(s):
    return json.loads(base64.urlsafe_b64decode(s.encode()))


# Every value the user can influence (PLAN, SQL, PARAMS) is embedded with
# repr(), which always yields a safe Python literal — untrusted text can never
# break out into executable source. The signal-table build is embedded too so
# the script is self-contained against a database made by build_db.py alone.
SCRIPT_TEMPLATE = '''#!/usr/bin/env python3
"""ChatViz exported analysis - runs the exact plan below on a full local copy
of the corpus, using the same SQL the live demo executed on its subset.

1. Build the database (see ChatViz README):  python build_db.py
2. Run:  python this_script.py [path/to/chats.db]

The script ensures the conv_signals helper table exists (creating it if your
database predates it), then runs the analysis. Output: a result table on
stdout and analysis_result.csv.
"""
import csv
import sqlite3
import sys

# --- analysis plan (fixed schema, validated before export) ---
PLAN = {plan!r}

# --- compiled SQL: read-only, parameterized; identical to the demo run ---
SQL = {sql!r}

PARAMS = {params!r}

# --- idempotent build/upgrade of the conv_signals helper table ---
ENSURE_SIGNALS = {build!r}
UPGRADE_SIGNALS = {upgrade!r}

db = sys.argv[1] if len(sys.argv) > 1 else "chats.db"
con = sqlite3.connect(db)
con.row_factory = sqlite3.Row
have = con.execute("SELECT 1 FROM sqlite_master WHERE type='table' "
                   "AND name='conv_signals'").fetchone()
if not have:
    for stmt in ENSURE_SIGNALS:
        con.execute(stmt)
    con.commit()
elif not any(r[1] == "is_complete"
             for r in con.execute("PRAGMA table_info(conv_signals)")):
    for stmt in UPGRADE_SIGNALS:    # table predates the completeness signal
        con.execute(stmt)
    con.commit()
rows = [dict(r) for r in con.execute(SQL, PARAMS).fetchall()]
con.close()

if not rows:
    print("no rows matched")
    sys.exit(0)

cols = list(rows[0].keys())
widths = [max(len(c), *(len(str(r[c])) for r in rows)) for c in cols]
print("  ".join(c.ljust(w) for c, w in zip(cols, widths)))
for r in rows:
    print("  ".join(str(r[c]).ljust(w) for c, w in zip(cols, widths)))

with open("analysis_result.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=cols)
    w.writeheader()
    w.writerows(rows)
print(f"\\n{{len(rows)}} row(s) -> analysis_result.csv")
'''


def to_script(plan):
    """Render a validated plan as a standalone, dependency-free script.

    All user-influenced values are embedded as repr() literals, so untrusted
    text can never escape into executable code."""
    import migrate
    sql, params, _meta = compile_plan(plan)
    return SCRIPT_TEMPLATE.format(
        plan=plan,
        sql=sql,
        params=params,
        build=migrate.BUILD_STATEMENTS,
        upgrade=migrate.COMPLETENESS_STATEMENTS,
    )