#!/usr/bin/env python3
"""Delete expired project workbench events in bounded batches."""
from __future__ import annotations

import argparse
import asyncio
from collections import defaultdict
from pathlib import Path
import sys
from uuid import UUID

from psycopg import AsyncConnection
from psycopg.conninfo import conninfo_to_dict

REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
if str(REPOSITORY_ROOT) not in sys.path:
    sys.path.insert(0, str(REPOSITORY_ROOT))

from app.auth.project_dependencies import DB_ROLE_BIZ_DATA
from app.core.config import settings
from app.infra.db.metadb.database import SessionLocal, close_metadata_engine
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository


def _database_identity(dsn: str) -> tuple[str, str, str]:
    values = conninfo_to_dict(dsn)
    return (
        values.get("host") or values.get("hostaddr") or "localhost",
        values.get("port") or "5432",
        values.get("dbname") or values.get("user") or "",
    )


async def cleanup_project(
    project_id: UUID, dsn: str, retention_days: int, batch_size: int, dry_run: bool
) -> int:
    async with await AsyncConnection.connect(dsn) as conn:
        if dry_run:
            cursor = await conn.execute(
                """SELECT count(*) FROM workbench.domain_event_log
                   WHERE project_id=%s
                     AND occurred_at < now() - make_interval(days => %s)""",
                (project_id, retention_days),
            )
            return int((await cursor.fetchone())[0])
        deleted = 0
        while True:
            cursor = await conn.execute(
                """WITH expired AS (
                       SELECT event_seq FROM workbench.domain_event_log
                       WHERE project_id=%s
                         AND occurred_at < now() - make_interval(days => %s)
                       ORDER BY event_seq LIMIT %s
                   )
                   DELETE FROM workbench.domain_event_log e
                   USING expired WHERE e.event_seq=expired.event_seq""",
                (project_id, retention_days, batch_size),
            )
            await conn.commit()
            count = cursor.rowcount
            deleted += count
            if count < batch_size:
                return deleted


async def cleanup(
    retention_days: int, batch_size: int, project_id: UUID | None, dry_run: bool
) -> dict[UUID, int]:
    async with SessionLocal() as session:
        repository = MetadataRepository(session)
        routings = await repository.list_active_project_db_routings(DB_ROLE_BIZ_DATA)
    if project_id is not None:
        routings = [routing for routing in routings if routing.project_id == project_id]
        if not routings:
            raise ValueError("active project business database routing was not found")
    by_database = defaultdict(list)
    for routing in routings:
        by_database[_database_identity(routing.dsn)].append(routing)
    shared = [items for items in by_database.values() if len(items) > 1]
    if shared:
        project_ids = ", ".join(str(item.project_id) for items in shared for item in items)
        raise ValueError(f"shared business database routing is unsupported: {project_ids}")
    return {
        routing.project_id: await cleanup_project(
            routing.project_id, routing.dsn, retention_days, batch_size, dry_run
        )
        for routing in routings
    }


async def main(args) -> None:
    try:
        results = await cleanup(
            args.retention_days, args.batch_size, args.project_id, args.dry_run
        )
        for project_id, count in results.items():
            action = "would delete" if args.dry_run else "deleted"
            print(f"{project_id}: {action} {count} events")
    finally:
        await close_metadata_engine()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--retention-days", type=int, default=settings.WORKBENCH_EVENT_RETENTION_DAYS
    )
    parser.add_argument("--batch-size", type=int, default=5000)
    parser.add_argument("--project-id", type=UUID)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    if args.retention_days < 1 or args.batch_size < 1:
        parser.error("retention days and batch size must be positive")
    asyncio.run(main(args))