import asyncio
from datetime import datetime

from app.repositories import parse_run_repo


class _Result:
    def __init__(self, modified_count=1):
        self.modified_count = modified_count


class _Collection:
    def __init__(self):
        self.query = None
        self.update = None

    async def update_one(self, query, update):
        self.query = query
        self.update = update
        return _Result(modified_count=1)


class _DB:
    def __init__(self):
        self.doc_parse_runs = _Collection()


def test_mark_run_queued_updates_status(monkeypatch):
    fake_db = _DB()
    monkeypatch.setattr(parse_run_repo, "get_db", lambda: fake_db)

    asyncio.run(parse_run_repo.mark_run_queued("run_1"))

    assert fake_db.doc_parse_runs.query == {"run_id": "run_1"}
    assert fake_db.doc_parse_runs.update["$set"]["status"] == "queued"


def test_mark_run_retrying_sets_retry_fields(monkeypatch):
    fake_db = _DB()
    monkeypatch.setattr(parse_run_repo, "get_db", lambda: fake_db)

    next_retry_at = datetime.utcnow()
    asyncio.run(
        parse_run_repo.mark_run_retrying(
            run_id="run_1",
            retry_count=2,
            error="timeout",
            next_retry_at=next_retry_at,
        )
    )

    payload = fake_db.doc_parse_runs.update["$set"]
    assert payload["status"] == "retrying"
    assert payload["retry_count"] == 2
    assert payload["error"] == "timeout"
    assert payload["last_error"] == "timeout"
    assert payload["next_retry_at"] == next_retry_at


def test_claim_run_for_processing_uses_atomic_status_filter(monkeypatch):
    fake_db = _DB()
    monkeypatch.setattr(parse_run_repo, "get_db", lambda: fake_db)

    claimed = asyncio.run(parse_run_repo.claim_run_for_processing("run_1", stale_after_seconds=60))

    assert claimed is True
    query = fake_db.doc_parse_runs.query
    assert query["run_id"] == "run_1"
    assert "$or" in query

    statuses = query["$or"][0]["status"]["$in"]
    assert "pending" in statuses
    assert "queued" in statuses
    assert "retrying" in statuses

    stale_processing_clause = query["$or"][1]
    assert stale_processing_clause["status"] == "processing"
    assert "$lte" in stale_processing_clause["started_at"]

    update = fake_db.doc_parse_runs.update["$set"]
    assert update["status"] == "processing"
    assert "started_at" in update