"""Remplit servers.ssh_method par defaut selon OS + environnement.

Regles SANEF :
  Linux   + Production     -> ssh_psmp   (PSMP SSH via CyberArk)
  Linux   + autre          -> ssh_key    (cybsecope cle SSH)
  Windows + Production     -> rdp_psmp   (RDP via PSMP CyberArk)
  Windows + autre          -> rdp_local  (RDP user@domain password)

N'applique pas de defaut si ssh_method est deja renseigne (sauf --overwrite).

Usage:
    python tools/fill_ssh_method_by_default.py [--dry-run] [--overwrite]
"""
import os
import argparse
from sqlalchemy import create_engine, text

DATABASE_URL = os.getenv("DATABASE_URL_DEMO") or os.getenv("DATABASE_URL") \
    or "postgresql://patchcenter:PatchCenter2026!@localhost:5432/patchcenter_demo"


def pick_method(os_family, environnement):
    is_prod = (environnement == "Production")
    os_f = (os_family or "").lower()
    if os_f == "linux":
        return "ssh_psmp" if is_prod else "ssh_key"
    if os_f == "windows":
        return "rdp_psmp" if is_prod else "rdp_local"
    return None


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--overwrite", action="store_true",
                        help="Remplace ssh_method existant (sinon ne rempli que les vides/a_definir)")
    args = parser.parse_args()

    engine = create_engine(DATABASE_URL)
    print(f"[INFO] DB: {DATABASE_URL.split('@')[-1]}")
    conn = engine.connect().execution_options(isolation_level="AUTOCOMMIT")

    # Etend le CHECK constraint pour accepter les nouveaux modes SANEF
    print("[INFO] Drop ancien CHECK ssh_method...")
    conn.execute(text("ALTER TABLE servers DROP CONSTRAINT IF EXISTS servers_ssh_method_check"))
    conn.execute(text(
        "ALTER TABLE servers ADD CONSTRAINT servers_ssh_method_check "
        "CHECK (ssh_method IN ('ssh_key', 'ssh_psmp', 'ssh_password', "
        "'rdp_local', 'rdp_psmp', 'winrm', 'a_definir') OR ssh_method IS NULL)"
    ))

    where = ""
    if not args.overwrite:
        where = "AND (ssh_method IS NULL OR ssh_method = '' OR ssh_method = 'a_definir' OR ssh_method = 'ssh_key')"

    rows = conn.execute(text(f"""
        SELECT id, hostname, os_family, environnement, ssh_method
        FROM servers
        WHERE os_family IS NOT NULL {where}
        ORDER BY hostname
    """)).fetchall()
    print(f"[INFO] {len(rows)} candidats")

    stats = {"updated": 0, "skipped": 0, "unchanged": 0}
    by_method = {}

    for r in rows:
        target = pick_method(r.os_family, r.environnement)
        if not target:
            stats["skipped"] += 1
            continue
        if r.ssh_method == target:
            stats["unchanged"] += 1
            continue
        by_method[target] = by_method.get(target, 0) + 1
        if args.dry_run:
            print(f"  DRY: {r.hostname:25s} [{r.os_family}/{r.environnement or '-'}] "
                  f"{r.ssh_method or 'NULL'} -> {target}")
        else:
            conn.execute(text("UPDATE servers SET ssh_method=:m WHERE id=:sid"),
                         {"m": target, "sid": r.id})
        stats["updated"] += 1

    conn.close()
    print(f"\n[DONE] Maj: {stats['updated']} | Inchanges: {stats['unchanged']} "
          f"| Skip: {stats['skipped']}")
    if by_method:
        print("\nRepartition :")
        for m, n in sorted(by_method.items(), key=lambda x: -x[1]):
            print(f"  {m:15s} {n}")


if __name__ == "__main__":
    main()