#!/usr/bin/env bash

THREADS_DB="${HOME}/Library/Application Support/Zed/threads/threads.db"

if ! command -v sqlite3 &> /dev/null; then
    echo "Error: sqlite3 is not installed. Please install SQLite to continue."
    exit 1
fi

if ! command -v duckdb &> /dev/null; then
    echo "Error: duckdb is not installed. Please install DuckDB to continue."
    exit 1
fi

if ! command -v zstd &> /dev/null; then
    echo "Error: zstd is not installed. Please install zstd to continue."
    exit 1
fi

if [ ! -f "${THREADS_DB}" ]; then
    echo "Error: Threads database not found at ${THREADS_DB}"
    echo "Make sure Zed is installed and has created the threads database."
    exit 1
fi

rm -rf "$(pwd)/threads/"
mkdir -p "$(pwd)/threads/"{0.1.0,0.2.0,0.3.0}

cp "${THREADS_DB}" .

sqlite3 threads.db "SELECT id FROM threads;" | \
  while read id; do

    sqlite3 threads.db "SELECT writefile('temp_${id}.zst', data) FROM threads WHERE id = '$id';" > /dev/null

    zstd --rm -q -d "temp_${id}.zst" -o "$(pwd)/threads/${id}.json" > /dev/null

    [[ $(jq 'has("version") and .version == "0.1.0"' "$(pwd)/threads/${id}.json") == "true" ]] && mv "$(pwd)/threads/${id}.json" "$(pwd)/threads/0.1.0/${id}.json" && continue

    [[ $(jq 'has("version") and .version == "0.2.0"' "$(pwd)/threads/${id}.json") == "true" ]] && mv "$(pwd)/threads/${id}.json" "$(pwd)/threads/0.2.0/${id}.json" && continue

    [[ $(jq 'has("version") and .version == "0.3.0"' "$(pwd)/threads/${id}.json") == "true" ]] && mv "$(pwd)/threads/${id}.json" "$(pwd)/threads/0.3.0/${id}.json"

  done

duckdb -csv <<EOQ
ATTACH 'threads.db' AS zed;
WITH v01 AS (
  FROM read_json('threads/0.1.0/*.json') AS t
  SELECT
    updated_at::DATE AS ts,
    parse_filename(t.filename, true) AS id,
    UNNEST(t.request_token_usage) AS u
),
v02 AS (
  FROM read_json('threads/0.2.0/*.json') AS t
  SELECT
    updated_at::DATE AS ts,
    parse_filename(t.filename, true) AS id,
    UNNEST(t.request_token_usage) AS u
),
v03 AS (
  FROM read_json('threads/0.3.0/*.json') AS t
  SELECT
    updated_at::DATE AS ts,
    parse_filename(t.filename, true) AS id,
    UNNEST(map_entries(t.request_token_usage)) AS entry
),
tokens AS (
  SELECT
    id,
    ts,
    SUM(input_tokens)    AS total_input_tokens,
    SUM(output_tokens)   AS total_output_tokens
  FROM (
    SELECT id, ts, u.input_tokens AS input_tokens, u.output_tokens AS output_tokens FROM v01
    UNION ALL
    SELECT id, ts, u.input_tokens, u.output_tokens FROM v02
    UNION ALL
    SELECT id, ts, entry.value.input_tokens, entry.value.output_tokens FROM v03
  ) AS all_flat
  GROUP BY id, ts
)
FROM
  tokens tok,
  zed.threads thr
SELECT
  tok.id,
  tok.ts,
  tok.total_input_tokens,
  tok.total_output_tokens,
  thr.summary
WHERE
  tok.id = thr.id
ORDER BY tok.ts
EOQ