#!/usr/bin/env python3
"""Capture non-secret apps-pg capacity evidence for resource-control.

Used by RAILIANCE-WP-0016. Never prints credentials or role passwords.
"""
from __future__ import annotations

import argparse
import json
import shlex
import subprocess
import sys
from datetime import datetime, timezone
from typing import Any

DEFAULT_REMOTE = "railiance01"
NAMESPACE = "databases"
CLUSTER = "apps-pg"
POD = "apps-pg-1"
RESOURCE_ID = "resource:railiance:apps-pg"

CONSUMERS = {
    "vergabe_db": {
        "name": "vergabe-teilnahme",
        "role": "vergabe",
        "cost_attribution_key": "platform:vergabe-teilnahme",
    },
    "coulomb_social_db": {
        "name": "coulomb-social",
        "role": "coulomb_social",
        "cost_attribution_key": "platform:coulomb-social",
    },
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n", 1)[0])
    parser.add_argument("--remote", default=DEFAULT_REMOTE)
    parser.add_argument("-o", "--output")
    return parser.parse_args()


def ssh_run(remote: str, remote_cmd: str) -> str:
    cmd = ["ssh", "-o", "BatchMode=yes", remote, remote_cmd]
    result = subprocess.run(cmd, check=False, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"{remote_cmd} failed: {(result.stderr or '').strip()}")
    return result.stdout


def ssh_kubectl(remote: str, *args: str) -> str:
    return ssh_run(remote, "kubectl " + " ".join(shlex.quote(arg) for arg in args))


def ssh_psql(remote: str, sql: str) -> str:
    remote_cmd = (
        "kubectl exec -n {ns} {pod} -c postgres -- "
        "psql -U postgres -d postgres -At -F , -c {sql}"
    ).format(ns=shlex.quote(NAMESPACE), pod=shlex.quote(POD), sql=shlex.quote(sql))
    return ssh_run(remote, remote_cmd)


def as_int(value: str) -> int:
    return int(value.strip())


def collect(remote: str) -> dict[str, Any]:
    cluster = json.loads(
        ssh_kubectl(remote, "get", "cluster", CLUSTER, "-n", NAMESPACE, "-o", "json")
    )
    pvc = json.loads(
        ssh_kubectl(
            remote,
            "get",
            "pvc",
            "-n",
            NAMESPACE,
            "-l",
            f"cnpg.io/cluster={CLUSTER}",
            "-o",
            "json",
        )
    )
    top = ssh_kubectl(
        remote, "top", "pod", "-n", NAMESPACE, "-l", f"cnpg.io/cluster={CLUSTER}", "--no-headers"
    ).strip()
    cpu_obs, mem_obs = None, None
    if top:
        parts = top.split()
        if len(parts) >= 3:
            cpu_obs, mem_obs = parts[1], parts[2]

    pgdata = ssh_kubectl(
        remote,
        "exec",
        "-n",
        NAMESPACE,
        POD,
        "-c",
        "postgres",
        "--",
        "du",
        "-sb",
        "/var/lib/postgresql/data/pgdata",
    ).split()[0]

    db_rows = []
    for line in ssh_psql(
        remote,
        "SELECT datname, pg_database_size(datname) FROM pg_database "
        "WHERE datistemplate = false ORDER BY 1",
    ).splitlines():
        if not line.strip():
            continue
        name, size = line.split(",", 1)
        db_rows.append({"name": name, "bytes": as_int(size)})

    stats = {}
    for line in ssh_psql(
        remote,
        "SELECT datname, numbackends, xact_commit, xact_rollback, blks_read, blks_hit, "
        "tup_inserted, tup_updated, tup_deleted FROM pg_stat_database "
        "WHERE datname IS NOT NULL AND datname NOT LIKE 'template%' ORDER BY 1",
    ).splitlines():
        if not line.strip():
            continue
        cols = line.split(",")
        stats[cols[0]] = {
            "backends": as_int(cols[1]),
            "xact_commit": as_int(cols[2]),
            "xact_rollback": as_int(cols[3]),
            "blks_read": as_int(cols[4]),
            "blks_hit": as_int(cols[5]),
            "tup_inserted": as_int(cols[6]),
            "tup_updated": as_int(cols[7]),
            "tup_deleted": as_int(cols[8]),
        }

    sessions = as_int(ssh_psql(remote, "SELECT count(*) FROM pg_stat_activity"))
    max_conn = as_int(ssh_psql(remote, "SELECT current_setting('max_connections')"))
    wal = ssh_psql(
        remote, "SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), '0/0')"
    ).strip()
    version = ssh_psql(remote, "SELECT version()").strip()

    backups = ssh_kubectl(
        remote, "get", "backup,scheduledbackup", "-n", NAMESPACE, "--ignore-not-found"
    ).strip()

    spec = cluster.get("spec") or {}
    status = cluster.get("status") or {}
    resources = spec.get("resources") or {}
    pvc_item = (pvc.get("items") or [{}])[0]
    provisioned = ((pvc_item.get("status") or {}).get("capacity") or {}).get("storage")
    storage_class = (pvc_item.get("spec") or {}).get("storageClassName")

    conditions = {c.get("type"): c for c in status.get("conditions") or []}
    consumers = []
    consumer_bytes = 0
    for db in db_rows:
        meta = CONSUMERS.get(db["name"])
        if not meta:
            continue
        stat = stats.get(db["name"], {})
        consumer_bytes += db["bytes"]
        consumers.append(
            {
                "database": db["name"],
                "workload": meta["name"],
                "role": meta["role"],
                "cost_attribution_key": meta["cost_attribution_key"],
                "bytes": db["bytes"],
                "xact_commit": stat.get("xact_commit", 0),
                "backends": stat.get("backends", 0),
            }
        )

    pgdata_bytes = as_int(pgdata)
    residual_bytes = max(pgdata_bytes - sum(db["bytes"] for db in db_rows), 0)
    for row in consumers:
        row["share_of_consumer_bytes"] = (
            round(row["bytes"] / consumer_bytes, 4) if consumer_bytes else None
        )

    return {
        "schema_version": "0.1",
        "record_type": "usage_observation",
        "resource_id": RESOURCE_ID,
        "source": "railiance-platform",
        "workplan_id": "RAILIANCE-WP-0016",
        "reef": "reef-railiance",
        "captured_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "identity": {
            "cluster": CLUSTER,
            "namespace": NAMESPACE,
            "pod": status.get("currentPrimary") or POD,
            "image": status.get("image") or spec.get("imageName"),
            "postgres_version": version,
            "created_at": cluster.get("metadata", {}).get("creationTimestamp"),
            "phase": status.get("phase"),
        },
        "capacity": {
            "instances_provisioned": spec.get("instances"),
            "instances_ready": status.get("readyInstances"),
            "cpu_request": (resources.get("requests") or {}).get("cpu"),
            "memory_request": (resources.get("requests") or {}).get("memory"),
            "cpu_limit": (resources.get("limits") or {}).get("cpu"),
            "memory_limit": (resources.get("limits") or {}).get("memory"),
            "storage_provisioned": provisioned or spec.get("storage", {}).get("size"),
            "storage_class": storage_class,
            "max_connections": max_conn,
        },
        "utilization": {
            "cpu_observed": cpu_obs,
            "memory_observed": mem_obs,
            "pgdata_bytes": pgdata_bytes,
            "wal_bytes_since_init": as_int(wal) if wal else None,
            "sessions": sessions,
            "databases": [{**db, **stats.get(db["name"], {})} for db in db_rows],
        },
        "consumers": consumers,
        "recovery": {
            "cluster_ready": (conditions.get("Ready") or {}).get("status") == "True",
            "continuous_archiving": (conditions.get("ContinuousArchiving") or {}).get(
                "status"
            )
            == "True",
            "scheduled_backup_present": bool(backups),
            "option_a_target": False,
            "restore_drill_recorded": False,
            "rpo": "unbounded",
            "rto": "undefined",
            "failure_domain": "single-host local-path on railiance01",
        },
        "allocation": {
            "method": "proportional",
            "driver": "database_gb",
            "method_version": "apps-pg-dbbytes-v1",
            "consumer_bytes": consumer_bytes,
            "residual_bytes": residual_bytes,
            "residual_key": "platform:apps-pg-overhead",
            "notes": [
                "Consumer share is pg_database_size of declared consumer databases only.",
                "Residual is PGDATA minus all non-template databases (WAL and catalogs).",
            ],
        },
        "provenance": {
            "commands": [
                f"kubectl get cluster {CLUSTER} -n {NAMESPACE} -o json",
                f"kubectl top pod -n {NAMESPACE} -l cnpg.io/cluster={CLUSTER}",
                f"kubectl exec -n {NAMESPACE} {POD} -c postgres -- du -sb /var/lib/postgresql/data/pgdata",
                "psql non-secret catalog and pg_stat_database queries",
            ],
            "authority": "reef-railiance Kubernetes API + local postgres catalog",
        },
    }


def main() -> int:
    args = parse_args()
    try:
        payload = collect(args.remote)
    except (RuntimeError, json.JSONDecodeError, ValueError) as exc:
        print(f"capture failed: {exc}", file=sys.stderr)
        return 1
    text = json.dumps(payload, indent=2) + "\n"
    if args.output:
        with open(args.output, "w", encoding="utf-8") as fh:
            fh.write(text)
    else:
        sys.stdout.write(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())
