#!/usr/bin/env python3
"""TSV registry helpers for Loop PG migration.

Columns:
  loop_id, domain, kube_context, yc_cluster, target_cluster,
  size_bytes, size_pretty, batch_id, status, notes
"""
from __future__ import annotations

import argparse
import csv
import fcntl
import os
import re
import sys
import tempfile
import time
from contextlib import contextmanager
from pathlib import Path

COLUMNS = [
    "loop_id",
    "domain",
    "kube_context",
    "yc_cluster",
    "target_cluster",
    "size_bytes",
    "size_pretty",
    "batch_id",
    "status",
    "notes",
]

# Fields registry-refresh may overwrite from live scan
REFRESH_FIELDS = ("domain", "kube_context", "yc_cluster")

# Fields never clobbered by refresh if already set (unless --force-*)
PRESERVE_IF_SET = ("target_cluster", "batch_id", "status", "size_bytes", "size_pretty", "notes")

STATUSES = {
    "pending",
    "sized",
    "migrating",
    "migrated",
    "verified",
    "cutover",
    "done",
    "failed",
    "size_error",
    "excluded",
}

DONE_LIKE = {"migrated", "verified", "cutover", "done"}


def registry_path(migrate_dir: Path) -> Path:
    return migrate_dir / "registry.tsv"


def empty_row(loop_id: str = "") -> dict[str, str]:
    return {c: "" for c in COLUMNS} | {"loop_id": loop_id, "status": "pending"}


@contextmanager
def registry_lock(path: Path, timeout_sec: float = 120.0):
    """Exclusive flock for registry.tsv (and its done/failed mirrors)."""
    path.parent.mkdir(parents=True, exist_ok=True)
    lock_path = path.parent / f"{path.name}.lock"
    fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o644)
    deadline = time.monotonic() + timeout_sec
    try:
        while True:
            try:
                fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
                break
            except BlockingIOError:
                if time.monotonic() >= deadline:
                    fcntl.flock(fd, fcntl.LOCK_EX)  # block past timeout rather than fail cutover
                    break
                time.sleep(0.05)
        yield
    finally:
        try:
            fcntl.flock(fd, fcntl.LOCK_UN)
        finally:
            os.close(fd)


def _count_data_lines(path: Path) -> int:
    if not path.exists():
        return 0
    n = 0
    with path.open() as f:
        for line in f:
            if not line.strip():
                continue
            if line.startswith("loop_id\t") or line.strip() == "loop_id":
                continue
            n += 1
    return n


def load(path: Path) -> dict[str, dict[str, str]]:
    """Load registry.tsv. Tolerates header missing or moved (e.g. after `sort`)."""
    if not path.exists():
        return {}
    size = path.stat().st_size
    rows: dict[str, dict[str, str]] = {}
    with path.open(newline="") as f:
        sample = f.readline()
        if not sample:
            return {}
        # Normal file: first line is the header
        if sample.startswith("loop_id\t") or sample.strip() == "loop_id":
            f.seek(0)
            reader = csv.DictReader(f, delimiter="\t")
        else:
            # Header missing or sorted to the bottom — parse every line as data
            f.seek(0)
            reader = csv.DictReader(f, fieldnames=COLUMNS, delimiter="\t")
        for raw in reader:
            if not raw:
                continue
            loop_id = (raw.get("loop_id") or "").strip()
            if not loop_id or loop_id.startswith("#"):
                continue
            # stray header row in the body
            if loop_id == "loop_id":
                continue
            row = empty_row(loop_id)
            for c in COLUMNS:
                if c in raw and raw[c] is not None:
                    row[c] = str(raw[c]).strip()
            rows[loop_id] = row
    # Non-empty file that parses to nothing = corruption / wrong format
    if size >= 500 and not rows:
        raise RuntimeError(f"registry {path} is {size} bytes but parsed 0 rows — refuse load")
    return rows


def _save_unlocked(path: Path, rows: dict[str, dict[str, str]]) -> None:
    """Atomic write via unique tempfile (safe under parallel cutover-list)."""
    path.parent.mkdir(parents=True, exist_ok=True)
    old_n = _count_data_lines(path)
    new_n = len(rows)
    # Hard guard: never collapse a full registry to a handful of rows (parallel race / empty load).
    # Override: REGISTRY_ALLOW_SHRINK=1
    allow_shrink = os.environ.get("REGISTRY_ALLOW_SHRINK", "").strip() in ("1", "true", "yes")
    if not allow_shrink and old_n >= 20 and new_n < max(10, old_n // 2):
        raise RuntimeError(
            f"refusing to shrink registry {path} from {old_n} to {new_n} rows "
            f"(set REGISTRY_ALLOW_SHRINK=1 to override)"
        )
    ordered = sorted(
        rows.values(),
        key=lambda r: (r.get("batch_id") or "zzz", r.get("kube_context") or "", r["loop_id"]),
    )
    fd, tmp_name = tempfile.mkstemp(
        prefix=f".{path.name}.",
        suffix=".tmp",
        dir=str(path.parent),
    )
    try:
        with os.fdopen(fd, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=COLUMNS, delimiter="\t", lineterminator="\n")
            w.writeheader()
            for r in ordered:
                w.writerow({c: r.get(c, "") for c in COLUMNS})
        os.replace(tmp_name, path)
    except Exception:
        try:
            os.unlink(tmp_name)
        except OSError:
            pass
        raise


def save(path: Path, rows: dict[str, dict[str, str]]) -> None:
    with registry_lock(path):
        _save_unlocked(path, rows)


def upsert(rows: dict[str, dict[str, str]], patch: dict[str, str], *, force: bool = False) -> dict[str, str]:
    loop_id = patch["loop_id"]
    cur = rows.get(loop_id) or empty_row(loop_id)
    for k, v in patch.items():
        if k not in COLUMNS or v is None:
            continue
        v = str(v).strip()
        if k == "loop_id":
            cur[k] = v
            continue
        if not force and k in PRESERVE_IF_SET and cur.get(k):
            # allow status overwrite only when explicitly forced or empty
            if k == "status" and not force:
                continue
            if k != "status":
                continue
        if v != "" or force:
            cur[k] = v
    rows[loop_id] = cur
    return cur


def _write_id_list(path: Path, ids: list[str]) -> None:
    text = "\n".join(ids) + ("\n" if ids else "")
    fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=str(path.parent))
    try:
        with os.fdopen(fd, "w") as f:
            f.write(text)
        os.replace(tmp_name, path)
    except Exception:
        try:
            os.unlink(tmp_name)
        except OSError:
            pass
        raise


def _sync_legacy_mirrors(registry: Path, loop_id: str, status: str) -> None:
    """done.txt / failed.txt are write-only mirrors of registry status. Caller holds registry_lock."""
    root = registry.parent
    done_p, fail_p = root / "done.txt", root / "failed.txt"

    def _read_ids(path: Path) -> list[str]:
        if not path.exists():
            return []
        return [x for x in path.read_text().split() if x]

    def _add(path: Path, lid: str) -> None:
        cur = _read_ids(path)
        if lid not in cur:
            cur.append(lid)
            _write_id_list(path, sorted(cur))

    def _del(path: Path, lid: str) -> None:
        if not path.exists():
            return
        cur = [x for x in _read_ids(path) if x != lid]
        _write_id_list(path, cur)

    if status in DONE_LIKE:
        _add(done_p, loop_id)
        _del(fail_p, loop_id)
    elif status == "failed":
        _add(fail_p, loop_id)
        _del(done_p, loop_id)
    elif status:
        _del(done_p, loop_id)
        _del(fail_p, loop_id)


def set_fields(path: Path, loop_id: str, fields: dict[str, str], *, force: bool = True) -> None:
    # load→upsert→save under one lock so parallel cutover-list does not clobber rows
    with registry_lock(path):
        size = path.stat().st_size if path.exists() else 0
        rows = load(path)
        # Empty parse of a previously-large file would wipe the registry on save — abort
        if size >= 500 and not rows:
            raise RuntimeError(f"registry {path} unreadable ({size} bytes, 0 rows)")
        if size >= 5000 and len(rows) < 10:
            raise RuntimeError(
                f"registry {path} looks truncated ({size} bytes, {len(rows)} rows) — refuse update"
            )
        patch = {"loop_id": loop_id, **fields}
        # Never let force=True clear unrelated columns when caller only patches status/etc.
        # Only apply non-empty values unless key is status (explicit empty not used).
        clean = {"loop_id": loop_id}
        for k, v in fields.items():
            if k == "loop_id":
                continue
            if v is None:
                continue
            v = str(v)
            if v == "" and k != "notes":
                continue
            clean[k] = v
        upsert(rows, clean, force=force)
        _save_unlocked(path, rows)
        if "status" in fields:
            _sync_legacy_mirrors(path, loop_id, fields["status"])


def is_done_status(status: str) -> bool:
    return status in DONE_LIKE


# Known YC MDB cluster ids (fallback when `yc` CLI unavailable)
KNOWN_CID_NAMES = {
    "c9q9hoccq1gs9jf2g3r9": "LP-PG-1",
    "c9q87tpl86oinfglqem4": "LP-PG-2",
    "c9q6tst6nutgn8ju2p5r": "LP-PG-3",
    "c9q9gqbgbnljg3apqkak": "LP-PG-4",
    "c9q0a21dumlqo9u3vmm0": "LP-PG-5",
    "c9qqfe1b3e43gkp4jju5": "LP-PG-6",
    "c9q4h551pgq3po55f7hu": "LP-PG-FF",
    "c9qfg20anmmknvmpqiuj": "LP-PG-LM",
    "c9qfkheuh46d1g731bqi": "postgresql120",
    "c9qnuaa4qqfrcj11h4rm": "4",
    "c9qbkifpjolmls057bij": "5",
    "c9qe5f6bo1c7isk8id9g": "6",
}

# Original shared pool → restore cluster (host rewrite after source cluster delete/move)
# old_name → (new_name, old_cid, new_cid)
RESTORE_REWRITES = {
    "LP-PG-1": ("postgresql120", "c9q9hoccq1gs9jf2g3r9", "c9qfkheuh46d1g731bqi"),
    "LP-PG-4": ("4", "c9q9gqbgbnljg3apqkak", "c9qnuaa4qqfrcj11h4rm"),
    "LP-PG-5": ("5", "c9q0a21dumlqo9u3vmm0", "c9qbkifpjolmls057bij"),
    "LP-PG-6": ("6", "c9qqfe1b3e43gkp4jju5", "c9qe5f6bo1c7isk8id9g"),
}


def rw_host(cid: str) -> str:
    return f"c-{cid}.rw.mdb.yandexcloud.net"


def restore_host_map() -> dict[str, tuple[str, str]]:
    """old_fqdn → (new_fqdn, new_cluster_name)."""
    out: dict[str, tuple[str, str]] = {}
    for old_name, (new_name, old_cid, new_cid) in RESTORE_REWRITES.items():
        out[rw_host(old_cid)] = (rw_host(new_cid), new_name)
    return out


def host_to_yc_cluster(host: str, cid_to_name: dict[str, str] | None = None) -> str:
    host = host or ""
    if "95.143.178.247" in host:
        return "CNPG-wilix"
    if "135.106.201.40" in host:
        return "CNPG-lm"
    if "135.106.201.38" in host:
        return "CNPG-ff"
    names = {**KNOWN_CID_NAMES, **(cid_to_name or {})}
    m = re.search(r"c-([a-z0-9]+)\.rw\.mdb", host)
    if m:
        cid = m.group(1)
        return names.get(cid, cid)
    if "mdb.yandexcloud.net" in host:
        return f"other-yc:{host.split(':')[0]}"
    if host:
        return host.split(":")[0]
    return ""


def cmd_get(args: argparse.Namespace) -> int:
    rows = load(Path(args.registry))
    row = rows.get(args.loop_id)
    if not row:
        return 1
    if args.field:
        print(row.get(args.field, ""))
    else:
        print("\t".join(row.get(c, "") for c in COLUMNS))
    return 0


def cmd_set(args: argparse.Namespace) -> int:
    fields = {}
    for item in args.field:
        if "=" not in item:
            raise SystemExit(f"bad field (want key=value): {item}")
        k, v = item.split("=", 1)
        if k not in COLUMNS or k == "loop_id":
            raise SystemExit(f"invalid field: {k}")
        fields[k] = v
    set_fields(Path(args.registry), args.loop_id, fields, force=True)
    return 0


def cmd_export(args: argparse.Namespace) -> int:
    rows = load(Path(args.registry))
    out = sys.stdout
    close = False
    if args.output and args.output != "-":
        out = open(args.output, "w")  # noqa: SIM115
        close = True
    try:
        n = 0
        for r in sorted(rows.values(), key=lambda x: x["loop_id"]):
            if args.batch and r.get("batch_id") != args.batch:
                continue
            if args.status:
                allowed = set(args.status.split(","))
                if r.get("status") not in allowed:
                    continue
            if args.target and r.get("target_cluster") != args.target:
                continue
            if args.skip_done and is_done_status(r.get("status", "")):
                continue
            ctx = r.get("kube_context") or ""
            lid = r.get("loop_id") or ""
            if not ctx or not lid:
                continue
            # migrate-list line; optional target hint as 3rd col if cnpg/selectel
            line = f"{ctx}  {lid}"
            if args.with_target_hint and r.get("target_cluster"):
                line += "  cnpg"
            out.write(line + "\n")
            n += 1
        if args.output and args.output != "-":
            print(f"exported {n} rows", file=sys.stderr)
    finally:
        if close:
            out.close()
    return 0


def cmd_summary(args: argparse.Namespace) -> int:
    rows = load(Path(args.registry))
    from collections import Counter

    by_status = Counter(r.get("status") or "" for r in rows.values())
    by_batch = Counter(r.get("batch_id") or "(none)" for r in rows.values())
    by_yc = Counter(r.get("yc_cluster") or "(none)" for r in rows.values())
    total_bytes = 0
    sized = 0
    for r in rows.values():
        b = r.get("size_bytes") or ""
        if b.isdigit():
            total_bytes += int(b)
            sized += 1
    print(f"rows={len(rows)} sized={sized} total≈{total_bytes / (1024**3):.1f} GiB")
    print("status:", dict(by_status))
    print("batch:", dict(sorted(by_batch.items())))
    print("yc_cluster top:", by_yc.most_common(15))
    return 0


def cmd_is_done(args: argparse.Namespace) -> int:
    rows = load(Path(args.registry))
    row = rows.get(args.loop_id)
    if row and is_done_status(row.get("status", "")):
        return 0
    return 1


def main() -> int:
    p = argparse.ArgumentParser(description="Loop migrate registry.tsv helpers")
    p.add_argument("--registry", default="", help="path to registry.tsv")
    sub = p.add_subparsers(dest="cmd", required=True)

    g = sub.add_parser("get", help="print row or one field")
    g.add_argument("loop_id")
    g.add_argument("field", nargs="?")
    g.set_defaults(func=cmd_get)

    s = sub.add_parser("set", help="set fields: key=value …")
    s.add_argument("loop_id")
    s.add_argument("field", nargs="+")
    s.set_defaults(func=cmd_set)

    e = sub.add_parser("export", help="export migrate-list lines")
    e.add_argument("--batch", default="")
    e.add_argument("--status", default="", help="comma-separated statuses")
    e.add_argument("--target", default="")
    e.add_argument("--skip-done", action="store_true")
    e.add_argument("--with-target-hint", action="store_true")
    e.add_argument("-o", "--output", default="-")
    e.set_defaults(func=cmd_export)

    u = sub.add_parser("summary")
    u.set_defaults(func=cmd_summary)

    d = sub.add_parser("is-done", help="exit 0 if status is migrated/verified/cutover/done")
    d.add_argument("loop_id")
    d.set_defaults(func=cmd_is_done)

    args = p.parse_args()
    if not args.registry:
        args.registry = str(Path(__file__).resolve().parent / "registry.tsv")
    return args.func(args)


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