#!/usr/bin/env bash
set -euo pipefail

# Delete all data for an org from the qc-trace database.
# Usage: ./delete_org.sh <org_name>
#
# Requires QC_TRACE_DSN env var or .prod.env in the repo root.
# Uses uv + psycopg (no psql needed).

ORG="${1:-}"
if [ -z "$ORG" ]; then
    echo "Usage: $0 <org_name>"
    echo "Example: $0 sagar-remote"
    exit 1
fi

# Load DSN from .prod.env if not set
if [ -z "${QC_TRACE_DSN:-}" ]; then
    SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
    ENV_FILE="$SCRIPT_DIR/../.prod.env"
    if [ -f "$ENV_FILE" ]; then
        QC_TRACE_DSN=$(grep '^QC_TRACE_DSN=' "$ENV_FILE" | cut -d'"' -f2)
    fi
fi

if [ -z "${QC_TRACE_DSN:-}" ]; then
    echo "Error: QC_TRACE_DSN not set and .prod.env not found"
    exit 1
fi

export QC_TRACE_DSN
export ORG

uv run --with psycopg-binary python3 << 'PYEOF'
import os, sys
import psycopg

dsn = os.environ["QC_TRACE_DSN"]
org = os.environ["ORG"]

conn = psycopg.connect(dsn)

with conn.cursor() as cur:
    cur.execute("SELECT COUNT(*) FROM sessions WHERE org = %s", (org,))
    sessions = cur.fetchone()[0]
    cur.execute(
        "SELECT COUNT(*) FROM messages WHERE session_id IN (SELECT id FROM sessions WHERE org = %s)",
        (org,),
    )
    messages = cur.fetchone()[0]

print(f"==> Org: {org}")
print(f"    Sessions: {sessions}")
print(f"    Messages: {messages}")

if sessions == 0:
    print(f"No data found for org '{org}'. Nothing to delete.")
    sys.exit(0)

print(f"\nDelete all data for org '{org}'? (type 'yes' to confirm): ", end="", flush=True)
with open("/dev/tty") as tty:
    confirm = tty.readline().strip()
if confirm != "yes":
    print("Aborted.")
    sys.exit(1)

with conn.cursor() as cur:
    print("==> Deleting file_progress...")
    cur.execute(
        """
        DELETE FROM file_progress
        WHERE raw_file_path IN (
            SELECT DISTINCT raw_file_path FROM messages
            WHERE session_id IN (SELECT id FROM sessions WHERE org = %s)
            AND raw_file_path IS NOT NULL
        )
        """,
        (org,),
    )
    print(f"    Deleted {cur.rowcount} file_progress rows")

    print("==> Deleting sessions (cascades to messages, token_usage, tool_calls, tool_results)...")
    cur.execute("DELETE FROM sessions WHERE org = %s", (org,))
    print(f"    Deleted {cur.rowcount} sessions")

conn.commit()
conn.close()
print(f"\n✓ All data for org '{org}' deleted.")
PYEOF
