#!/usr/bin/env python3
"""Delete memories that were imported from Neo4j (tagged metadata.source='neo4j_legacy_import').

Usage:
    docker compose exec backend python /tmp/cleanup.py --user alice
    docker compose exec backend python /tmp/cleanup.py --users alice,hetashree,manju
    docker compose exec backend python /tmp/cleanup.py --dry-run --users alice
"""
import argparse
import json
import os
import re
import sys

import httpx
from qdrant_client import QdrantClient
from qdrant_client.http import models

QDRANT_HOST = os.environ.get("QDRANT_HOST", "qdrant")
QDRANT_PORT = int(os.environ.get("QDRANT_PORT", "6333"))
QDRANT_COLLECTION = os.environ.get("QDRANT_COLLECTION_NAME", "mem0_v3")
MEM0_URL = os.environ.get("MEM0_URL", "http://localhost:8000")


def load_api_keys() -> dict:
    raw = os.environ.get("API_KEYS", "{}")
    keys = json.loads(raw)
    inv: dict = {}
    for k, u in keys.items():
        inv.setdefault(u, k)
    return inv


def scroll_legacy_for(client: QdrantClient, user_id: str) -> list:
    """Return all memory IDs in mem0_v3 with source='neo4j_legacy_import' for the user."""
    ids = []
    offset = None
    while True:
        points, offset = client.scroll(
            collection_name=QDRANT_COLLECTION,
            scroll_filter=models.Filter(
                must=[
                    models.FieldCondition(key="user_id", match=models.MatchValue(value=user_id)),
                    models.FieldCondition(key="source", match=models.MatchValue(value="neo4j_legacy_import")),
                ]
            ),
            limit=128,
            with_payload=False,
            with_vectors=False,
            offset=offset,
        )
        ids.extend(p.id for p in points)
        if offset is None:
            break
    return ids


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--user", help="Single user to clean")
    ap.add_argument("--users", help="Comma-separated user list")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    targets = []
    if args.user:
        targets.append(args.user)
    if args.users:
        targets.extend(u.strip() for u in args.users.split(","))
    if not targets:
        print("Usage: --user X or --users a,b,c", file=sys.stderr)
        return 2

    api_keys = load_api_keys()
    client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
    http = httpx.Client(timeout=60.0)

    for user_id in targets:
        ids = scroll_legacy_for(client, user_id)
        print(f"{user_id}: found {len(ids)} legacy-import memories")
        if args.dry_run:
            continue
        key = api_keys.get(user_id)
        if not key:
            print(f"  no API key for {user_id} — skipping")
            continue
        deleted = errors = 0
        for mid in ids:
            r = http.delete(f"{MEM0_URL}/memories/{mid}", headers={"X-API-Key": key})
            if r.status_code == 200:
                deleted += 1
            else:
                errors += 1
                print(f"  failed {mid}: HTTP {r.status_code} {r.text[:120]}")
        print(f"  deleted={deleted} errors={errors}")
    return 0


if __name__ == "__main__":
    sys.exit(main())