#!/bin/bash
# Usage: ./check_id.sh <N> [identities.yaml]
# Prints the Nth ID (0-indexed) from the S3 source and checks if it exists in Snowflake.
N=${1:?Usage: ./check_id.sh <N> [identities.yaml]}
IDENTITY=${2:-identities.yaml}

uv run python -c "
import sys
from generate_data import make_source, load_identity, _ID_LINE_LENGTH
import snowflake.connector

n = int(sys.argv[1])
identity = load_identity(sys.argv[2])
source = make_source('s3://q-chalk-test-data/surveyors-guild/adventurer-ids.txt', identity, skip=n)
id_val = next(source.iter_ids(), None)
if id_val is None:
    print(f'No ID at index {n}')
    sys.exit(1)
print(f'Row {n}: {id_val}')

sf = identity['snowflake']
conn = snowflake.connector.connect(account=sf['account'], user=sf['user'], password=sf['password'], warehouse=sf.get('warehouse'), role=sf.get('role'))
cur = conn.cursor()
cur.execute('SELECT COUNT(*) FROM CHALK_Q.SURVEYORS_GUILD.ADVENTURER WHERE id = %s', (id_val,))
count = cur.fetchone()[0]
print(f'Exists in Snowflake: {\"yes\" if count > 0 else \"no\"} (count={count})')
cur.close()
conn.close()
" "$N" "$IDENTITY"