#!/usr/bin/env python3
"""Quick utility to query the production postgres database.

Usage:
    python scripts/query_db.py stats
    python scripts/query_db.py sessions [--source codex_cli] [--limit 10]
    python scripts/query_db.py messages --session <session_id> [--limit 20]
    python scripts/query_db.py sources
    python scripts/query_db.py sql "SELECT count(*) FROM messages"
"""

import argparse
import json
import os
import sys

import psycopg


DSN = os.environ.get("QC_TRACE_DSN", "")


def get_conn():
    if not DSN:
        print("Error: QC_TRACE_DSN env var not set")
        print("  export QC_TRACE_DSN='postgresql://...'")
        sys.exit(1)
    return psycopg.connect(DSN)


def cmd_stats(args):
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT count(*) FROM messages")
            total_msgs = cur.fetchone()[0]

            cur.execute("SELECT count(*) FROM sessions")
            total_sessions = cur.fetchone()[0]

            cur.execute(
                "SELECT source, count(*) as cnt FROM messages "
                "GROUP BY source ORDER BY cnt DESC"
            )
            by_source = cur.fetchall()

            cur.execute(
                "SELECT source, count(*) as cnt FROM sessions "
                "GROUP BY source ORDER BY cnt DESC"
            )
            sessions_by_source = cur.fetchall()

    print(f"Total messages:  {total_msgs:,}")
    print(f"Total sessions:  {total_sessions:,}")
    print()
    print("Messages by source:")
    for source, count in by_source:
        print(f"  {source:<20} {count:>10,}")
    print()
    print("Sessions by source:")
    for source, count in sessions_by_source:
        print(f"  {source:<20} {count:>10,}")


def cmd_sessions(args):
    with get_conn() as conn:
        with conn.cursor() as cur:
            query = "SELECT id, source, org, device_name, first_seen, last_updated FROM sessions"
            params = []
            if args.source:
                query += " WHERE source = %s"
                params.append(args.source)
            query += " ORDER BY last_updated DESC LIMIT %s"
            params.append(args.limit)
            cur.execute(query, params)
            rows = cur.fetchall()

    for row in rows:
        sid, source, org, device, first_seen, last_updated = row
        print(f"{sid[:20]}  {source:<15} {org or '-':<15} {device or '-':<25} {last_updated}")


def cmd_messages(args):
    if not args.session:
        print("Error: --session required")
        sys.exit(1)
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT id, source, msg_type, timestamp, "
                "LEFT(content, 80) as content_preview "
                "FROM messages WHERE session_id = %s "
                "ORDER BY raw_line_number LIMIT %s",
                (args.session, args.limit),
            )
            rows = cur.fetchall()

    if not rows:
        print(f"No messages found for session {args.session}")
        return

    for row in rows:
        mid, source, msg_type, ts, preview = row
        preview = (preview or "").replace("\n", " ")
        print(f"  {msg_type:<15} {preview[:60]}")
    print(f"\n  {len(rows)} messages shown")


def cmd_sources(args):
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT source, count(*) as msg_count, "
                "count(DISTINCT session_id) as session_count, "
                "max(ingested_at) as last_ingested "
                "FROM messages GROUP BY source ORDER BY msg_count DESC"
            )
            rows = cur.fetchall()

    print(f"{'Source':<20} {'Messages':>10} {'Sessions':>10} {'Last ingested'}")
    print("-" * 70)
    for source, msg_count, session_count, last_ingested in rows:
        print(f"{source:<20} {msg_count:>10,} {session_count:>10,} {last_ingested}")


def cmd_sql(args):
    query = args.query
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(query)
            if cur.description:
                cols = [d.name for d in cur.description]
                rows = cur.fetchall()
                print("\t".join(cols))
                for row in rows:
                    print("\t".join(str(v) for v in row))
            else:
                print(f"OK ({cur.rowcount} rows affected)")


def main():
    parser = argparse.ArgumentParser(description="Query qc-trace postgres")
    sub = parser.add_subparsers(dest="command")

    sub.add_parser("stats")

    p_sessions = sub.add_parser("sessions")
    p_sessions.add_argument("--source", default=None)
    p_sessions.add_argument("--limit", type=int, default=10)

    p_messages = sub.add_parser("messages")
    p_messages.add_argument("--session", required=True)
    p_messages.add_argument("--limit", type=int, default=50)

    sub.add_parser("sources")

    p_sql = sub.add_parser("sql")
    p_sql.add_argument("query")

    args = parser.parse_args()
    if not args.command:
        parser.print_help()
        return

    cmds = {
        "stats": cmd_stats,
        "sessions": cmd_sessions,
        "messages": cmd_messages,
        "sources": cmd_sources,
        "sql": cmd_sql,
    }
    cmds[args.command](args)


if __name__ == "__main__":
    main()
