from datetime import UTC, datetime, timedelta
from uuid import uuid4

import pytest
from sqlalchemy import delete, update

from app.infrastructure.db import SessionFactory
from app.model.conversation import ConversationMessage
from app.model.platform import AgentRun, RequestIdempotency
from app.repository.agent_run_repository import AgentRunRepository


@pytest.mark.integration
@pytest.mark.asyncio
async def test_worker_lease_renew_and_takeover() -> None:
    now = datetime.now(UTC).replace(tzinfo=None)
    session_id = f"lease-{uuid4()}"
    run_id = str(uuid4())
    trace_id = str(uuid4())
    async with SessionFactory() as session:
        idem = RequestIdempotency(
            user_id=1, session_id=session_id, agent_type="customer_service",
            idempotency_key=f"lease-key-{uuid4()}", request_hash="a" * 64,
            trace_id=trace_id, status="processing", expire_at=now + timedelta(hours=1),
            created_at=now, updated_at=now,
        )
        message = ConversationMessage(
            session_id=session_id, customer_id=1, portal="api", role="user",
            content="lease test", trace_id=trace_id, created_at=now,
        )
        session.add_all([idem, message])
        await session.flush()
        run = AgentRun(
            run_id=run_id, idempotency_id=idem.id, session_id=session_id, user_id=1,
            agent_type="customer_service", trace_id=trace_id, request_message_id=message.id,
            created_at=now, updated_at=now,
        )
        session.add(run)
        await session.commit()
        try:
            repo = AgentRunRepository(session)
            assert await repo.claim(run_id, "worker-a", 60)
            assert await repo.renew(run_id, "worker-a", 60)
            assert not await repo.claim(run_id, "worker-b", 60)
            await session.execute(
                update(AgentRun).where(AgentRun.run_id == run_id).values(
                    locked_until=now - timedelta(seconds=1)
                )
            )
            await session.commit()
            assert await repo.claim(run_id, "worker-b", 60)
        finally:
            await session.execute(delete(AgentRun).where(AgentRun.run_id == run_id))
            await session.execute(
                delete(RequestIdempotency).where(RequestIdempotency.id == idem.id)
            )
            await session.execute(
                delete(ConversationMessage).where(ConversationMessage.session_id == session_id)
            )
            await session.commit()