# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright (c) 2022 ScyllaDB

import json
from pathlib import Path
from time import sleep, time

from cassandra.auth import PlainTextAuthProvider
from cassandra.cluster import Cluster
import pytest

from sdcm.gemini_thread import GeminiStressThread
from sdcm.utils.docker_utils import running_in_docker
from unit_tests.integration.conftest import configure_scylla_node
from unit_tests.lib.dummy_remote import LocalLoaderSetDummy
from unit_tests.lib.fake_cluster import DummyDbCluster

pytestmark = [
    pytest.mark.usefixtures("events"),
    pytest.mark.integration,
    pytest.mark.xdist_group("docker_heavy"),
]

GEMINI_KEYSPACE = "ks1"
GEMINI_TABLE = "table1"


@pytest.fixture(scope="session")
def gemini_schemas_dir(test_data_dir: Path) -> Path:
    """Return the path to unit_tests/test_data/gemini_schemas/."""
    return test_data_dir / "gemini_schemas"


@pytest.fixture(name="docker_scylla_oracle", scope="function")
def fixture_docker_scylla_oracle(params):
    """A standalone Scylla instance used as the oracle cluster in gemini tests.

    This is intentionally a fresh, independent container — not seeded from
    ``docker_scylla`` — so that gemini can treat it as a separate cluster.
    """
    scylla = configure_scylla_node({}, params)
    yield scylla
    scylla.kill()


@pytest.fixture(name="cql_session")
def fixture_cql_session(docker_scylla):
    """Yield an open CQL session connected to ``docker_scylla``, then shut it down."""
    if running_in_docker():
        address = f"{docker_scylla.internal_ip_address}:9042"
    else:
        address = docker_scylla.get_port("9042")
    node_ip, port = address.split(":")
    cluster = Cluster(
        [node_ip],
        port=int(port),
        auth_provider=PlainTextAuthProvider(username="cassandra", password="cassandra"),
    )
    session = cluster.connect()
    try:
        yield session
    finally:
        session.shutdown()
        cluster.shutdown()


def load_statement_ratios(gemini_schemas_dir: Path) -> str:
    ratios_path = gemini_schemas_dir / "no_delete_statement_ratios.json"
    return json.dumps(json.loads(ratios_path.read_text(encoding="utf-8")), separators=(",", ":"))


@pytest.fixture(name="gemini_thread")
def fixture_gemini_thread(request, params, docker_scylla, docker_scylla_oracle, gemini_schemas_dir):
    """Build and teardown a GeminiStressThread for the standard oracle case.

    Test functions that need to vary mode, duration, or other options should
    call :func:`build_gemini_thread` directly instead.
    """
    thread = build_gemini_thread(params, docker_scylla, docker_scylla_oracle, gemini_schemas_dir=gemini_schemas_dir)
    request.addfinalizer(thread.kill)
    return thread


@pytest.fixture(name="gemini_thread_no_oracle")
def fixture_gemini_thread_no_oracle(request, params, docker_scylla, gemini_schemas_dir):
    """Build and teardown a GeminiStressThread with no oracle cluster."""
    thread = build_gemini_thread(params, docker_scylla, oracle_node=None, gemini_schemas_dir=gemini_schemas_dir)
    request.addfinalizer(thread.kill)
    return thread


def build_gemini_thread(
    params,
    test_node,
    oracle_node=None,
    *,
    gemini_schemas_dir: Path,
    mode="mixed",
    duration="1m",
    extra_options=None,
    use_schema=False,
):
    """Construct a :class:`GeminiStressThread` for integration tests.

    Args:
        params: SCT params dict (from the ``params`` fixture).
        test_node: The primary Scylla Docker node to test against.
        oracle_node: Optional oracle Scylla Docker node; ``None`` disables oracle.
        gemini_schemas_dir: Path to the gemini_schemas test data directory.
        mode: Gemini run mode (``"mixed"``, ``"write"``, etc.).
        duration: Gemini run duration string (e.g. ``"30s"``).
        extra_options: Additional CLI flags to append to the stress command.
        use_schema: When ``True``, sets ``gemini_schema_url`` to the bundled
            two-partition-key schema fixture file.

    Returns:
        GeminiStressThread: Configured but not yet started thread.
    """
    # params is a function-scoped fixture — each test gets its own fresh SCTConfiguration,
    # so mutating it directly is safe and avoids the deprecated pydantic .copy() which can
    # drop dynamically-merged state (e.g. stress_image set by load_docker_images_defaults).
    params.update({"gemini_table_options": ["gc_grace_seconds=60"]})
    if use_schema:
        params["gemini_schema_url"] = str((gemini_schemas_dir / "simple_two_partition_keys_schema.json").resolve())

    loader_set = LocalLoaderSetDummy(params=params)
    test_cluster = DummyDbCluster([test_node])
    oracle_cluster = DummyDbCluster([oracle_node]) if oracle_node else None
    options = [
        f"--duration={duration}",
        "--warmup=0",
        "--concurrency=5",
        f"--mode={mode}",
        "--cql-features=normal",
        "--max-mutation-retries=100",
        "--max-mutation-retries-backoff=1s",
        "--replication-strategy=\"{'class': 'SimpleStrategy', 'replication_factor': '1'}\"",
        "--oracle-replication-strategy=\"{'class': 'SimpleStrategy', 'replication_factor': '1'}\"",
        # Large partition count reduces the probability that concurrent reader/writer goroutines
        # collide on the same partition.  With --concurrency=5 (5R+5W) and only 100 partitions
        # the collision rate was ~5% per write, making write-timeout-induced divergence nearly
        # certain within a 30-second run.  10 000 partitions drops that rate to ~0.05%.
        "--partition-count=10000",
        "--max-errors-to-store=1",
        # Override the aggressive SCT default (3 s) with a more generous value.  On a local
        # Docker host the latency is sub-millisecond, but under CPU/memory pressure a 3-second
        # deadline causes write timeouts.  A timed-out write that actually committed on the test
        # cluster but not the oracle (or vice versa) is the primary source of "different
        # clustering-key values" divergence errors because gemini retries with new data.
        "--request-timeout=10s",
        f"--statement-ratios='{load_statement_ratios(gemini_schemas_dir)}'",
    ]
    if extra_options:
        options.extend(extra_options)
    return GeminiStressThread(
        loaders=loader_set,
        stress_cmd=" ".join(options),
        test_cluster=test_cluster,
        oracle_cluster=oracle_cluster,
        timeout=240,
        params=params,
    )


def wait_for_table(session, keyspace=GEMINI_KEYSPACE, table=GEMINI_TABLE, timeout=60):
    """Poll until the given table exists in ``system_schema.tables``.

    Args:
        session: Open CQL session.
        keyspace: Target keyspace name.
        table: Target table name.
        timeout: Maximum seconds to wait before raising.

    Raises:
        AssertionError: If the table is not visible within *timeout* seconds.
    """
    deadline = time() + timeout
    while time() < deadline:
        row = session.execute(
            "SELECT table_name FROM system_schema.tables WHERE keyspace_name=%s AND table_name=%s",
            (keyspace, table),
        ).one()
        if row:
            return
        session.cluster.refresh_schema_metadata()
        sleep(1)
    raise AssertionError(f"Gemini table {keyspace}.{table} was not created in time")


def get_table_metadata(session, keyspace=GEMINI_KEYSPACE, table=GEMINI_TABLE):
    """Return the Cassandra driver table metadata for *keyspace.table*.

    Args:
        session: Open CQL session.
        keyspace: Target keyspace name.
        table: Target table name.

    Returns:
        cassandra.metadata.TableMetadata: Live table metadata object.
    """
    session.cluster.refresh_schema_metadata()
    return session.cluster.metadata.keyspaces[keyspace].tables[table]


def assert_result_structure(result):
    """Assert that a gemini result dict has the expected structure and sane values.

    Args:
        result: Parsed gemini result dict (from ``get_gemini_results``).
    """
    assert result is not None, "Gemini result must not be None"

    for key in ("write_ops", "read_ops", "write_errors", "read_errors"):
        assert key in result, f"Result must contain {key}"

    assert result["write_ops"] >= 0, "write_ops must be non-negative"
    assert result["read_ops"] >= 0, "read_ops must be non-negative"
    assert result["write_errors"] == 0, f"Expected no write errors, got: {result['write_errors']}"
    assert result["read_errors"] == 0, f"Expected no read errors, got: {result['read_errors']}"

    errors = result.get("errors")
    assert not errors, f"Expected no validation errors, got: {errors}"


def assert_write_result(result):
    """Assert write-mode result: write_ops must be positive, read_ops can be zero.

    Args:
        result: Parsed gemini result dict.
    """
    assert_result_structure(result)
    assert result["write_ops"] > 0, "Write run produced zero write_ops — did gemini actually run?"


def assert_mixed_result(result):
    """Assert mixed-mode result: both write_ops and read_ops must be positive.

    Args:
        result: Parsed gemini result dict.
    """
    assert_result_structure(result)
    assert result["write_ops"] > 0, "Mixed run produced zero write_ops — did gemini actually run?"
    assert result["read_ops"] > 0, "Mixed run produced zero read_ops — did gemini actually run?"


def _format_debug_report(thread: GeminiStressThread) -> str:
    """Read all gemini output files and return a formatted debug report.

    Reads the result JSON, summary, test-statement and oracle-statement log
    files that were written by the loader.  Designed to be called *after*
    :meth:`get_gemini_results` — the underlying ``get_results()`` call is
    safe to invoke again because completed futures return cached values.

    Args:
        thread: A finished :class:`GeminiStressThread` instance.

    Returns:
        Multi-section string, one section per loader, ready to embed in a
        ``pytest.fail()`` message.
    """
    raw_results = thread.get_results()
    sections = []

    for loader_idx, (_docker, _result, local_result_file) in enumerate(raw_results):
        logdir = Path(local_result_file).parent
        section = [
            f"\n{'=' * 70}",
            f"Loader #{loader_idx}  logdir={logdir}",
            f"{'=' * 70}",
        ]

        # ── Result JSON ───────────────────────────────────────────────────────
        section.append("\n[Result File]")
        try:
            content = Path(local_result_file).read_text(encoding="utf-8")
            try:
                section.append(json.dumps(json.loads(content), indent=2))
            except json.JSONDecodeError:
                section.append(content or "(empty)")
        except FileNotFoundError:
            section.append(f"(not found: {local_result_file})")

        # ── Summary ───────────────────────────────────────────────────────────
        summary_path = logdir / Path(thread.gemini_summary_file).name
        section.append("\n[Summary File]")
        try:
            section.append(summary_path.read_text(encoding="utf-8") or "(empty)")
        except FileNotFoundError:
            section.append(f"(not found: {summary_path})")

        # ── Test statement log ────────────────────────────────────────────────
        test_stmts_path = logdir / Path(thread.gemini_test_statements_file).name
        section.append("\n[Test Statement Log]")
        try:
            section.append(test_stmts_path.read_text(encoding="utf-8") or "(empty)")
        except FileNotFoundError:
            section.append(f"(not found: {test_stmts_path})")

        # ── Oracle statement log ──────────────────────────────────────────────
        oracle_stmts_path = logdir / Path(thread.gemini_oracle_statements_file).name
        section.append("\n[Oracle Statement Log]")
        try:
            section.append(oracle_stmts_path.read_text(encoding="utf-8") or "(empty)")
        except FileNotFoundError:
            section.append(f"(not found: {oracle_stmts_path})")

        sections.append("\n".join(section))

    return "\n".join(sections) if sections else "(no loader results found)"


def test_gemini_mixed_mode_with_oracle(request, docker_scylla, docker_scylla_oracle, params, gemini_schemas_dir):
    """A full mixed-mode run against an oracle cluster must complete without errors.

    This is the most representative integration scenario: gemini writes data and
    validates it against the oracle cluster simultaneously. Both write_ops and
    read_ops must be positive, and the result must carry no errors.
    """
    thread = build_gemini_thread(
        params, docker_scylla, docker_scylla_oracle, gemini_schemas_dir=gemini_schemas_dir, mode="mixed", duration="30s"
    )
    request.addfinalizer(thread.kill)
    thread.run()
    results = thread.get_gemini_results()
    stats = thread.verify_gemini_results(results)

    def _fail(msg: str) -> None:
        """Fail the test and append a full debug report (files, logs) to *msg*."""
        pytest.fail(f"{msg}\n{_format_debug_report(thread)}")

    if stats["status"] != "PASSED":
        _fail(f"Gemini mixed-mode run failed\n  status : {stats['status']}\n  errors : {stats.get('errors', {})}")

    if len(results) != 1:
        _fail(f"Expected exactly 1 result from a single-loader run, got {len(results)}")

    result = results[0]
    issues = []

    for field in ("write_ops", "read_ops", "write_errors", "read_errors"):
        if field not in result:
            issues.append(f"missing required field '{field}'")

    if result.get("write_ops", 0) <= 0:
        issues.append(f"write_ops={result.get('write_ops')} — did gemini actually run?")
    if result.get("read_ops", 0) <= 0:
        issues.append(f"read_ops={result.get('read_ops')} — did gemini actually run?")
    if result.get("write_errors"):
        issues.append(f"write_errors={result['write_errors']}")
    if result.get("read_errors"):
        issues.append(f"read_errors={result['read_errors']}")
    errors = result.get("errors")
    if errors:
        issues.append(f"{len(errors)} validation error(s) found in 'errors' field")

    if issues:
        _fail("Gemini mixed-mode result has issues:\n  " + "\n  ".join(issues))


def test_gemini_write_only_with_oracle(request, docker_scylla, docker_scylla_oracle, params, gemini_schemas_dir):
    """A write-only run against an oracle cluster must succeed and produce non-zero write_ops.

    Write-only mode is used for CDC tests and write-only workloads where reads are not
    performed during the run. An oracle cluster is still present so that a subsequent
    read-mode run could validate consistency if needed.
    """
    thread = build_gemini_thread(
        params, docker_scylla, docker_scylla_oracle, gemini_schemas_dir=gemini_schemas_dir, mode="write", duration="30s"
    )
    request.addfinalizer(thread.kill)
    thread.run()
    results = thread.get_gemini_results()
    stats = thread.verify_gemini_results(results)

    assert stats["status"] == "PASSED", f"Gemini write-only run failed: {stats}"
    assert len(results) == 1, "Expected exactly one result from a single-loader run"
    assert_write_result(results[0])


def test_gemini_with_schema(request, docker_scylla, docker_scylla_oracle, cql_session, params, gemini_schemas_dir):
    """A mixed-mode run using an explicit schema file must succeed and produce non-zero ops.

    Using a schema file gives deterministic table structure (known partition key names,
    column types, etc.) which enables targeted CQL assertions after the run. The schema
    used here defines a two-partition-key table (pk1: int, pk2: varchar) with two
    regular columns, so we can verify that the keyspace and table were created as expected.
    """
    thread = build_gemini_thread(
        params,
        docker_scylla,
        docker_scylla_oracle,
        gemini_schemas_dir=gemini_schemas_dir,
        mode="mixed",
        duration="30s",
        use_schema=True,
    )
    request.addfinalizer(thread.kill)

    # Verify the schema path is correctly injected into the gemini command
    gemini_cmd = thread.generate_gemini_command(schema_path="/tmp/schema.json")
    assert '--schema="/tmp/schema.json"' in gemini_cmd, "Schema path must appear in the generated command"
    assert "--partition-count=10000" in gemini_cmd, "partition-count must be forwarded from stress_cmd"
    assert "--duration=30s" in gemini_cmd, "duration must appear in the generated command"
    assert "--warmup=0" in gemini_cmd, "warmup must appear in the generated command"
    assert f"--statement-ratios='{load_statement_ratios(gemini_schemas_dir)}'" in gemini_cmd, (
        "statement-ratios must appear in the generated command"
    )

    thread.run()
    results = thread.get_gemini_results()
    stats = thread.verify_gemini_results(results)

    assert stats["status"] == "PASSED", f"Gemini schema run failed: {stats}"
    assert len(results) == 1, "Expected exactly one result from a single-loader run"
    assert_mixed_result(results[0])

    # Verify the schema was actually applied: the expected keyspace and table must exist
    wait_for_table(cql_session, keyspace=GEMINI_KEYSPACE, table=GEMINI_TABLE)
    metadata = get_table_metadata(cql_session, keyspace=GEMINI_KEYSPACE, table=GEMINI_TABLE)

    # The schema defines two partition keys: pk1 (int) and pk2 (varchar)
    pk_names = [col.name for col in metadata.partition_key]
    assert pk_names == ["pk1", "pk2"], f"Expected partition keys [pk1, pk2] from schema file, got {pk_names}"

    # The schema defines two regular columns: col1 (int) and col2 (varchar)
    regular_col_names = sorted(
        col_name
        for col_name, col in metadata.columns.items()
        if col_name not in {pk.name for pk in metadata.partition_key}
        and col_name not in {ck.name for ck in metadata.clustering_key}
    )
    assert regular_col_names == ["col1", "col2"], (
        f"Expected regular columns [col1, col2] from schema file, got {regular_col_names}"
    )


def test_gemini_without_oracle(request, docker_scylla, params, gemini_schemas_dir):
    """A write-only run with no oracle cluster must succeed.

    The oracle-less code path (``oracle_cluster is None``) is used when no
    consistency verification is needed, e.g. pure write-load or smoke tests.
    No ``--oracle-cluster`` flag must appear in the generated command.
    """
    thread = build_gemini_thread(
        params, docker_scylla, oracle_node=None, gemini_schemas_dir=gemini_schemas_dir, mode="write", duration="30s"
    )
    request.addfinalizer(thread.kill)

    # Confirm --oracle-cluster is absent from the generated command
    gemini_cmd = thread.generate_gemini_command()
    assert "--oracle-cluster" not in gemini_cmd, "oracle-cluster flag must not appear when oracle_cluster is None"

    thread.run()
    results = thread.get_gemini_results()
    stats = thread.verify_gemini_results(results)

    assert stats["status"] == "PASSED", f"Gemini no-oracle run failed: {stats}"
    assert len(results) == 1, "Expected exactly one result from a single-loader run"
    assert_write_result(results[0])
