import argparse
import hashlib
import json
import subprocess
from pathlib import Path

import pytest

from ops import restore_gcp_generated_postgres_snapshot as restore

SERVICE_STATE = {
    "ActiveState": "active",
    "SubState": "running",
    "MainPID": "148735",
    "NRestarts": "0",
    "ExecMainStartTimestamp": "Thu 2026-07-09 07:00:04 UTC",
}


def manifest_rows(database: str, *, target: bool, row_count: int = 1) -> list[dict]:
    identity = {
        "kind": "identity",
        "database": database,
        "server_version_num": 160014,
        "transaction_read_only": "on",
        "server_address": "10.61.0.3" if target else None,
        "server_port": 5432 if target else None,
        "ssl": target,
        "ssl_version": "TLSv1.3" if target else None,
        "database_collation": "en_US.UTF8" if target else "en_US.utf8",
        "database_ctype": "en_US.UTF8" if target else "en_US.utf8",
        "server_encoding": "UTF8",
    }
    rows = [
        identity,
        {"kind": "schemas", "items": ["kb_stage", "public"]},
        {"kind": "extensions", "items": [{"name": "pgcrypto", "version": "1.3"}]},
        {
            "kind": "application_roles",
            "items": [
                {
                    "name": "kb_apply",
                    "can_login": True,
                    "superuser": False,
                    "create_db": False,
                    "create_role": False,
                    "inherit": True,
                    "replication": False,
                    "bypass_rls": False,
                }
            ],
        },
    ]
    rows.extend(
        {"kind": kind, "items": []}
        for kind in (
            "columns",
            "constraints",
            "indexes",
            "sequences",
            "views",
            "functions",
            "triggers",
            "types",
            "policies",
        )
    )
    rows.extend(
        [
            {
                "kind": "table",
                "schema": "public",
                "table": "claims",
                "row_count": row_count,
                "rowset_md5": "same" if row_count == 1 else "different",
            },
            {
                "kind": "performance",
                "query": "count_claims",
                "elapsed_ms": 1.0,
                "observed_rows": row_count,
            },
        ]
    )
    return rows


def manifest_text(database: str, *, target: bool, row_count: int = 1) -> str:
    return "".join(json.dumps(row) + "\n" for row in manifest_rows(database, target=target, row_count=row_count))


def args_for(tmp_path: Path, **overrides) -> argparse.Namespace:
    dump = tmp_path / "teleo-canonical.dump"
    dump.write_bytes(b"PGDMPfixture")
    source_manifest = tmp_path / "source-manifest.jsonl"
    source_manifest.write_text(manifest_text("teleo", target=False), encoding="utf-8")
    values = {
        "operation": "restore",
        "execute": True,
        "target_db": "teleo_clone_dbfirst_test",
        "run_id": "gcp-restore-test12345",
        "host": restore.DEFAULT_HOST,
        "project": restore.DEFAULT_PROJECT,
        "password_secret": restore.DEFAULT_SECRET,
        "service": restore.DEFAULT_SERVICE,
        "command_timeout": 2.0,
        "psql_bin": "psql",
        "run_root": tmp_path / "runs",
        "dump": dump,
        "expected_dump_sha256": hashlib.sha256(dump.read_bytes()).hexdigest(),
        "source_manifest": source_manifest,
        "manifest_sql": Path("ops/postgres_parity_manifest.sql").resolve(),
        "pg_restore_bin": "pg_restore",
        "restore_timeout": 5.0,
        "manifest_timeout": 5.0,
        "max_target_query_ms": 2000.0,
        "max_target_source_ratio": 20.0,
    }
    values.update(overrides)
    return argparse.Namespace(**values)


def install_restore_fakes(
    monkeypatch: pytest.MonkeyPatch,
    *,
    target_manifest: str,
) -> dict[str, bool]:
    state = {"exists": False, "dropped": False}
    monkeypatch.setattr(restore.os, "geteuid", lambda: 0)
    monkeypatch.setattr(restore, "service_state", lambda *_args, **_kwargs: dict(SERVICE_STATE))
    monkeypatch.setattr(restore, "resolve_password", lambda *_args, **_kwargs: "private-password")
    monkeypatch.setattr(
        restore,
        "validate_pg_restore_version",
        lambda *_args, **_kwargs: {
            "pg_restore_bin": "pg_restore",
            "pg_restore_version": "pg_restore (PostgreSQL) 16.14",
            "pg_restore_major": 16,
            "source_postgres_major": 16,
            "compatible": True,
        },
    )
    monkeypatch.setattr(restore, "database_exists", lambda *_args, **_kwargs: state["exists"])

    def fake_create(*_args, **_kwargs):
        state["exists"] = True

    def fake_drop(*_args, **_kwargs):
        existed = state["exists"]
        state["exists"] = False
        state["dropped"] = True
        return {"existed_before": existed, "clone_database_remaining": 0}

    monkeypatch.setattr(restore, "create_database", fake_create)
    monkeypatch.setattr(restore, "restore_dump", lambda *_args, **_kwargs: None)
    monkeypatch.setattr(restore, "capture_target_manifest", lambda *_args, **_kwargs: target_manifest)
    monkeypatch.setattr(
        restore,
        "rollback_database_state",
        lambda *_args, **_kwargs: {"exists": True, "datallowconn": False, "connections": 0},
    )
    monkeypatch.setattr(restore, "drop_clone", fake_drop)
    return state


def test_restore_success_retains_only_exact_parity_clone(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    args = args_for(tmp_path)
    state = install_restore_fakes(
        monkeypatch,
        target_manifest=manifest_text(args.target_db, target=True),
    )

    result = restore.run_restore(args)

    assert result["status"] == "pass"
    assert result["phase"] == "retained_for_no_send_replay"
    assert result["parity"]["problems"] == []
    assert result["parity"]["details"]["source_total_rows"] == 1
    assert result["target"]["connectivity"]["private_connectivity"] is True
    assert result["target"]["connectivity"]["ssl"] is True
    assert result["live_service"]["unchanged"] is True
    assert result["cleanup"]["clone_database_remaining"] == 1
    assert state == {"exists": True, "dropped": False}
    run_dir = args.run_root / args.run_id
    assert (run_dir / "target-manifest.jsonl").stat().st_mode & 0o777 == 0o600
    assert (run_dir / "restore-receipt.json").stat().st_mode & 0o777 == 0o600


def test_parity_failure_drops_generated_clone_and_receipts_both_failures(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    args = args_for(tmp_path)
    state = install_restore_fakes(
        monkeypatch,
        target_manifest=manifest_text(args.target_db, target=True, row_count=2),
    )

    result = restore.run_restore(args)

    assert result["status"] == "fail"
    assert result["error"]["phase"] == "parity_compare"
    assert result["parity"]["problems"] == ["table_content_mismatches=1"]
    assert result["cleanup"]["attempted"] is True
    assert result["cleanup"]["clone_database_remaining"] == 0
    assert state == {"exists": False, "dropped": True}


def test_create_command_failure_still_drops_a_committed_clone(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    args = args_for(tmp_path)
    state = install_restore_fakes(
        monkeypatch,
        target_manifest=manifest_text(args.target_db, target=True),
    )

    def committed_then_failed(*_args, **_kwargs):
        state["exists"] = True
        raise restore.RestoreError("connection dropped after create committed")

    monkeypatch.setattr(restore, "create_database", committed_then_failed)

    result = restore.run_restore(args)

    assert result["status"] == "fail"
    assert result["error"]["phase"] == "database_create"
    assert result["cleanup"]["attempted"] is True
    assert result["cleanup"]["clone_database_remaining"] == 0
    assert state == {"exists": False, "dropped": True}


def test_changed_live_service_turns_successful_restore_into_cleanup(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    args = args_for(tmp_path)
    state = install_restore_fakes(
        monkeypatch,
        target_manifest=manifest_text(args.target_db, target=True),
    )
    calls = 0

    def changing_service(*_args, **_kwargs):
        nonlocal calls
        calls += 1
        result = dict(SERVICE_STATE)
        if calls > 1:
            result["NRestarts"] = "1"
        return result

    monkeypatch.setattr(restore, "service_state", changing_service)

    result = restore.run_restore(args)

    assert result["status"] == "fail"
    assert result["error"]["phase"] == "service_and_rollback_readback"
    assert result["live_service"]["unchanged"] is False
    assert result["cleanup"]["clone_database_remaining"] == 0
    assert state["dropped"] is True


def test_cleanup_is_bound_to_passing_restore_receipt(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    args = args_for(tmp_path)
    state = install_restore_fakes(
        monkeypatch,
        target_manifest=manifest_text(args.target_db, target=True),
    )
    assert restore.run_restore(args)["status"] == "pass"

    cleanup_args = argparse.Namespace(
        **{
            key: getattr(args, key)
            for key in (
                "target_db",
                "run_id",
                "host",
                "project",
                "password_secret",
                "service",
                "command_timeout",
                "psql_bin",
                "run_root",
            )
        },
        operation="cleanup",
        execute=True,
    )
    result = restore.run_cleanup(cleanup_args)

    assert result["status"] == "pass"
    assert result["database"] == {"existed_before": True, "clone_database_remaining": 0}
    assert result["rollback_database"] == {"exists": True, "datallowconn": False, "connections": 0}
    assert result["live_service"]["unchanged"] is True
    assert state == {"exists": False, "dropped": True}


def test_restore_dump_uses_tls_no_owner_no_acl_and_keeps_secret_out_of_argv(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    args = args_for(tmp_path)
    captured = {}

    def fake_run(command, *, timeout, env=None, input_text=None):
        captured.update(command=command, timeout=timeout, env=env, input_text=input_text)
        return subprocess.CompletedProcess(command, 0, "", "")

    monkeypatch.setattr(restore, "_run", fake_run)
    restore.restore_dump(args, "private-password")

    assert captured["command"][0] == "pg_restore"
    assert captured["command"][1] == "-d"
    assert "sslmode=require" in captured["command"][2]
    assert "--no-owner" in captured["command"]
    assert "--no-acl" in captured["command"]
    assert "--exit-on-error" in captured["command"]
    assert all("private-password" not in argument for argument in captured["command"])
    assert captured["env"]["PGPASSWORD"] == "private-password"


def test_pg_restore_preflight_rejects_client_older_than_snapshot(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    manifest_path = tmp_path / "source.jsonl"
    manifest_path.write_text(manifest_text("teleo", target=False), encoding="utf-8")
    source_manifest = restore.load_manifest(manifest_path)

    def fake_run(command, *, timeout, env=None, input_text=None):
        assert command == ["pg_restore", "--version"]
        assert timeout > 0
        assert env is None
        assert input_text is None
        return subprocess.CompletedProcess(command, 0, "pg_restore (PostgreSQL) 15.18\n", "")

    monkeypatch.setattr(restore, "_run", fake_run)
    with pytest.raises(restore.RestoreError, match="major 15 is older than source PostgreSQL major 16"):
        restore.validate_pg_restore_version("pg_restore", source_manifest, timeout=2.0)


@pytest.mark.parametrize("target", ["teleo_canonical", "teleo_clone_UPPER", "other_clone_test"])
def test_parse_args_rejects_non_disposable_target(target: str) -> None:
    with pytest.raises(SystemExit):
        restore.parse_args(
            [
                "cleanup",
                "--execute",
                "--target-db",
                target,
                "--run-id",
                "gcp-restore-test12345",
            ]
        )


def test_parse_args_requires_explicit_execute() -> None:
    with pytest.raises(SystemExit):
        restore.parse_args(
            [
                "cleanup",
                "--target-db",
                "teleo_clone_test",
                "--run-id",
                "gcp-restore-test12345",
            ]
        )
