from datetime import datetime
from typing import Any, Dict, List, Optional

from config import POSTGRES_CONNECTION_STRING


async def delete_conversation_by_chat_id(chat_id: str) -> Dict[str, Any]:
    """Delete wx-agent LangGraph checkpoint state for one chat_id."""
    chat_id = str(chat_id or "").strip()
    if not chat_id:
        raise ValueError("chat_id is empty")

    counts = {
        "checkpoint_writes": 0,
        "checkpoint_blobs": 0,
        "checkpoints": 0,
    }
    async with await _connect_pool() as pool:
        async with pool.connection() as conn:
            async with conn.cursor() as cur:
                await _delete_checkpoint_threads(cur, chat_id, counts)
    return {"success": True, "chat_id": chat_id, "deleted": counts}


async def delete_conversations_by_time(
    before: Optional[datetime] = None,
    after: Optional[datetime] = None,
) -> Dict[str, Any]:
    """Delete LangGraph checkpoint state whose checkpoint timestamp falls in a time range."""
    if before is None and after is None:
        raise ValueError("before or after is required")

    where_parts = []
    params: List[Any] = []
    if before is not None:
        where_parts.append("(checkpoint->>'ts')::timestamptz < %s")
        params.append(before)
    if after is not None:
        where_parts.append("(checkpoint->>'ts')::timestamptz >= %s")
        params.append(after)
    where_sql = " AND ".join(where_parts)

    counts = {
        "checkpoint_writes": 0,
        "checkpoint_blobs": 0,
        "checkpoints": 0,
    }
    deleted_thread_ids: List[str] = []

    async with await _connect_pool() as pool:
        async with pool.connection() as conn:
            async with conn.cursor() as cur:
                await cur.execute(
                    f"SELECT DISTINCT thread_id FROM checkpoints WHERE {where_sql}",
                    params,
                )
                deleted_thread_ids = [str(row[0]) for row in await cur.fetchall() if row and row[0]]

                for thread_id in deleted_thread_ids:
                    await _delete_exact_thread(cur, thread_id, counts)

    return {
        "success": True,
        "before": before.isoformat() if before else "",
        "after": after.isoformat() if after else "",
        "thread_count": len(deleted_thread_ids),
        "thread_ids": deleted_thread_ids,
        "deleted": counts,
    }


async def _delete_checkpoint_threads(cur, chat_id: str, counts: Dict[str, int]) -> None:
    thread_like = f"{chat_id}:%"
    await cur.execute(
        "DELETE FROM checkpoint_writes WHERE thread_id = %s OR thread_id LIKE %s",
        (chat_id, thread_like),
    )
    counts["checkpoint_writes"] += cur.rowcount or 0
    await cur.execute(
        "DELETE FROM checkpoint_blobs WHERE thread_id = %s OR thread_id LIKE %s",
        (chat_id, thread_like),
    )
    counts["checkpoint_blobs"] += cur.rowcount or 0
    await cur.execute(
        "DELETE FROM checkpoints WHERE thread_id = %s OR thread_id LIKE %s",
        (chat_id, thread_like),
    )
    counts["checkpoints"] += cur.rowcount or 0


async def _delete_exact_thread(cur, thread_id: str, counts: Dict[str, int]) -> None:
    await cur.execute("DELETE FROM checkpoint_writes WHERE thread_id = %s", (thread_id,))
    counts["checkpoint_writes"] += cur.rowcount or 0
    await cur.execute("DELETE FROM checkpoint_blobs WHERE thread_id = %s", (thread_id,))
    counts["checkpoint_blobs"] += cur.rowcount or 0
    await cur.execute("DELETE FROM checkpoints WHERE thread_id = %s", (thread_id,))
    counts["checkpoints"] += cur.rowcount or 0


async def _connect_pool():
    from psycopg_pool import AsyncConnectionPool

    return AsyncConnectionPool(
        POSTGRES_CONNECTION_STRING,
        kwargs={"autocommit": True},
    )