# -*- coding: utf-8 -*-
"""Location: ./tests/populate/cleanup.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Clean up populated test data via REST API DELETE endpoints.

Usage:
    python -m tests.populate.cleanup --confirm
    python -m tests.populate.cleanup --dry-run

Note: Teams and users cleanup is best-effort. Teams require owner tokens
(not admin), and user deletion may fail due to foreign key constraints.
For a complete wipe, delete the database file: rm mcp.db
"""

# Standard
import argparse
import asyncio
import logging
import os
import re
import sys
import time
from typing import Any, Dict, List

# Third-Party
from rich.console import Console

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Deletion order (respects dependencies - delete leaf entities first)
# (entity_name, list_endpoint, id_field, items_key_or_none)
# items_key: key in response dict containing items, or None for plain-list endpoints
DELETION_ORDER = [
    ("a2a_agents", "/a2a", "id", "agents"),
    ("servers", "/servers", "id", "servers"),
    ("prompts", "/prompts", "id", "prompts"),
    ("resources", "/resources", "id", "resources"),
    ("tools", "/tools", "id", "tools"),
    ("gateways", "/gateways", "id", "gateways"),
]

# Teams and users require special handling (ownership, cascading deps)
LOADTEST_PASSWORD = "LoadTest1234!"


async def cleanup_entities(
    base_url: str,
    email_domain: str = "loadtest.example.com",
    dry_run: bool = False,
    batch_size: int = 50,
) -> Dict[str, Any]:
    """Delete all entities with @loadtest.example.com via API."""
    # Third-Party
    import httpx

    # Generate admin token
    try:
        # First-Party
        from mcpgateway.utils.create_jwt_token import _create_jwt_token

        admin_token = os.environ.get("MCPGATEWAY_BEARER_TOKEN") or _create_jwt_token(
            data={"sub": "admin@example.com", "username": "admin@example.com"},
            expires_in_minutes=60,
            user_data={"email": "admin@example.com", "full_name": "Admin", "is_admin": True},
            teams=None,
        )
    except ImportError:
        admin_token = os.environ.get("MCPGATEWAY_BEARER_TOKEN", "")
        if not admin_token:
            raise RuntimeError("Set MCPGATEWAY_BEARER_TOKEN or install mcpgateway")

    headers = {"Authorization": f"Bearer {admin_token}", "Content-Type": "application/json"}
    results: Dict[str, Any] = {}

    async with httpx.AsyncClient(base_url=base_url, timeout=60.0, follow_redirects=True) as client:
        for entity_name, list_endpoint, id_field, items_key in DELETION_ORDER:
            deleted = 0
            errors = 0

            # Fetch all entities
            entity_ids: List[str] = []
            cursor = None
            page_size = 100
            loadtest_name_re = re.compile(r"^.+-(server|prompt|resource|agent|tool)-\d+-\d+$")

            while True:
                try:
                    params: Dict[str, Any] = {"limit": page_size}
                    if items_key:
                        params["include_pagination"] = "true"
                        if cursor:
                            params["cursor"] = cursor

                    resp = await client.get(list_endpoint, headers=headers, params=params)
                    if resp.status_code != 200:
                        break

                    data = resp.json()

                    # Extract items from response
                    if isinstance(data, dict):
                        items = data.get(items_key, []) if items_key else []
                        # Cursor-based uses nextCursor, offset-based has no cursor
                        cursor = data.get("nextCursor")
                    elif isinstance(data, list):
                        items = data
                        cursor = None
                    else:
                        break

                    if not items:
                        break

                    for item in items:
                        # Filter to loadtest entities by checking multiple identifying fields
                        searchable = " ".join(str(item.get(f, "")) for f in ("email", "created_by", "user_email", "name", "url", "uri", "endpointUrl", "description", "slug"))
                        item_name = str(item.get("name", ""))
                        is_loadtest = email_domain in searchable or loadtest_name_re.match(item_name)
                        if is_loadtest:
                            entity_ids.append(str(item.get(id_field, item.get("id", ""))))

                    # Stop if no more pages
                    if not items_key or not cursor or len(items) < page_size:
                        break

                except Exception as exc:
                    logger.warning(f"Failed to list {entity_name}: {exc}")
                    break

            if not entity_ids:
                results[entity_name] = {"found": 0, "deleted": 0, "errors": 0}
                continue

            if dry_run:
                logger.info(f"[DRY RUN] Would delete {len(entity_ids)} {entity_name}")
                results[entity_name] = {"found": len(entity_ids), "deleted": 0, "dry_run": True}
                continue

            # Delete in batches
            sem = asyncio.Semaphore(batch_size)

            async def _delete_one(eid: str, ep: str = list_endpoint, hdrs: Dict[str, str] = headers):
                nonlocal deleted, errors
                async with sem:
                    try:
                        delete_endpoint = f"{ep.rstrip('/')}/{eid}"
                        resp = await client.delete(delete_endpoint, headers=hdrs)
                        if resp.status_code in (200, 204, 404):
                            deleted += 1
                        else:
                            errors += 1
                    except Exception:
                        errors += 1

            # Process deletion in chunks
            for i in range(0, len(entity_ids), batch_size * 2):
                chunk = entity_ids[i : i + batch_size * 2]
                await asyncio.gather(*[_delete_one(eid) for eid in chunk])

            results[entity_name] = {"found": len(entity_ids), "deleted": deleted, "errors": errors}
            logger.info(f"Deleted {deleted}/{len(entity_ids)} {entity_name} (errors: {errors})")

        # --- Teams: delete via user tokens (owners only) ---
        teams_deleted = 0
        teams_errors = 0
        teams_found = 0

        # List all teams
        team_map: Dict[str, str] = {}  # team_id -> created_by email
        cursor = None
        while True:
            params_t: Dict[str, Any] = {"limit": 100, "include_pagination": "true"}
            if cursor:
                params_t["cursor"] = cursor
            try:
                resp = await client.get("/teams/", headers=headers, params=params_t)
                if resp.status_code != 200:
                    break
                data = resp.json()
                teams_list = data.get("teams", []) if isinstance(data, dict) else data
                for t in teams_list:
                    created_by = t.get("created_by", "")
                    if email_domain in str(created_by) or email_domain in str(t.get("slug", "")):
                        team_map[t["id"]] = created_by
                cursor = data.get("nextCursor") if isinstance(data, dict) else None
                if not cursor or len(teams_list) < 100:
                    break
            except Exception:
                break

        teams_found = len(team_map)
        if teams_found > 0 and not dry_run:
            # Group teams by owner
            owner_teams: Dict[str, List[str]] = {}
            for tid, owner in team_map.items():
                owner_teams.setdefault(owner, []).append(tid)

            # Login as each owner and delete their teams
            for owner_email, team_ids_to_delete in owner_teams.items():
                try:
                    login_resp = await client.post(
                        "/auth/email/login",
                        json={"email": owner_email, "password": LOADTEST_PASSWORD},
                        headers={"Content-Type": "application/json"},
                    )
                    if login_resp.status_code != 200:
                        teams_errors += len(team_ids_to_delete)
                        continue
                    user_jwt = login_resp.json().get("access_token") or login_resp.json().get("token")
                    if not user_jwt:
                        teams_errors += len(team_ids_to_delete)
                        continue
                    user_hdrs = {"Authorization": f"Bearer {user_jwt}", "Content-Type": "application/json"}
                    for tid in team_ids_to_delete:
                        try:
                            resp = await client.delete(f"/teams/{tid}", headers=user_hdrs)
                            if resp.status_code in (200, 204, 404):
                                teams_deleted += 1
                            else:
                                teams_errors += 1
                        except Exception:
                            teams_errors += 1
                except Exception:
                    teams_errors += len(team_ids_to_delete)
        elif teams_found > 0 and dry_run:
            logger.info(f"[DRY RUN] Would delete {teams_found} teams")

        results["teams"] = {"found": teams_found, "deleted": teams_deleted, "errors": teams_errors}
        if teams_found > 0 and not dry_run:
            logger.info(f"Deleted {teams_deleted}/{teams_found} teams (errors: {teams_errors})")

        # --- Users: delete via admin endpoint ---
        users_deleted = 0
        users_errors = 0
        try:
            resp = await client.get("/auth/email/admin/users", headers=headers, params={"limit": 500})
            user_list = resp.json() if resp.status_code == 200 else []
            user_emails = [u["email"] for u in user_list if isinstance(u, dict) and email_domain in str(u.get("email", ""))]
        except Exception:
            user_emails = []

        if user_emails and not dry_run:
            for ue in user_emails:
                try:
                    resp = await client.delete(f"/auth/email/admin/users/{ue}", headers=headers)
                    if resp.status_code in (200, 204, 404):
                        users_deleted += 1
                    else:
                        users_errors += 1
                except Exception:
                    users_errors += 1
        elif user_emails and dry_run:
            logger.info(f"[DRY RUN] Would delete {len(user_emails)} users")

        results["users"] = {"found": len(user_emails), "deleted": users_deleted, "errors": users_errors}
        if user_emails and not dry_run:
            logger.info(f"Deleted {users_deleted}/{len(user_emails)} users (errors: {users_errors})")

    return results


def main():
    parser = argparse.ArgumentParser(description="Clean up loadtest data via REST API")
    parser.add_argument("--base-url", type=str, default=os.environ.get("MCPGATEWAY_BASE_URL", "http://localhost:8080"))
    parser.add_argument("--email-domain", type=str, default="loadtest.example.com")
    parser.add_argument("--confirm", action="store_true", help="Confirm deletion (required)")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be deleted")
    parser.add_argument("--batch-size", type=int, default=50, help="Concurrent delete batch size")
    args = parser.parse_args()

    console = Console()

    if not args.confirm and not args.dry_run:
        console.print("[red]Must specify --confirm to delete data (or use --dry-run)[/red]")
        sys.exit(1)

    console.print(f"\n[bold cyan]Cleaning up loadtest data at {args.base_url}[/bold cyan]")
    console.print(f"[dim]Email domain filter: @{args.email_domain}[/dim]\n")

    if args.dry_run:
        console.print("[yellow]DRY RUN - No entities will be deleted[/yellow]\n")

    start_time = time.time()
    results = asyncio.run(cleanup_entities(args.base_url, args.email_domain, args.dry_run, args.batch_size))
    duration = time.time() - start_time

    # Print summary
    total_deleted = sum(r.get("deleted", 0) for r in results.values())
    total_errors = sum(r.get("errors", 0) for r in results.values())
    total_found = sum(r.get("found", 0) for r in results.values())

    console.print(f"\n{'=' * 60}")
    console.print("[bold]Cleanup Summary[/bold]")
    console.print(f"{'=' * 60}")

    for entity, info in results.items():
        found = info.get("found", 0)
        deleted = info.get("deleted", 0)
        errs = info.get("errors", 0)
        if found > 0:
            console.print(f"  {entity:20s}  found: {found:>8,}  deleted: {deleted:>8,}  errors: {errs:>4,}")

    console.print(f"\n  {'Total':20s}  found: {total_found:>8,}  deleted: {total_deleted:>8,}  errors: {total_errors:>4,}")
    console.print(f"  Duration: {duration:.2f}s")
    console.print(f"  Dry Run: {args.dry_run}")
    console.print(f"{'=' * 60}\n")

    sys.exit(0)


if __name__ == "__main__":
    main()
