"""copy_dynamodb retries UnprocessedItems instead of counting them as written."""

import importlib.util
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]


def _load():
    spec = importlib.util.spec_from_file_location(
        "copy_dynamodb", ROOT / "scripts" / "cutover" / "copy_dynamodb.py"
    )
    mod = importlib.util.module_from_spec(spec)
    sys.modules["copy_dynamodb"] = mod
    spec.loader.exec_module(mod)
    return mod


mod = _load()


class FakeDdb:
    def __init__(self, unprocessed_first=None):
        self.calls = []
        self.unprocessed_first = list(unprocessed_first or [])
        self._first = True

    def batch_write_item(self, RequestItems):
        batch = RequestItems[mod.TABLE]
        self.calls.append(batch)
        if self._first and self.unprocessed_first:
            self._first = False
            leftover = [req for req in batch if req in self.unprocessed_first]
            return {"UnprocessedItems": {mod.TABLE: leftover} if leftover else {}}
        return {"UnprocessedItems": {}}


def test_batch_write_retries_unprocessed_items():
    items = [{"PK": {"S": "a"}}, {"PK": {"S": "b"}}]
    first = [{"PutRequest": {"Item": items[1]}}]
    client = FakeDdb(unprocessed_first=first)
    written = mod._batch_write_all(client, items, sleep=lambda _s: None)
    assert written == 2
    assert len(client.calls) == 2
    assert client.calls[1] == first


def test_batch_write_raises_if_unprocessed_remain():
    items = [{"PK": {"S": "a"}}]
    stuck = [{"PutRequest": {"Item": items[0]}}]
    client = FakeDdb(unprocessed_first=stuck)
    client._always = True

    def always_unprocessed(RequestItems):
        client.calls.append(RequestItems[mod.TABLE])
        return {"UnprocessedItems": {mod.TABLE: stuck}}

    client.batch_write_item = always_unprocessed
    try:
        mod._batch_write_all(client, items, sleep=lambda _s: None, max_attempts=3)
        raise AssertionError("expected RuntimeError")
    except RuntimeError as exc:
        assert "unprocessed" in str(exc)