"""Opt-in end-to-end archive test for a disposable PostgreSQL database.

Run with DEERFLOW_TEST_POSTGRES_URL set to an isolated database. This test
archives, purges, transparently reads, and restores one uniquely named thread.
"""

from __future__ import annotations

import os
import uuid
from datetime import datetime

import pytest
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from psycopg.rows import dict_row
from psycopg_pool import AsyncConnectionPool
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from app.gateway.routers.chat_archive import ChatArchiveRunner, _save_job, _upsert_month
from deerflow.persistence.base import Base
from deerflow.persistence.thread_meta.model import ThreadMetaRow
from deerflow.runtime.checkpointer.chat_archive import archive_db_path, ensure_catalog
from deerflow.runtime.checkpointer.checkpoint_migration import DualReadCheckpointer


async def _checkpoint_ids(saver, thread_id: str) -> list[str]:
    return [
        item.config["configurable"]["checkpoint_id"]
        async for item in saver.alist({"configurable": {"thread_id": thread_id}})
    ]


@pytest.mark.asyncio
async def test_postgres_month_archive_purge_fallback_and_rollback(tmp_path):
    conninfo = os.getenv("DEERFLOW_TEST_POSTGRES_URL")
    if not conninfo:
        pytest.skip("set DEERFLOW_TEST_POSTGRES_URL to an isolated PostgreSQL database")

    thread_id = f"archive-qa-{uuid.uuid4().hex}"
    month = "2025-02"
    pool = AsyncConnectionPool(
        conninfo=conninfo,
        min_size=1,
        max_size=3,
        kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row},
        open=False,
    )
    business_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'business.db'}")
    try:
        await pool.open(wait=True)
        primary = AsyncPostgresSaver(pool)
        await primary.setup()

        async with business_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all, tables=[ThreadMetaRow.__table__])
        session_factory = async_sessionmaker(business_engine, expire_on_commit=False)
        created_at = datetime(2025, 2, 12, 12, 0)
        async with session_factory() as session:
            session.add(ThreadMetaRow(
                thread_id=thread_id,
                assistant_id="archive-qa",
                user_id="archive-qa",
                display_name="isolated PostgreSQL archive test",
                status="idle",
                metadata_json={},
                created_at=created_at,
                updated_at=created_at,
            ))
            await session.commit()

        checkpoint = empty_checkpoint()
        await primary.aput(
            {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
            checkpoint,
            {"source": "isolated-postgres-qa"},
            checkpoint["channel_versions"],
        )
        assert await _checkpoint_ids(primary, thread_id) == [checkpoint["id"]]

        root = tmp_path / "chat-archives"
        ensure_catalog(root)
        job_id = f"archive-qa-{uuid.uuid4().hex}"
        now = "2026-10-25T00:00:00+08:00"
        _save_job(root, {
            "job_id": job_id,
            "month": month,
            "action": "archive",
            "status": "running",
            "total_threads": 1,
            "processed_threads": 0,
            "failed_threads": 0,
            "current_thread_id": None,
            "message": "",
            "error": None,
            "created_at": now,
            "updated_at": now,
        })
        _upsert_month(root, month, status="copying", total=1, job_id=job_id)
        runner = ChatArchiveRunner(root, primary, session_factory)
        await runner._archive_month(job_id, month)

        archived_path = archive_db_path(root, month)
        async with AsyncSqliteSaver.from_conn_string(str(archived_path)) as archived:
            await archived.setup()
            assert await _checkpoint_ids(archived, thread_id) == [checkpoint["id"]]

        await runner._purge_month(job_id, month)
        assert await _checkpoint_ids(primary, thread_id) == []

        proxy = DualReadCheckpointer(primary, [], archive_dir=root)
        recovered = await proxy.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
        assert recovered is not None
        assert recovered.checkpoint["id"] == checkpoint["id"]
        await proxy.aclose_sources()

        await runner._rollback_month(job_id, month)
        assert await _checkpoint_ids(primary, thread_id) == [checkpoint["id"]]
    finally:
        await business_engine.dispose()
        await pool.close()