#
# Copyright (C) 2025-present ScyllaDB
#
# SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.0
#

from test.pylib.manager_client import ManagerClient
from test.pylib.repair import load_tablet_sstables_repaired_at, create_table_insert_data_for_repair
from test.pylib.tablets import get_all_tablet_replicas
from test.cluster.tasks.task_manager_client import TaskManagerClient
from test.cluster.util import reconnect_driver, find_server_by_host_id, get_topology_coordinator, new_test_keyspace, new_test_table, trigger_stepdown

from cassandra.query import ConsistencyLevel

import pytest
import asyncio
import logging
import time
import glob
import os
import subprocess
import json
import socket
import random
import requests
import re

logger = logging.getLogger(__name__)

async def inject_error_on(manager, error_name, servers, params = {}):
    errs = [manager.api.enable_injection(s.ip_addr, error_name, False, params) for s in servers]
    await asyncio.gather(*errs)

async def inject_error_off(manager, error_name, servers):
    errs = [manager.api.disable_injection(s.ip_addr, error_name) for s in servers]
    await asyncio.gather(*errs)

def get_sstables_repaired_at(m, token):
    return m[str(token)]

def get_sstables(workdir, ks, table):
    base_pattern = os.path.join(workdir, "data", f"{ks}*", f"{table}-*","*-Data.db")
    sstables = glob.glob(base_pattern)
    return sstables

async def get_sst_status(run, log):
    sst_add = await log.grep(rf'.*Added sst.*for incremental repair')
    sst_skip = await log.grep(rf'.*Skipped adding sst.*for incremental repair')
    sst_mark = await log.grep(rf'.*Marking.*for incremental repair')
    logging.info(f'{run=}: {sst_add=} {sst_skip=} {sst_mark=}');
    logging.info(f'{run=}: {len(sst_add)=} {len(sst_skip)=} {len(sst_mark)=}');
    return sst_add, sst_skip, sst_mark

async def insert_keys(cql, ks, start_key, end_key):
    await asyncio.gather(*[cql.run_async(f"INSERT INTO {ks}.test (pk, c) VALUES ({k}, {k});") for k in range(start_key, end_key)])

def get_repaired_at_from_sst(sst_file, scylla_path):
    try:
        cmd = [scylla_path, "sstable", "dump-statistics", sst_file]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        output = json.loads(result.stdout)
        repaired_at = output.get("sstables", {}).get(sst_file, {}).get("stats", {}).get("repaired_at")
        return repaired_at
    except subprocess.CalledProcessError as e:
        logging.error(f"Command failed with error: {e}")
    except json.JSONDecodeError as e:
        logging.error(f"Failed to parse JSON output: {e}")
    except Exception as e:
        logging.error(f"An unexpected error occurred: {e}")

    return None

def get_keys_from_sst(sst_file, scylla_path):
    try:
        cmd = [scylla_path, "sstable", "dump-data", sst_file]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        output = json.loads(result.stdout)
        keys = []
        for sstable_data in output.get("sstables", {}).values():
            for entry in sstable_data:
                key_value = entry.get("key", {}).get("value")
                if key_value:
                    keys.append(int(key_value))
        return keys
    except subprocess.CalledProcessError as e:
        logging.error(f"Command failed with error: {e}")
    except json.JSONDecodeError as e:
        logging.error(f"Failed to parse JSON output: {e}")
    except Exception as e:
        logging.error(f"An unexpected error occurred: {e}")
    return []

def get_metrics(server, metric_name):
    num = 0
    metrics = requests.get(f"http://{server.ip_addr}:9180/metrics").text
    logger.info(f"Got for {metric_name=} {metrics=}");
    pattern = re.compile(rf"^{metric_name}")
    for metric in metrics.split('\n'):
        if pattern.match(metric) is not None:
            num += int(metric.split()[1])
    return num

def get_incremental_repair_sst_skipped_bytes(server):
    return get_metrics(server, "scylla_repair_inc_sst_skipped_bytes")

def get_incremental_repair_sst_read_bytes(server):
    return get_metrics(server, "scylla_repair_inc_sst_read_bytes")

def get_repair_tablet_time_ms(server):
    return get_metrics(server, "scylla_repair_tablet_time_ms")

async def get_sstables_for_server(manager, server, ks):
    node_workdir = await manager.server_get_workdir(server.server_id)
    sstables = get_sstables(node_workdir, ks, 'test')
    return sstables

async def verify_repaired_and_unrepaired_keys(manager, scylla_path, servers, ks, repaired_keys, unrepaired_keys):
    async def process_server_keys(server):
        node_workdir = await manager.server_get_workdir(server.server_id)
        sstables = get_sstables(node_workdir, ks, 'test')
        nr_sstables = len(sstables)
        async def process_sst(sst):
            keys = get_keys_from_sst(sst, scylla_path)
            repaired_at = get_repaired_at_from_sst(sst, scylla_path)
            logger.info(f"Got {node_workdir=} {nr_sstables=} {sst=} {repaired_at=} {keys=} {repaired_keys=} {unrepaired_keys=}")
            if repaired_at == 0 or repaired_at is None:
                for key in keys:
                    assert node_workdir and sst and key in unrepaired_keys
            elif repaired_at > 0:
                for key in keys:
                    assert node_workdir and sst and key in repaired_keys
        await asyncio.gather(*(process_sst(sst) for sst in sstables))
    await asyncio.gather(*[process_server_keys(server) for server in servers])

async def verify_max_repaired_at(manager, scylla_path, servers, ks, max_repaired_at):
    async def process_server_keys(server):
        node_workdir = await manager.server_get_workdir(server.server_id)
        sstables = get_sstables(node_workdir, ks, 'test')
        nr_sstables = len(sstables)
        for sst in sstables:
            keys = get_keys_from_sst(sst, scylla_path)
            repaired_at = get_repaired_at_from_sst(sst, scylla_path)
            logger.info(f"Got {node_workdir=} {nr_sstables=} {sst=} {repaired_at=}")
            assert repaired_at <= max_repaired_at
    await asyncio.gather(*[process_server_keys(server) for server in servers])

async def trigger_tablet_merge(manager, servers, logs):
    s1_log = logs[0]
    s1_mark = await s1_log.mark()
    await inject_error_on(manager, "tablet_force_tablet_count_decrease", servers)
    await s1_log.wait_for('Detected tablet merge for table', from_mark=s1_mark)
    await inject_error_off(manager, "tablet_force_tablet_count_decrease", servers)

async def prepare_cluster_for_incremental_repair(manager, nr_keys = 100 , cmdline = []):
    servers, cql, hosts, ks, table_id = await create_table_insert_data_for_repair(manager, nr_keys=nr_keys, cmdline=cmdline)
    repaired_keys = set(range(0, nr_keys))
    unrepaired_keys = set()
    current_key = nr_keys
    token = 'all'
    logs = []
    for s in servers:
        logs.append(await manager.server_open_log(s.server_id))
    return servers, cql, hosts, ks, table_id, logs, repaired_keys,  unrepaired_keys, current_key, token

@pytest.mark.asyncio
async def test_tablet_repair_sstable_skipped_read_metrics(manager: ManagerClient):
    servers, cql, hosts, ks, table_id, logs, _, _, _, token = await prepare_cluster_for_incremental_repair(manager)

    await insert_keys(cql, ks, 0, 100)

    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')
    skipped_bytes = get_incremental_repair_sst_skipped_bytes(servers[0])
    read_bytes = get_incremental_repair_sst_read_bytes(servers[0])
    # Nothing to skip. Repair all data.
    assert skipped_bytes == 0
    assert read_bytes > 0

    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')
    skipped_bytes2 = get_incremental_repair_sst_skipped_bytes(servers[0])
    read_bytes2 = get_incremental_repair_sst_read_bytes(servers[0])
    # Skip all. Nothing to repair
    assert skipped_bytes2 > 0
    assert read_bytes2 == read_bytes

    await insert_keys(cql, ks, 200, 300)

    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')
    skipped_bytes3 = get_incremental_repair_sst_skipped_bytes(servers[0])
    read_bytes3 = get_incremental_repair_sst_read_bytes(servers[0])
    # Both skipped and read bytes should grow
    assert skipped_bytes3 > skipped_bytes2
    assert read_bytes3 > read_bytes2

@pytest.mark.asyncio
async def test_tablet_incremental_repair(manager: ManagerClient):
    servers, cql, hosts, ks, table_id = await create_table_insert_data_for_repair(manager, fast_stats_refresh=False, disable_flush_cache_time=True)
    token = -1
    logs = []
    for s in servers:
        logs.append(await manager.server_open_log(s.server_id))

    sstables_repaired_at = 0

    async def get_compaction_status(run, log):
        disable = await log.grep(rf'.*Disabled compaction for.*for incremental repair')
        enable = await log.grep(rf'.*Re-enabled compaction for.*for incremental repair')
        logging.info(f'{run=}: {disable=} {enable=}');
        return disable, enable

    async def get_sst_status(run, log):
        sst_add = await log.grep(rf'.*Added sst.*for incremental repair')
        sst_skip = await log.grep(rf'.*Skipped adding sst.*for incremental repair')
        sst_mark = await log.grep(rf'.*Marking.*for incremental repair')
        logging.info(f'{run=}: {sst_add=} {sst_skip=} {sst_mark=}');
        logging.info(f'{run=}: {len(sst_add)=} {len(sst_skip)=} {len(sst_mark)=}');
        return sst_add, sst_skip, sst_mark

    map0 = await load_tablet_sstables_repaired_at(manager, cql, servers[0], hosts[0], table_id)
    logging.info(f'map1={map0}')
    assert get_sstables_repaired_at(map0, token) == sstables_repaired_at

    # First repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')
    map1 = await load_tablet_sstables_repaired_at(manager, cql, servers[0], hosts[0], table_id)
    logging.info(f'map1={map1}')
    # Check sstables_repaired_at is increased by 1
    assert get_sstables_repaired_at(map1, token) == sstables_repaired_at + 1
    # Check one sstable is added and marked as repaired
    for log in logs:
        sst_add, sst_skip, sst_mark = await get_sst_status("First", log)
        assert len(sst_add) == 1
        assert len(sst_skip) == 0
        assert len(sst_mark) == 1
        disable, enable = await get_compaction_status("First", log)
        assert len(disable) == 1
        assert len(enable) == 1

    # Second repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')
    map2 = await load_tablet_sstables_repaired_at(manager, cql, servers[0], hosts[0], table_id)
    logging.info(f'map2={map2}')
    # Check sstables_repaired_at is increased by 1
    assert get_sstables_repaired_at(map2, token) == sstables_repaired_at + 2
    # Check one sstable is skipped
    for log in logs:
        sst_add, sst_skip, sst_mark = await get_sst_status("Second", log)
        assert len(sst_add) == 1
        assert len(sst_skip) == 1
        assert len(sst_mark) == 1
        disable, enable = await get_compaction_status("Second", log)
        assert len(disable) == 2
        assert len(enable) == 2

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_error(manager: ManagerClient):
    servers, cql, hosts, ks, table_id = await create_table_insert_data_for_repair(manager)
    token = -1
    map0 = await load_tablet_sstables_repaired_at(manager, cql, servers[0], hosts[0], table_id)

    # Repair should not finish with error
    await inject_error_on(manager, "repair_tablet_fail_on_rpc_call", servers)
    try:
        await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental', timeout=10)
        assert False # Check the tablet repair is not supposed to finish
    except TimeoutError:
        logger.info("Repair timeout as expected")
    await inject_error_off(manager, "repair_tablet_fail_on_rpc_call", servers)

    map1 = await load_tablet_sstables_repaired_at(manager, cql, servers[0], hosts[0], table_id)

    # Check the sstables_repaired_at is not increased
    assert get_sstables_repaired_at(map0, token) == get_sstables_repaired_at(map1, token)

async def do_tablet_incremental_repair_and_ops(manager: ManagerClient, ops: str):
    nr_keys = 100
    servers, cql, hosts, ks, table_id, logs, repaired_keys, unrepaired_keys, current_key, token = await prepare_cluster_for_incremental_repair(manager, nr_keys, cmdline=['--logger-log-level', 'compaction=debug'])
    token = -1

    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')
    # 1 add 0 skip 1 mark
    for log in logs:
        sst_add, sst_skip, sst_mark = await get_sst_status("First", log)
        assert len(sst_add) == 1
        assert len(sst_skip) == 0
        assert len(sst_mark) == 1

    await insert_keys(cql, ks, current_key, current_key + nr_keys)
    current_key += nr_keys

    # Test ops does not mix up repaired and unrepaired sstables together
    for server in servers:
        if (ops == 'scrub_abort'):
            await manager.api.keyspace_scrub_sstables(server.ip_addr, ks, 'ABORT')
        elif (ops == 'scrub_validate'):
            await manager.api.keyspace_scrub_sstables(server.ip_addr, ks, 'VALIDATE')
        elif (ops == 'cleanup'):
            await manager.api.cleanup_keyspace(server.ip_addr, ks)
        elif (ops == 'upgradesstables'):
            await manager.api.keyspace_upgrade_sstables(server.ip_addr, ks)
        elif (ops == 'major'):
            await manager.api.keyspace_compaction(server.ip_addr, ks, 'test')
        else:
            assert False # Wrong ops

    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')

    # 1 add 1 skip 1 mark
    for log in logs:
        sst_add, sst_skip, sst_mark = await get_sst_status("Second", log)
        assert len(sst_add) == 2
        assert len(sst_mark) == 2
        assert len(sst_skip) == 1

@pytest.mark.asyncio
async def test_tablet_incremental_repair_and_scrubsstables_abort(manager: ManagerClient):
    await do_tablet_incremental_repair_and_ops(manager, 'scrub_abort')

@pytest.mark.asyncio
async def test_tablet_incremental_repair_and_scrubsstables_validate(manager: ManagerClient):
    await do_tablet_incremental_repair_and_ops(manager, 'scrub_validate')

@pytest.mark.asyncio
async def test_tablet_incremental_repair_and_cleanup(manager: ManagerClient):
    await do_tablet_incremental_repair_and_ops(manager, 'cleanup')

@pytest.mark.asyncio
async def test_tablet_incremental_repair_and_upgradesstables(manager: ManagerClient):
    await do_tablet_incremental_repair_and_ops(manager, 'upgradesstables')

@pytest.mark.asyncio
async def test_tablet_incremental_repair_and_major(manager: ManagerClient):
    await do_tablet_incremental_repair_and_ops(manager, 'major')

@pytest.mark.asyncio
async def test_tablet_incremental_repair_and_minor(manager: ManagerClient):
    nr_keys = 100
    servers, cql, hosts, ks, table_id, logs, repaired_keys, unrepaired_keys, current_key, token = await prepare_cluster_for_incremental_repair(manager, nr_keys)

    # Disable autocompaction
    for server in servers:
        await manager.api.disable_autocompaction(server.ip_addr, ks, 'test')

    # First repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')

    # Insert more keys
    await insert_keys(cql, ks, current_key, current_key + nr_keys)
    repaired_keys = repaired_keys | set(range(current_key, current_key + nr_keys))
    current_key += nr_keys

    # Second repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')

    # Insert more keys and flush to get 2 more sstables
    for _ in range(2):
        await insert_keys(cql, ks, current_key, current_key + nr_keys)
        unrepaired_keys = unrepaired_keys | set(range(current_key, current_key + nr_keys))
        current_key += nr_keys
        for server in servers:
            await manager.api.flush_keyspace(server.ip_addr, ks)

    # Enable autocompaction
    for server in servers:
        await manager.api.enable_autocompaction(server.ip_addr, ks, 'test')

    tm = TaskManagerClient(manager.api)
    for server in servers:
        module_name = "compaction"
        tasks = await tm.list_tasks(server.ip_addr, module_name)
        logger.info(f"Compaction {tasks=}")
        await tm.drain_module_tasks(server.ip_addr, module_name)

    # Verify repaired and unrepaired keys
    scylla_path = await manager.server_get_exe(servers[0].server_id)

    for server in servers:
        await manager.server_stop_gracefully(server.server_id)

    await verify_repaired_and_unrepaired_keys(manager, scylla_path, servers, ks, repaired_keys, unrepaired_keys)

async def do_test_tablet_incremental_repair_with_split_and_merge(manager, do_split, do_merge):
    nr_keys = 100
    servers, cql, hosts, ks, table_id, logs, repaired_keys, unrepaired_keys, current_key, token = await prepare_cluster_for_incremental_repair(manager, nr_keys)

    # First repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 1
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token) # sstables_repaired_at 2

    # Insert more keys
    await insert_keys(cql, ks, current_key, current_key + nr_keys)
    repaired_keys = repaired_keys | set(range(current_key, current_key + nr_keys))
    current_key += nr_keys

    # Second repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 3

    # Insert more keys and flush to get 2 more sstables
    for _ in range(2):
        await insert_keys(cql, ks, current_key, current_key + nr_keys)
        unrepaired_keys = unrepaired_keys | set(range(current_key, current_key + nr_keys))
        current_key += nr_keys
        for server in servers:
            await manager.api.flush_keyspace(server.ip_addr, ks)

    if do_split:
        s1_mark = await logs[0].mark()
        await inject_error_on(manager, "tablet_force_tablet_count_increase", servers)
        await logs[0].wait_for('Detected tablet split for table', from_mark=s1_mark)
        await inject_error_off(manager, "tablet_force_tablet_count_increase", servers)

    if do_merge:
        await trigger_tablet_merge(manager, servers, logs)

    scylla_path = await manager.server_get_exe(servers[0].server_id)

    await asyncio.sleep(random.randint(1, 5))

    for server in servers:
        await manager.server_stop_gracefully(server.server_id)

    await verify_repaired_and_unrepaired_keys(manager, scylla_path, servers, ks, repaired_keys, unrepaired_keys)

@pytest.mark.skip(reason="https://github.com/scylladb/scylladb/issues/24153")
@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_with_split_and_merge(manager: ManagerClient):
    await do_test_tablet_incremental_repair_with_split_and_merge(manager, do_split=True, do_merge=True)

@pytest.mark.skip(reason="https://github.com/scylladb/scylladb/issues/24153")
@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_with_split(manager: ManagerClient):
    await do_test_tablet_incremental_repair_with_split_and_merge(manager, do_split=True, do_merge=False)

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_with_merge(manager: ManagerClient):
    await do_test_tablet_incremental_repair_with_split_and_merge(manager, do_split=False, do_merge=True)

@pytest.mark.asyncio
async def test_tablet_incremental_repair_existing_and_repair_produced_sstable(manager: ManagerClient):
    nr_keys = 100
    cmdline = ["--hinted-handoff-enabled", "0"]
    servers, cql, hosts, ks, table_id, logs, repaired_keys, unrepaired_keys, current_key, token = await prepare_cluster_for_incremental_repair(manager, nr_keys, cmdline)

    await manager.server_stop_gracefully(servers[1].server_id)

    # Insert more keys
    await insert_keys(cql, ks, current_key, current_key + nr_keys)
    repaired_keys = repaired_keys | set(range(current_key, current_key + nr_keys))
    current_key += nr_keys

    await manager.server_start(servers[1].server_id)

    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')

    scylla_path = await manager.server_get_exe(servers[0].server_id)

    for server in servers:
        await manager.server_stop_gracefully(server.server_id)

    await verify_repaired_and_unrepaired_keys(manager, scylla_path, servers, ks, repaired_keys, unrepaired_keys)

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_merge_higher_repaired_at_number(manager):
    nr_keys = 100
    servers, cql, hosts, ks, table_id, logs, repaired_keys, unrepaired_keys, current_key, token = await prepare_cluster_for_incremental_repair(manager, nr_keys)

    # First repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 1
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 2

    # Insert more keys
    await insert_keys(cql, ks, current_key, current_key + nr_keys)
    unrepaired_keys = unrepaired_keys | set(range(current_key, current_key + nr_keys))
    current_key += nr_keys

    # Second repair
    await inject_error_on(manager, "repair_tablet_no_update_sstables_repair_at", servers)
    # some sstable repaired_at = 3, but sstables_repaired_at = 2
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 2
    await inject_error_off(manager, "repair_tablet_no_update_sstables_repair_at", servers)

    scylla_path = await manager.server_get_exe(servers[0].server_id)

    s1_mark = await logs[0].mark()
    await trigger_tablet_merge(manager, servers, logs)
    # The merge process will set the unrepaired sstable with repaired_at=3 to repaired_at=0 during merge
    await logs[0].wait_for('Updating repaired_at for tablet merge .* old=3 new=0 sstables_repaired_at=2', from_mark=s1_mark)
    await logs[0].wait_for('Completed updating repaired_at=.* for tablet merge', from_mark=s1_mark)

    for server in servers:
        await manager.server_stop_gracefully(server.server_id)

    # Verify max repaired_at of sstables should be 2
    max_repaired_at = 2
    await verify_max_repaired_at(manager, scylla_path, servers, ks, 2)

    for server in servers:
        await manager.server_stop_gracefully(server.server_id)

    await verify_repaired_and_unrepaired_keys(manager, scylla_path, servers, ks, repaired_keys, unrepaired_keys)

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_merge_correct_repaired_at_number_after_merge(manager):
    nr_keys = 100
    servers, cql, hosts, ks, table_id, logs, repaired_keys, unrepaired_keys, current_key, token = await prepare_cluster_for_incremental_repair(manager, nr_keys)

    # First repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 1
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 2

    # Insert more keys
    await insert_keys(cql, ks, current_key, current_key + nr_keys)
    unrepaired_keys = unrepaired_keys | set(range(current_key, current_key + nr_keys))
    current_key += nr_keys

    # Second repair
    replicas = await get_all_tablet_replicas(manager, servers[0], ks, 'test')
    last_tokens = [t.last_token for t in replicas]
    for t in last_tokens[0::2]:
        logging.info(f"Start repair for token={t}");
        await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", t, incremental_mode='incremental') # sstables_repaired_at 3

    scylla_path = await manager.server_get_exe(servers[0].server_id)

    # Trigger merge
    await trigger_tablet_merge(manager, servers, logs)

    # Verify sstables_repaired_at should be 3 after merge
    for server, host in zip(servers, hosts):
        token_repaired_at_map = await load_tablet_sstables_repaired_at(manager, cql, server, host, table_id)
        for last_token, sstables_repaired_at in token_repaired_at_map.items():
            logging.info(f"{last_token=} {sstables_repaired_at=}")
            assert sstables_repaired_at == 3

async def do_test_tablet_incremental_repair_merge_error(manager, error):
    nr_keys = 100
    # Make sure no data commit log replay after force server stop
    cmdline = ['--enable-commitlog', '0']
    servers, cql, hosts, ks, table_id, logs, repaired_keys, unrepaired_keys, current_key, token = await prepare_cluster_for_incremental_repair(manager, nr_keys, cmdline)

    # First repair
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental') # sstables_repaired_at 1

    # Insert more keys
    await insert_keys(cql, ks, current_key, current_key + nr_keys)
    unrepaired_keys = unrepaired_keys | set(range(current_key, current_key + nr_keys))
    current_key += nr_keys

    for server in servers:
        await manager.api.flush_keyspace(server.ip_addr, ks)

    scylla_path = await manager.server_get_exe(server.server_id)

    coord = await get_topology_coordinator(manager)
    coord_serv = await find_server_by_host_id(manager, servers, coord)
    coord_log = await manager.server_open_log(coord_serv.server_id)

    # Trigger merge and error in merge
    mark = await coord_log.mark()
    await inject_error_on(manager, error, [coord_serv])
    await inject_error_on(manager, "tablet_force_tablet_count_decrease", servers)
    await inject_error_on(manager, "tablet_force_tablet_count_decrease_once", servers)
    await coord_log.wait_for(f'Got {error}', from_mark=mark)
    await inject_error_off(manager, "tablet_force_tablet_count_decrease", servers)
    await manager.server_stop(coord_serv.server_id)
    await manager.server_start(coord_serv.server_id)

    for server in servers:
        await manager.server_stop_gracefully(server.server_id)

    await verify_repaired_and_unrepaired_keys(manager, scylla_path, servers, ks, repaired_keys, unrepaired_keys)

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_merge_error_in_merge_finalization(manager):
    await do_test_tablet_incremental_repair_merge_error(manager, 'handle_tablet_resize_finalization_for_merge_error')

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_merge_error_in_merge_completion_fiber(manager):
    await do_test_tablet_incremental_repair_merge_error(manager, 'merge_completion_fiber_error')

@pytest.mark.asyncio
async def test_tablet_repair_with_incremental_option(manager: ManagerClient):
    servers, cql, hosts, ks, table_id, logs, _, _, _, token = await prepare_cluster_for_incremental_repair(manager)
    token = -1

    sstables_repaired_at = 0

    async def do_repair_and_check(incremental_mode, expected_increment, expected_log_regex, check=None):
        nonlocal sstables_repaired_at
        mark = await logs[0].mark()
        if check:
            skip1 = get_incremental_repair_sst_skipped_bytes(servers[0])
            read1 = get_incremental_repair_sst_read_bytes(servers[0])
        await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode=incremental_mode)
        sstables_repaired_at += expected_increment
        m = await load_tablet_sstables_repaired_at(manager, cql, servers[0], hosts[0], table_id)
        assert get_sstables_repaired_at(m, token) == sstables_repaired_at
        await logs[0].wait_for(expected_log_regex, from_mark=mark)
        if check:
            skip2 = get_incremental_repair_sst_skipped_bytes(servers[0])
            read2 = get_incremental_repair_sst_read_bytes(servers[0])
            check(skip1, read1, skip2, read2)

    def check1(skip1, read1, skip2, read2):
        assert skip1 == 0
        assert read1 == 0
        assert skip2 == 0
        assert read2 > 0
    await do_repair_and_check(None, 1, rf'Starting tablet repair by API .* incremental_mode=incremental.*', check1)

    def check2(skip1, read1, skip2, read2):
        assert skip1 == skip2
        assert read1 == read2
    await do_repair_and_check('disabled', 0, rf'Starting tablet repair by API .* incremental_mode=disabled.*', check2)

    def check3(skip1, read1, skip2, read2):
        assert skip1 < skip2
        assert read1 == read2
    await do_repair_and_check('incremental', 1, rf'Starting tablet repair by API .* incremental_mode=incremental.*', check3)

    def check4(skip1, read1, skip2, read2):
        assert skip1 == skip2
        assert read1 < read2
    await do_repair_and_check('full', 1, rf'Starting tablet repair by API .* incremental_mode=full.*', check4)

@pytest.mark.asyncio
async def test_incremental_repair_tablet_time_metrics(manager: ManagerClient):
    servers, _, _, ks, _, _, _, _, _, token = await prepare_cluster_for_incremental_repair(manager)
    time1 = 0
    time2 = 0

    for s in servers:
        time1 += get_repair_tablet_time_ms(s)
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token, incremental_mode='incremental')
    for s in servers:
        time2 += get_repair_tablet_time_ms(s)

    assert time1 == 0
    assert time2 > 0

# Reproducer for https://github.com/scylladb/scylladb/issues/26346
@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_incremental_repair_finishes_when_tablet_skips_end_repair_stage(manager):
    servers = await manager.servers_add(3, auto_rack_dc="dc1")

    async with new_test_keyspace(manager, "WITH replication = {'class': 'NetworkTopologyStrategy', 'replication_factor': 3} AND tablets = {'initial': 1}") as ks:
        async with new_test_table(manager, ks, "pk int PRIMARY KEY, t text") as cf:
            table = cf.split('.')[-1]

            coord = await get_topology_coordinator(manager)
            coord_serv = await find_server_by_host_id(manager, servers, coord)
            coord_log = await manager.server_open_log(coord_serv.server_id)
            coord_mark = await coord_log.mark()

            await manager.api.enable_injection(coord_serv.ip_addr, "delay_end_repair_update", one_shot=False)
            response = await manager.api.tablet_repair(servers[0].ip_addr, ks, table, "all", await_completion=False, incremental_mode="incremental")
            task_id = response['tablet_task_id']

            await coord_log.wait_for("Finished tablet repair", from_mark=coord_mark)
            await trigger_stepdown(manager, coord_serv)

            # Disable injection in case, the same node is elected as coordinator
            await manager.api.disable_injection(coord_serv.ip_addr, "delay_end_repair_update")
            await manager.api.wait_task(servers[0].ip_addr, task_id)

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_incremental_repair_rejoin_do_tablet_operation(manager):
    cmdline = ['--logger-log-level', 'raft_topology=debug']
    servers = await manager.servers_add(3, auto_rack_dc="dc1", cmdline=cmdline)

    async with new_test_keyspace(manager, "WITH replication = {'class': 'NetworkTopologyStrategy', 'replication_factor': 3} AND tablets = {'initial': 1}") as ks:
        async with new_test_table(manager, ks, "pk int PRIMARY KEY, t text") as cf:
            table = cf.split('.')[-1]

            async def get_coord():
                coord = await get_topology_coordinator(manager)
                coord_serv = await find_server_by_host_id(manager, servers, coord)
                coord_log = await manager.server_open_log(coord_serv.server_id)
                return coord, coord_serv, coord_log

            for s in servers:
                await manager.api.enable_injection(s.ip_addr, "repair_finish_wait", one_shot=False)

            coord, coord_serv, coord_log = await get_coord()

            response = await manager.api.tablet_repair(servers[0].ip_addr, ks, table, "all", await_completion=False, incremental_mode="incremental")
            task_id = response['tablet_task_id']

            await coord_log.wait_for("Initiating tablet repair")
            await trigger_stepdown(manager, coord_serv)

            coord, coord_serv, coord_log = await get_coord()
            await coord_log.wait_for("Initiating tablet repair")

            found = False
            for run in range(10):
                for s in servers:
                    log = await manager.server_open_log(s.server_id)
                    res = await log.grep('Repair retry joining with existing session for tablet')
                    if len(res) > 0:
                        found = True
                        break
                if found:
                    break
                await asyncio.sleep(3)
            assert found

            for s in servers:
                await manager.api.message_injection(s.ip_addr, "repair_finish_wait")
                await manager.api.disable_injection(s.ip_addr, "repair_finish_wait")
            await coord_log.wait_for("Finished tablet repair")

@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_incremental_retry_end_repair_stage(manager):
    config = {'tablet_load_stats_refresh_interval_in_seconds': 1}
    servers = await manager.servers_add(3, auto_rack_dc="dc1", config=config)

    async with new_test_keyspace(manager, "WITH replication = {'class': 'NetworkTopologyStrategy', 'replication_factor': 3} AND tablets = {'initial': 1}") as ks:
        async with new_test_table(manager, ks, "pk int PRIMARY KEY, t text") as cf:
            table = cf.split('.')[-1]

            async def get_coord():
                coord = await get_topology_coordinator(manager)
                coord_serv = await find_server_by_host_id(manager, servers, coord)
                coord_log = await manager.server_open_log(coord_serv.server_id)
                return coord, coord_serv, coord_log

            for s in servers:
                await manager.api.enable_injection(s.ip_addr, "fail_rpc_repair_update_compaction_ctrl", one_shot=False)
                await manager.api.enable_injection(s.ip_addr, "log_tablet_transition_stage_end_repair", one_shot=False)

            coord, coord_serv, coord_log = await get_coord()

            response = await manager.api.tablet_repair(servers[0].ip_addr, ks, table, "all", await_completion=False, incremental_mode="incremental")
            task_id = response['tablet_task_id']

            await coord_log.wait_for("Finished tablet repair")
            await trigger_stepdown(manager, coord_serv)

            coord, coord_serv, coord_log = await get_coord()
            await coord_log.wait_for(r"Failed to perform repair_update_compaction_ctrl for tablet repair .* nr_retried=3")

            for s in servers:
                await manager.api.disable_injection(s.ip_addr, "fail_rpc_repair_update_compaction_ctrl")

            await coord_log.wait_for(r"The end_repair stage finished for tablet repair")

            try:
                await manager.api.wait_task(servers[0].ip_addr, task_id)
            except Exception as e:
                # Currently wait_task throws task not found in case the task_id has finished
                if "not found" not in str(e):
                    raise e

            # Run the second repair after the first is finished
            await manager.api.tablet_repair(servers[0].ip_addr, ks, table, "all", await_completion=True, incremental_mode="incremental")

# Reproducer for https://github.com/scylladb/scylladb/issues/27666.
# This test checks both tablet and vnode table. It tests the code path dealing
# with appending sstables produced by repair to a list work correctly with
# multishard writer when the shard count is different.
@pytest.mark.asyncio
@pytest.mark.parametrize("use_tablet", [False, True])
async def test_repair_sigsegv_with_diff_shard_count(manager: ManagerClient, use_tablet):
    cmdline0 = [ '--smp', '2']
    cmdline1 = [ '--smp', '3']
    servers = await manager.servers_add(1, cmdline=cmdline0, auto_rack_dc="dc1")
    servers.append(await manager.server_add(cmdline=cmdline1, property_file={'dc': 'dc1', 'rack': 'rack2'}))

    cql = manager.get_cql()

    async with new_test_keyspace(manager, f"WITH replication = {{'class': 'NetworkTopologyStrategy', 'replication_factor': 2}} AND TABLETS = {{ 'enabled': {str(use_tablet).lower()} }}  ") as ks:
        await cql.run_async(f"CREATE TABLE {ks}.test (pk int PRIMARY KEY, c int);")

        async def write_with_cl_one(range_start, range_end):
            insert_stmt = cql.prepare(f"INSERT INTO {ks}.test (pk, c) VALUES (?, ?)")
            insert_stmt.consistency_level = ConsistencyLevel.ONE
            pks = range(range_start, range_end)
            await asyncio.gather(*[cql.run_async(insert_stmt, (k, k)) for k in pks])

        logger.info("Adding data only on first node")
        await manager.api.flush_keyspace(servers[1].ip_addr, ks)
        await manager.server_stop(servers[1].server_id)
        manager.driver_close()
        cql = await reconnect_driver(manager)
        await write_with_cl_one(0, 10)
        await manager.api.flush_keyspace(servers[0].ip_addr, ks)

        logger.info("Adding data only on second node")
        await manager.server_start(servers[1].server_id)
        await manager.api.flush_keyspace(servers[0].ip_addr, ks)
        await manager.server_stop(servers[0].server_id)
        manager.driver_close()
        cql = await reconnect_driver(manager)
        await write_with_cl_one(10, 20)

        await manager.server_start(servers[0].server_id)
        await manager.servers_see_each_other(servers)

        if use_tablet:
            logger.info("Starting tablet repair")
            await manager.api.tablet_repair(servers[1].ip_addr, ks, "test", token='all', incremental_mode="incremental")
        else:
            logger.info("Starting vnode repair")
            await manager.api.repair(servers[1].ip_addr, ks, "test")

# Reproducer for https://github.com/scylladb/scylladb/issues/27365
# Incremental repair vs table drop
@pytest.mark.asyncio
@pytest.mark.skip_mode(mode='release', reason='error injections are not supported in release mode')
async def test_tablet_incremental_repair_table_drop_compaction_group_gone(manager: ManagerClient):
    cmdline = ['--logger-log-level', 'repair=debug']
    servers, cql, hosts, ks, table_id, logs, _, _, _, _ = await prepare_cluster_for_incremental_repair(manager, cmdline=cmdline)

    coord = await get_topology_coordinator(manager)
    coord_serv = await find_server_by_host_id(manager, servers, coord)
    coord_log = await manager.server_open_log(coord_serv.server_id)

    # Trigger merge and wait until the merge fiber starts
    s1_mark = await coord_log.mark()

    # Trigger repair and wait for the inc repair prepare preparation to start
    s1_mark = await coord_log.mark()
    await inject_error_on(manager, "wait_after_prepare_sstables_for_incremental_repair", servers)
    await manager.api.tablet_repair(servers[0].ip_addr, ks, "test", token=-1, await_completion=False, incremental_mode='incremental')
    # Wait for preparation to finish.
    await coord_log.wait_for('Re-enabled compaction for range', from_mark=s1_mark)

    s1_mark = await coord_log.mark()
    drop_future = cql.run_async(f"DROP TABLE {ks}.test;")
    await coord_log.wait_for(f'Stopping.*ongoing compactions for table {ks}.test', from_mark=s1_mark)
    await asyncio.sleep(0.2)

    # Continue the repair to trigger use-after-free
    for s in servers:
        await manager.api.message_injection(s.ip_addr, "wait_after_prepare_sstables_for_incremental_repair")

    await drop_future
