#!/usr/bin/env python3
"""Plan or apply the Teleo GCP runtime baseline.

Default mode is dry-run. Execute mode is intentionally explicit and proof
oriented: it records every non-secret operation, redacts password-bearing
commands, requires a /32 admin SSH CIDR, and stops before Cloud SQL password
work unless the password is present in the configured environment variable.
"""

from __future__ import annotations

import argparse
import ipaddress
import json
import os
import subprocess
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path

PROJECT = "teleo-501523"
REGION = "europe-west6"
ZONE = "europe-west6-a"
NETWORK = "teleo-staging-net"
SUBNET = "teleo-staging-subnet"
SUBNET_RANGE = "10.70.0.0/24"
PRIVATE_RANGE = "teleo-cloudsql-private-range"
SNAPSHOT_POLICY = "teleo-daily-7d-snapshots"
BACKUP_BUCKETS = ("teleo-501523-prod-backups", "teleo-501523-leoclean-backups")
CLOUDSQL_INSTANCE = "teleo-pgvector-standby"
CLOUDSQL_DATABASE = "teleo_kb"
CLOUDSQL_PASSWORD_SECRET = "gcp-teleo-pgvector-standby-postgres-password"

SERVICE_ACCOUNTS = {
    "sa-teleo-cloudbuild": {
        "email": f"sa-teleo-cloudbuild@{PROJECT}.iam.gserviceaccount.com",
        "display_name": "Teleo Cloud Build image publisher",
        "roles": ("roles/artifactregistry.writer", "roles/logging.logWriter", "roles/storage.objectViewer"),
    },
    "sa-teleo-prod-vm": {
        "email": f"sa-teleo-prod-vm@{PROJECT}.iam.gserviceaccount.com",
        "display_name": "Teleo production VM runtime",
        "roles": ("roles/artifactregistry.reader", "roles/logging.logWriter", "roles/monitoring.metricWriter"),
    },
    "sa-teleo-staging-vm": {
        "email": f"sa-teleo-staging-vm@{PROJECT}.iam.gserviceaccount.com",
        "display_name": "Teleo staging VM runtime",
        "roles": ("roles/artifactregistry.reader", "roles/logging.logWriter", "roles/monitoring.metricWriter"),
    },
}

VMS = {
    "teleo-prod-1": {
        "service_account": SERVICE_ACCOUNTS["sa-teleo-prod-vm"]["email"],
        "tag": "teleo-prod-ssh",
        "machine_type": "e2-standard-2",
        "boot_disk_size": "80GB",
    },
    "teleo-staging-1": {
        "service_account": SERVICE_ACCOUNTS["sa-teleo-staging-vm"]["email"],
        "tag": "teleo-staging-ssh",
        "machine_type": "e2-standard-2",
        "boot_disk_size": "80GB",
    },
}

SSH_FIREWALL_RULES = {
    "teleo-prod-allow-ssh-current-ip": "teleo-prod-ssh",
    "teleo-staging-allow-ssh-current-ip": "teleo-staging-ssh",
}

REQUIRED_APIS = (
    "artifactregistry.googleapis.com",
    "cloudbuild.googleapis.com",
    "compute.googleapis.com",
    "logging.googleapis.com",
    "monitoring.googleapis.com",
    "secretmanager.googleapis.com",
    "servicenetworking.googleapis.com",
    "sqladmin.googleapis.com",
)


@dataclass
class OperationResult:
    name: str
    command: list[str]
    action: str
    exitcode: int | None = None
    stdout_tail: str = ""
    stderr_tail: str = ""


def tail(value: str, limit: int = 2000) -> str:
    return value[-limit:]


def run_command(
    command: list[str],
    *,
    execute: bool,
    name: str,
    action: str,
    recorded_command: list[str] | None = None,
    stdin: str | None = None,
) -> OperationResult:
    safe_command = recorded_command or command
    if not execute:
        return OperationResult(name=name, command=safe_command, action=action)
    completed = subprocess.run(command, input=stdin, text=True, capture_output=True, check=False)
    return OperationResult(
        name=name,
        command=safe_command,
        action=action,
        exitcode=completed.returncode,
        stdout_tail=tail(completed.stdout),
        stderr_tail=tail(completed.stderr),
    )


def is_reauth_failure(result: OperationResult) -> bool:
    return "Reauthentication failed" in result.stderr_tail or "cannot prompt during non-interactive execution" in result.stderr_tail


def validate_admin_ssh_cidr(value: str | None, *, execute: bool) -> str:
    if not value:
        if execute:
            raise ValueError("--admin-ssh-cidr is required in execute mode")
        return "<operator-ip>/32"
    network = ipaddress.ip_network(value, strict=False)
    if network.version != 4 or network.prefixlen != 32:
        raise ValueError("--admin-ssh-cidr must be a single IPv4 /32")
    return str(network)


def create_if_missing(
    *,
    probe: list[str],
    create: list[str],
    execute: bool,
    name: str,
    recorded_create: list[str] | None = None,
) -> list[OperationResult]:
    if execute:
        probe_result = run_command(probe, execute=True, name=f"probe:{name}", action="probe")
        if probe_result.exitcode == 0:
            probe_result.action = "already_exists"
            return [probe_result]
        if is_reauth_failure(probe_result):
            return [probe_result]
    return [
        run_command(
            create,
            execute=execute,
            name=name,
            action="create_if_missing",
            recorded_command=recorded_create,
        )
    ]


def ensure_service_account(account_id: str, email: str, display_name: str, *, execute: bool) -> list[OperationResult]:
    return create_if_missing(
        probe=["gcloud", "iam", "service-accounts", "describe", email, "--project", PROJECT, "--format=json"],
        create=[
            "gcloud",
            "iam",
            "service-accounts",
            "create",
            account_id,
            "--project",
            PROJECT,
            "--display-name",
            display_name,
        ],
        execute=execute,
        name=f"create_service_account:{email}",
    )


def ensure_project_binding(email: str, role: str, *, execute: bool, prefix: str) -> OperationResult:
    return run_command(
        [
            "gcloud",
            "projects",
            "add-iam-policy-binding",
            PROJECT,
            "--member",
            f"serviceAccount:{email}",
            "--role",
            role,
            "--quiet",
        ],
        execute=execute,
        name=f"{prefix}:{role}",
        action="ensure_binding",
    )


def ensure_network(*, execute: bool) -> list[OperationResult]:
    return create_if_missing(
        probe=["gcloud", "compute", "networks", "describe", NETWORK, "--project", PROJECT, "--format=json"],
        create=[
            "gcloud",
            "compute",
            "networks",
            "create",
            NETWORK,
            "--project",
            PROJECT,
            "--subnet-mode=custom",
            "--bgp-routing-mode=regional",
        ],
        execute=execute,
        name=f"create_network:{NETWORK}",
    )


def ensure_subnet(*, execute: bool) -> list[OperationResult]:
    operations = create_if_missing(
        probe=[
            "gcloud",
            "compute",
            "networks",
            "subnets",
            "describe",
            SUBNET,
            "--region",
            REGION,
            "--project",
            PROJECT,
            "--format=json",
        ],
        create=[
            "gcloud",
            "compute",
            "networks",
            "subnets",
            "create",
            SUBNET,
            "--project",
            PROJECT,
            "--network",
            NETWORK,
            "--region",
            REGION,
            "--range",
            SUBNET_RANGE,
            "--enable-private-ip-google-access",
        ],
        execute=execute,
        name=f"create_subnet:{SUBNET}",
    )
    operations.append(
        run_command(
            [
                "gcloud",
                "compute",
                "networks",
                "subnets",
                "update",
                SUBNET,
                "--project",
                PROJECT,
                "--region",
                REGION,
                "--enable-private-ip-google-access",
            ],
            execute=execute,
            name=f"enable_private_google_access:{SUBNET}",
            action="ensure_setting",
        )
    )
    return operations


def ensure_private_services_access(*, execute: bool) -> list[OperationResult]:
    operations = create_if_missing(
        probe=[
            "gcloud",
            "compute",
            "addresses",
            "describe",
            PRIVATE_RANGE,
            "--project",
            PROJECT,
            "--global",
            "--format=json",
        ],
        create=[
            "gcloud",
            "compute",
            "addresses",
            "create",
            PRIVATE_RANGE,
            "--project",
            PROJECT,
            "--global",
            "--prefix-length=16",
            "--purpose=VPC_PEERING",
            "--network",
            NETWORK,
        ],
        execute=execute,
        name=f"create_private_service_range:{PRIVATE_RANGE}",
    )
    operations.append(
        run_command(
            [
                "gcloud",
                "services",
                "vpc-peerings",
                "connect",
                "--project",
                PROJECT,
                "--service=servicenetworking.googleapis.com",
                "--network",
                NETWORK,
                "--ranges",
                PRIVATE_RANGE,
            ],
            execute=execute,
            name=f"ensure_private_services_access:{NETWORK}",
            action="ensure_vpc_peering",
        )
    )
    return operations


def ensure_firewall_rule(name: str, target_tag: str, admin_ssh_cidr: str, *, execute: bool) -> list[OperationResult]:
    return [
        *create_if_missing(
            probe=["gcloud", "compute", "firewall-rules", "describe", name, "--project", PROJECT, "--format=json"],
            create=[
                "gcloud",
                "compute",
                "firewall-rules",
                "create",
                name,
                "--project",
                PROJECT,
                "--network",
                NETWORK,
                "--direction=INGRESS",
                "--priority=1000",
                "--allow=tcp:22",
                "--source-ranges",
                admin_ssh_cidr,
                "--target-tags",
                target_tag,
                "--enable-logging",
            ],
            execute=execute,
            name=f"create_firewall_rule:{name}",
        ),
        run_command(
            [
                "gcloud",
                "compute",
                "firewall-rules",
                "update",
                name,
                "--project",
                PROJECT,
                "--source-ranges",
                admin_ssh_cidr,
                "--target-tags",
                target_tag,
                "--allow=tcp:22",
                "--enable-logging",
            ],
            execute=execute,
            name=f"update_firewall_rule:{name}",
            action="ensure_setting",
        ),
    ]


def ensure_vm(name: str, *, execute: bool) -> list[OperationResult]:
    vm = VMS[name]
    return create_if_missing(
        probe=["gcloud", "compute", "instances", "describe", name, "--zone", ZONE, "--project", PROJECT, "--format=json"],
        create=[
            "gcloud",
            "compute",
            "instances",
            "create",
            name,
            "--project",
            PROJECT,
            "--zone",
            ZONE,
            "--machine-type",
            vm["machine_type"],
            "--image-family=debian-12",
            "--image-project=debian-cloud",
            "--boot-disk-size",
            vm["boot_disk_size"],
            "--boot-disk-type=pd-balanced",
            "--network",
            NETWORK,
            "--subnet",
            SUBNET,
            "--no-address",
            "--tags",
            vm["tag"],
            "--service-account",
            vm["service_account"],
            "--scopes=cloud-platform",
            "--metadata=enable-oslogin=TRUE",
        ],
        execute=execute,
        name=f"create_vm:{name}",
    )


def ensure_snapshot_policy(*, execute: bool) -> list[OperationResult]:
    operations = create_if_missing(
        probe=[
            "gcloud",
            "compute",
            "resource-policies",
            "describe",
            SNAPSHOT_POLICY,
            "--region",
            REGION,
            "--project",
            PROJECT,
            "--format=json",
        ],
        create=[
            "gcloud",
            "compute",
            "resource-policies",
            "create",
            "snapshot-schedule",
            SNAPSHOT_POLICY,
            "--project",
            PROJECT,
            "--region",
            REGION,
            "--start-time=03:00",
            "--daily-schedule",
            "--max-retention-days=7",
            "--on-source-disk-delete=apply-retention-policy",
            "--storage-location",
            REGION,
        ],
        execute=execute,
        name=f"create_snapshot_policy:{SNAPSHOT_POLICY}",
    )
    for disk in VMS:
        operations.append(
            run_command(
                [
                    "gcloud",
                    "compute",
                    "disks",
                    "add-resource-policies",
                    disk,
                    "--project",
                    PROJECT,
                    "--zone",
                    ZONE,
                    "--resource-policies",
                    SNAPSHOT_POLICY,
                ],
                execute=execute,
                name=f"attach_snapshot_policy:{disk}",
                action="ensure_binding",
            )
        )
    return operations


def ensure_bucket(bucket: str, *, execute: bool) -> list[OperationResult]:
    operations = create_if_missing(
        probe=["gcloud", "storage", "buckets", "describe", f"gs://{bucket}", "--format=json"],
        create=[
            "gcloud",
            "storage",
            "buckets",
            "create",
            f"gs://{bucket}",
            "--project",
            PROJECT,
            "--location",
            REGION,
            "--uniform-bucket-level-access",
            "--public-access-prevention",
            "--soft-delete-duration=7d",
        ],
        execute=execute,
        name=f"create_bucket:{bucket}",
    )
    operations.append(
        run_command(
            [
                "gcloud",
                "storage",
                "buckets",
                "update",
                f"gs://{bucket}",
                "--versioning",
                "--uniform-bucket-level-access",
                "--public-access-prevention",
            ],
            execute=execute,
            name=f"harden_bucket:{bucket}",
            action="ensure_setting",
        )
    )
    return operations


def ensure_secret(password: str | None, *, execute: bool, password_env: str) -> list[OperationResult]:
    operations = create_if_missing(
        probe=["gcloud", "secrets", "describe", CLOUDSQL_PASSWORD_SECRET, "--project", PROJECT, "--format=json"],
        create=[
            "gcloud",
            "secrets",
            "create",
            CLOUDSQL_PASSWORD_SECRET,
            "--project",
            PROJECT,
            "--replication-policy=automatic",
        ],
        execute=execute,
        name=f"create_secret:{CLOUDSQL_PASSWORD_SECRET}",
    )
    if execute and password is None:
        operations.append(
            OperationResult(
                name=f"secret_version_missing_env:{password_env}",
                command=[],
                action="blocked_missing_secret_env",
                exitcode=1,
                stderr_tail=f"{password_env} is required to create the Cloud SQL password secret version",
            )
        )
        return operations
    operations.append(
        run_command(
            [
                "gcloud",
                "secrets",
                "versions",
                "add",
                CLOUDSQL_PASSWORD_SECRET,
                "--project",
                PROJECT,
                "--data-file=-",
            ],
            execute=execute,
            name=f"add_secret_version:{CLOUDSQL_PASSWORD_SECRET}",
            action="ensure_secret_version",
            recorded_command=[
                "gcloud",
                "secrets",
                "versions",
                "add",
                CLOUDSQL_PASSWORD_SECRET,
                "--project",
                PROJECT,
                "--data-file=-",
                f"# stdin from ${password_env}",
            ],
            stdin=password,
        )
    )
    return operations


def ensure_cloudsql(password: str | None, *, execute: bool, password_env: str) -> list[OperationResult]:
    recorded_create = [
        "gcloud",
        "sql",
        "instances",
        "create",
        CLOUDSQL_INSTANCE,
        "--project",
        PROJECT,
        "--database-version=POSTGRES_16",
        "--region",
        REGION,
        "--cpu=2",
        "--memory=8GiB",
        "--storage-type=SSD",
        "--storage-size=50GB",
        "--availability-type=zonal",
        "--network",
        f"projects/{PROJECT}/global/networks/{NETWORK}",
        "--no-assign-ip",
        "--ssl-mode=ENCRYPTED_ONLY",
        "--database-flags=cloudsql.iam_authentication=on",
        "--deletion-protection",
        "--backup-start-time=03:00",
        "--enable-point-in-time-recovery",
        "--retained-backups-count=7",
        "--retained-transaction-log-days=7",
        f"--root-password=${password_env}",
    ]
    if execute and password is None:
        return [
            OperationResult(
                name=f"cloudsql_password_missing_env:{password_env}",
                command=recorded_create,
                action="blocked_missing_secret_env",
                exitcode=1,
                stderr_tail=f"{password_env} is required to create {CLOUDSQL_INSTANCE}",
            )
        ]
    create = [part if part != f"--root-password=${password_env}" else f"--root-password={password}" for part in recorded_create]
    operations = create_if_missing(
        probe=[
            "gcloud",
            "sql",
            "instances",
            "describe",
            CLOUDSQL_INSTANCE,
            "--project",
            PROJECT,
            "--format=json",
        ],
        create=create,
        execute=execute,
        name=f"create_cloudsql_instance:{CLOUDSQL_INSTANCE}",
        recorded_create=recorded_create,
    )
    operations.append(
        run_command(
            [
                "gcloud",
                "sql",
                "databases",
                "create",
                CLOUDSQL_DATABASE,
                "--instance",
                CLOUDSQL_INSTANCE,
                "--project",
                PROJECT,
            ],
            execute=execute,
            name=f"create_cloudsql_database:{CLOUDSQL_DATABASE}",
            action="create_if_missing",
        )
    )
    return operations


def build_operations(*, execute: bool, admin_ssh_cidr: str | None, password_env: str) -> list[OperationResult]:
    try:
        ssh_cidr = validate_admin_ssh_cidr(admin_ssh_cidr, execute=execute)
    except ValueError as exc:
        return [OperationResult(name="validate_admin_ssh_cidr", command=[], action="validate", exitcode=1, stderr_tail=str(exc))]

    password = os.environ.get(password_env)
    operations: list[OperationResult] = []
    for api in REQUIRED_APIS:
        operations.append(
            run_command(
                ["gcloud", "services", "enable", api, "--project", PROJECT],
                execute=execute,
                name=f"enable_api:{api}",
                action="ensure_enabled",
            )
        )
        if operations[-1].exitcode not in (None, 0):
            return operations

    for account_id, account in SERVICE_ACCOUNTS.items():
        operations.extend(
            ensure_service_account(
                account_id,
                str(account["email"]),
                str(account["display_name"]),
                execute=execute,
            )
        )
        if operations[-1].exitcode not in (None, 0):
            return operations
        for role in account["roles"]:
            operations.append(ensure_project_binding(str(account["email"]), str(role), execute=execute, prefix=f"{account_id}_project_role"))
            if operations[-1].exitcode not in (None, 0):
                return operations

    operations.extend(ensure_network(execute=execute))
    if operations[-1].exitcode not in (None, 0):
        return operations
    operations.extend(ensure_subnet(execute=execute))
    if operations[-1].exitcode not in (None, 0):
        return operations
    operations.extend(ensure_private_services_access(execute=execute))
    if operations[-1].exitcode not in (None, 0):
        return operations

    for broad_rule in ("default-allow-ssh", "default-allow-rdp"):
        operations.append(
            run_command(
                ["gcloud", "compute", "firewall-rules", "update", broad_rule, "--project", PROJECT, "--disabled"],
                execute=execute,
                name=f"disable_broad_firewall_rule:{broad_rule}",
                action="disable_if_present",
            )
        )
        if operations[-1].exitcode not in (None, 0) and "not found" not in operations[-1].stderr_tail.lower():
            return operations

    for name, tag in SSH_FIREWALL_RULES.items():
        operations.extend(ensure_firewall_rule(name, tag, ssh_cidr, execute=execute))
        if operations[-1].exitcode not in (None, 0):
            return operations

    for name in VMS:
        operations.extend(ensure_vm(name, execute=execute))
        if operations[-1].exitcode not in (None, 0):
            return operations

    operations.extend(ensure_snapshot_policy(execute=execute))
    if operations[-1].exitcode not in (None, 0):
        return operations

    for bucket in BACKUP_BUCKETS:
        operations.extend(ensure_bucket(bucket, execute=execute))
        if operations[-1].exitcode not in (None, 0):
            return operations

    operations.extend(ensure_secret(password, execute=execute, password_env=password_env))
    if operations[-1].exitcode not in (None, 0):
        return operations
    operations.extend(ensure_cloudsql(password, execute=execute, password_env=password_env))
    return operations


def write_output(path: Path, payload: dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--execute", action="store_true", help="apply resources; default is dry-run")
    parser.add_argument("--admin-ssh-cidr", help="single trusted operator IPv4 /32 allowed to SSH")
    parser.add_argument("--postgres-password-env", default="TELEO_CLOUDSQL_POSTGRES_PASSWORD")
    parser.add_argument(
        "--output",
        default="outputs/gcp-infra-hardening-20260707/proofs/gcp-runtime-baseline-apply.json",
        help="proof JSON path",
    )
    args = parser.parse_args()

    operations = build_operations(
        execute=args.execute,
        admin_ssh_cidr=args.admin_ssh_cidr,
        password_env=args.postgres_password_env,
    )
    failed = [operation for operation in operations if operation.exitcode not in (None, 0)]
    payload = {
        "artifact": "teleo_gcp_runtime_baseline_apply",
        "generated_at_utc": datetime.now(UTC).isoformat(),
        "mode": "execute" if args.execute else "dry_run",
        "project": PROJECT,
        "region": REGION,
        "zone": ZONE,
        "network": NETWORK,
        "subnet": SUBNET,
        "admin_ssh_cidr": args.admin_ssh_cidr if args.admin_ssh_cidr else ("required_in_execute" if args.execute else "<operator-ip>/32"),
        "password_env": args.postgres_password_env,
        "resources": {
            "service_accounts": {account_id: account["email"] for account_id, account in SERVICE_ACCOUNTS.items()},
            "vms": VMS,
            "backup_buckets": BACKUP_BUCKETS,
            "snapshot_policy": SNAPSHOT_POLICY,
            "cloudsql_instance": CLOUDSQL_INSTANCE,
            "cloudsql_database": CLOUDSQL_DATABASE,
            "cloudsql_password_secret": CLOUDSQL_PASSWORD_SECRET,
        },
        "operations": [asdict(operation) for operation in operations],
        "status": "failed" if failed else ("applied" if args.execute else "planned"),
        "failed_operation_count": len(failed),
        "not_proven_by_this_artifact": [
            "GCP runtime readiness until execute mode succeeds and gcp-readiness passes",
            "Cloud SQL row-count parity until restore/import/readback verifier passes",
            "application cutover or Telegram bot parity",
        ],
        "next_checks": [
            "python3 ops/apply_gcp_iam_split.py --execute",
            "python3 ops/apply_gcp_runtime_baseline.py --execute --admin-ssh-cidr <operator-ip>/32",
            "gh workflow run gcp-readiness.yml --repo living-ip/teleo-infrastructure --ref main -f service_account=sa-teleo-readiness@teleo-501523.iam.gserviceaccount.com",
            "EXECUTE=1 ops/run_gcp_cloudsql_restore_drill.sh",
            "python3 ops/verify_gcp_cloudsql_restore_readback.py --drill-proof <proof.json> --target-counts-csv <cloudsql-counts.csv> --output <verification.json>",
        ],
    }
    write_output(Path(args.output), payload)
    print(json.dumps(payload, indent=2, sort_keys=True))
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())
