import json
import subprocess
from pathlib import Path

import pytest

from pingcap_datapipelines.config import PipelineConfig
from pingcap_datapipelines.extract import (
    chunk_tidb_export_patterns,
    extract_tidbcloud,
    plan_tidb_export_batches,
    run_extract,
)


def filter_payload_bytes(patterns: list[str]) -> int:
    payload = {"table": {"patterns": patterns, "where": "", "partitions": []}}
    return len(json.dumps(payload, separators=(",", ":")).encode())


def tidb_config(
    patterns: list[str],
    *,
    max_filter_bytes: int,
    fault_domains: list[dict[str, object]] | None = None,
) -> PipelineConfig:
    export = {"tidb_filter_max_bytes": max_filter_bytes, "parquet_compression": "ZSTD"}
    if fault_domains is not None:
        export["tidb_fault_domains"] = fault_domains
    return PipelineConfig(
        path=Path("source.yaml"),
        data={
            "source": {
                "name": "finance",
                "type": "tidbcloud",
                "database": "finance",
                "tidb_cloud_cli_profile": "finance",
                "tidb_cloud": {
                    "cluster_id": "cluster-1",
                    "gcs_service_account_key_env": "TEST_TIDB_GCS_KEY",
                },
            },
            "target": {
                "project": "gcp-digital-bi",
                "location": "us-west2",
                "datasets": {"raw": "raw", "final": "final", "control": "control"},
            },
            "tables": {"include": patterns},
            "export": export,
            "gcs": {
                "bucket": "bucket",
                "snapshot_prefix": "snapshots",
                "manifest_prefix": "manifests",
                "runtime_log_prefix": "logs",
            },
        },
    )


def test_chunk_tidb_export_patterns_preserves_order_within_byte_budget() -> None:
    patterns = [
        "finance.monthly_accounts_receivable",
        "finance.sf_revrec",
        "finance.marketplace_collection",
    ]
    max_filter_bytes = 100

    chunks = chunk_tidb_export_patterns(patterns, max_filter_bytes=max_filter_bytes)

    assert [pattern for chunk in chunks for pattern in chunk] == patterns
    assert len(chunks) > 1
    assert all(chunk for chunk in chunks)
    assert all(filter_payload_bytes(chunk) <= max_filter_bytes for chunk in chunks)


def test_chunk_tidb_export_patterns_rejects_one_oversized_pattern() -> None:
    pattern = f"finance.{'x' * 100}"

    with pytest.raises(ValueError, match="exceeds 100 bytes"):
        chunk_tidb_export_patterns([pattern], max_filter_bytes=100)


def test_chunk_tidb_export_patterns_rejects_later_oversized_pattern() -> None:
    pattern = f"finance.{'x' * 100}"

    with pytest.raises(ValueError, match="exceeds 100 bytes"):
        chunk_tidb_export_patterns(["finance.sf_revrec", pattern], max_filter_bytes=100)


def test_fault_domain_batches_preserve_every_pattern_exactly_once() -> None:
    patterns = [
        "nexuscrm.Account",
        "nexuscrm.Order",
        "nexuscrm._System_OutboxEvent",
        "nexuscrm._System_Log",
    ]
    fault_domains = [
        {"name": "business", "include": ["*"], "exclude": ["_System_*"]},
        {"name": "system", "include": ["_System_*"]},
    ]

    batches = plan_tidb_export_batches(
        patterns,
        fault_domains=fault_domains,
        max_filter_bytes=100,
    )

    exported = [pattern for batch in batches for pattern in batch["filter_patterns"]]
    assert sorted(exported) == sorted(patterns)
    assert len(exported) == len(set(exported))
    patterns_by_domain: dict[str, set[str]] = {}
    for batch in batches:
        patterns_by_domain.setdefault(str(batch["fault_domain"]), set()).update(batch["filter_patterns"])
    assert patterns_by_domain == {
        "business": {"nexuscrm.Account", "nexuscrm.Order"},
        "system": {"nexuscrm._System_OutboxEvent", "nexuscrm._System_Log"},
    }


@pytest.mark.parametrize(
    ("fault_domains", "match"),
    [
        (
            [
                {"name": "all", "include": ["*"]},
                {"name": "system", "include": ["_System_*"]},
            ],
            "must match exactly one fault domain",
        ),
        ([{"name": "system", "include": ["_System_*"]}], "must match exactly one fault domain"),
        (
            [
                {"name": "business", "include": ["*"], "exclude": ["_System_*"]},
                {"name": "system", "include": ["_System_*"]},
                {"name": "unused", "include": ["NeverPresent"]},
            ],
            "matched no patterns",
        ),
    ],
)
def test_fault_domain_batches_fail_closed_on_overlap_omission_or_empty_domain(
    fault_domains: list[dict[str, object]],
    match: str,
) -> None:
    with pytest.raises(ValueError, match=match):
        plan_tidb_export_batches(
            ["nexuscrm.Account", "nexuscrm._System_OutboxEvent"],
            fault_domains=fault_domains,
            max_filter_bytes=1000,
        )


def test_extract_tidbcloud_runs_every_filter_chunk_before_returning(monkeypatch: pytest.MonkeyPatch) -> None:
    config = tidb_config(
        ["monthly_accounts_receivable", "sf_revrec", "marketplace_collection"],
        max_filter_bytes=100,
    )
    expected_patterns = [
        "finance.monthly_accounts_receivable",
        "finance.sf_revrec",
        "finance.marketplace_collection",
    ]
    expected_chunks = chunk_tidb_export_patterns(expected_patterns, max_filter_bytes=100)
    commands: list[list[str]] = []
    descriptions: list[str] = []

    monkeypatch.setenv("TEST_TIDB_GCS_KEY", "base64-service-account")
    monkeypatch.setattr("pingcap_datapipelines.extract.ensure_ticloud_profile", lambda _: None)
    monkeypatch.setattr("pingcap_datapipelines.extract.ticloud_bin", lambda: "ticloud")

    def fake_run(
        command: list[str],
        *,
        description: str,
        sensitive_values: tuple[str, ...],
    ) -> subprocess.CompletedProcess[str]:
        commands.append(command)
        descriptions.append(description)
        assert sensitive_values == ("base64-service-account",)
        export_id = f"exp-{len(commands)}"
        return subprocess.CompletedProcess(command, 0, stdout=json.dumps({"exportId": export_id}), stderr="")

    monkeypatch.setattr("pingcap_datapipelines.extract.run_sensitive_command", fake_run)
    monkeypatch.setattr(
        "pingcap_datapipelines.extract.wait_for_ticloud_export",
        lambda _, export_id: {"exportId": export_id, "state": "SUCCEEDED"},
    )

    result = extract_tidbcloud(config, run_date="2026-07-23", run_id="run-1")

    assert len(commands) == len(expected_chunks)
    assert descriptions == [
        f"TiDB Cloud CreateExport chunk {index}/{len(expected_chunks)}" for index in range(1, len(expected_chunks) + 1)
    ]
    assert [command[command.index("--filter") + 1] for command in commands] == [
        ",".join(chunk) for chunk in expected_chunks
    ]
    assert [command[command.index("--gcs.uri") + 1] for command in commands] == [
        f"gs://bucket/snapshots/source=finance/run_date=2026-07-23/run_id=run-1/chunk={index:04d}/"
        for index in range(1, len(expected_chunks) + 1)
    ]
    assert result["export_count"] == len(expected_chunks)
    assert result["export_ids"] == [f"exp-{index}" for index in range(1, len(expected_chunks) + 1)]
    assert [chunk["filter_patterns"] for chunk in result["export_chunks"]] == expected_chunks


def test_run_extract_does_not_publish_manifest_when_a_later_chunk_fails(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    config = tidb_config(
        ["monthly_accounts_receivable", "sf_revrec", "marketplace_collection"],
        max_filter_bytes=100,
    )
    commands: list[list[str]] = []
    uploads: list[tuple[str, str]] = []
    local_writes: list[str] = []

    monkeypatch.setenv("TEST_TIDB_GCS_KEY", "base64-service-account")
    monkeypatch.setattr("pingcap_datapipelines.extract.ensure_ticloud_profile", lambda _: None)
    monkeypatch.setattr("pingcap_datapipelines.extract.ticloud_bin", lambda: "ticloud")

    def fake_run(command: list[str], **_: object) -> subprocess.CompletedProcess[str]:
        commands.append(command)
        export_id = f"exp-{len(commands)}"
        return subprocess.CompletedProcess(command, 0, stdout=json.dumps({"exportId": export_id}), stderr="")

    monkeypatch.setattr("pingcap_datapipelines.extract.run_sensitive_command", fake_run)
    monkeypatch.setattr(
        "pingcap_datapipelines.extract.wait_for_ticloud_export",
        lambda _, export_id: {
            "exportId": export_id,
            "state": "SUCCEEDED" if export_id == "exp-1" else "FAILED",
            "reason": "Export failed with unknown error" if export_id == "exp-2" else "",
        },
    )
    monkeypatch.setattr(
        "pingcap_datapipelines.extract.upload_text",
        lambda uri, text, **__: uploads.append((uri, text)),
    )
    monkeypatch.setattr(
        "pingcap_datapipelines.extract.write_local_run_file",
        lambda _, __, ___, filename, ____: local_writes.append(filename),
    )

    with pytest.raises(RuntimeError, match=r"chunk 2/3 .*exp-2.* state FAILED"):
        run_extract(config, run_date="2026-07-23", run_id="run-1")

    assert len(commands) == 2
    assert not any(uri.endswith("/manifest.json") for uri, _ in uploads)
    assert [uri for uri, _ in uploads] == [
        "gs://bucket/logs/source=finance/run_date=2026-07-23/run_id=run-1/events.jsonl"
    ]
    event = json.loads(uploads[0][1])
    assert event["event_type"] == "extract_failed"
    assert event["provider_reason"] == "Export failed with unknown error"
    assert event["provider_export"] == {"exportId": "exp-2", "state": "FAILED"}
    assert local_writes == ["events.jsonl"]


def test_run_extract_publishes_no_manifest_until_every_fault_domain_succeeds(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    config = tidb_config(
        ["Account", "_System_OutboxEvent"],
        max_filter_bytes=1000,
        fault_domains=[
            {"name": "business", "include": ["*"], "exclude": ["_System_*"]},
            {"name": "system", "include": ["_System_*"]},
        ],
    )
    commands: list[list[str]] = []
    uploads: list[str] = []

    monkeypatch.setenv("TEST_TIDB_GCS_KEY", "base64-service-account")
    monkeypatch.setattr("pingcap_datapipelines.extract.ensure_ticloud_profile", lambda _: None)
    monkeypatch.setattr("pingcap_datapipelines.extract.ticloud_bin", lambda: "ticloud")

    def fake_run(command: list[str], **_: object) -> subprocess.CompletedProcess[str]:
        commands.append(command)
        export_id = f"exp-{len(commands)}"
        return subprocess.CompletedProcess(command, 0, stdout=json.dumps({"exportId": export_id}), stderr="")

    monkeypatch.setattr("pingcap_datapipelines.extract.run_sensitive_command", fake_run)
    monkeypatch.setattr(
        "pingcap_datapipelines.extract.wait_for_ticloud_export",
        lambda _, export_id: {
            "exportId": export_id,
            "state": "SUCCEEDED" if export_id == "exp-1" else "FAILED",
            "errorMessage": "Export failed with unknown error",
        },
    )
    monkeypatch.setattr("pingcap_datapipelines.extract.upload_text", lambda uri, _, **__: uploads.append(uri))
    monkeypatch.setattr("pingcap_datapipelines.extract.write_local_run_file", lambda *_: None)

    with pytest.raises(RuntimeError, match="fault domain system.*provider reason"):
        run_extract(config, run_date="2026-09-02", run_id="run-1")

    assert len(commands) == 2
    assert commands[0][commands[0].index("--gcs.uri") + 1].endswith("fault_domain=business/chunk=0001/")
    assert commands[1][commands[1].index("--gcs.uri") + 1].endswith("fault_domain=system/chunk=0001/")
    assert not any(uri.endswith("/manifest.json") for uri in uploads)


def test_run_extract_publishes_manifest_after_every_fault_domain_succeeds(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    config = tidb_config(
        ["Account", "_System_OutboxEvent"],
        max_filter_bytes=1000,
        fault_domains=[
            {"name": "business", "include": ["*"], "exclude": ["_System_*"]},
            {"name": "system", "include": ["_System_*"]},
        ],
    )
    commands: list[list[str]] = []
    uploads: list[str] = []

    monkeypatch.setenv("TEST_TIDB_GCS_KEY", "base64-service-account")
    monkeypatch.setattr("pingcap_datapipelines.extract.ensure_ticloud_profile", lambda _: None)
    monkeypatch.setattr("pingcap_datapipelines.extract.ticloud_bin", lambda: "ticloud")

    def fake_run(command: list[str], **_: object) -> subprocess.CompletedProcess[str]:
        commands.append(command)
        return subprocess.CompletedProcess(
            command,
            0,
            stdout=json.dumps({"exportId": f"exp-{len(commands)}"}),
            stderr="",
        )

    monkeypatch.setattr("pingcap_datapipelines.extract.run_sensitive_command", fake_run)
    monkeypatch.setattr(
        "pingcap_datapipelines.extract.wait_for_ticloud_export",
        lambda _, export_id: {"exportId": export_id, "state": "SUCCEEDED"},
    )
    monkeypatch.setattr(
        "pingcap_datapipelines.extract.build_manifest",
        lambda *_args, **_kwargs: {"run_id": "run-1", "status": "exported"},
    )
    monkeypatch.setattr("pingcap_datapipelines.extract.upload_text", lambda uri, *_args, **_kwargs: uploads.append(uri))
    monkeypatch.setattr("pingcap_datapipelines.extract.write_local_run_file", lambda *_: None)

    result = run_extract(config, run_date="2026-09-02", run_id="run-1")

    assert len(commands) == 2
    assert result["manifest_uri"].endswith("/manifest.json")
    assert uploads == [result["manifest_uri"]]


@pytest.mark.parametrize(('states', 'attempt_count', 'succeeds'), [
    (['FAILED', 'SUCCEEDED'], 2, True),
    (['FAILED', 'FAILED', 'FAILED'], 3, False),
    (['CANCELED'], 1, False),
])
def test_failed_exports_retry_with_disjoint_prefixes_and_bounded_attempts(
    monkeypatch, states, attempt_count, succeeds,
):
    config = tidb_config(['sf_revrec'], max_filter_bytes=400)
    config.data['export']['tidb_max_attempts'] = 3
    monkeypatch.setenv('TEST_TIDB_GCS_KEY', 'fake-key')
    monkeypatch.setattr('pingcap_datapipelines.extract.ensure_ticloud_profile', lambda _: None)
    monkeypatch.setattr('pingcap_datapipelines.extract.ticloud_bin', lambda: 'ticloud')
    monkeypatch.setattr('pingcap_datapipelines.extract.time.sleep', lambda _: None)
    commands = []
    def run(command, **kwargs):
        commands.append(list(command))
        return subprocess.CompletedProcess(command, 0, stdout=json.dumps({'exportId': f'exp-{len(commands)}'}))
    monkeypatch.setattr('pingcap_datapipelines.extract.run_sensitive_command', run)
    monkeypatch.setattr('pingcap_datapipelines.extract.wait_for_ticloud_export',
                        lambda _, eid: {'exportId': eid, 'state': states[len(commands)-1]})
    if succeeds:
        result = extract_tidbcloud(config, run_date='2026-09-08', run_id='retry')
        assert result['export_ids'] == ['exp-2']
        assert result['export_chunks'][0]['gcs_uri'].endswith('attempt=0002/')
        assert len(result['export_chunks'][0]['failed_attempts']) == 1
    else:
        with pytest.raises(RuntimeError, match='ended in state'):
            extract_tidbcloud(config, run_date='2026-09-08', run_id='retry')
    assert len(commands) == attempt_count
    assert len({c[c.index('--gcs.uri')+1] for c in commands}) == attempt_count