#!/usr/bin/env python3
"""Перенос баз PostgreSQL со стенда на локальный сервер: pg_dump -Fc + pg_restore.

Два режима: интерактивное меню (запуск без аргументов) и команды CLI.
"""

from __future__ import annotations

import argparse
import getpass
import json
import os
import re
import shutil
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

DEFAULT_CONFIG = Path(__file__).with_name("config.json")
STAMP_RE = re.compile(r"-\d{8}-\d{6}$")

# Windows-консоль по умолчанию отдаёт cp866, русский лог в ней превращается в мусор
for stream in (sys.stdout, sys.stderr):
    if hasattr(stream, "reconfigure"):
        stream.reconfigure(encoding="utf-8", errors="replace")


class ConfigError(Exception):
    pass


class StepError(Exception):
    pass


class Cancelled(Exception):
    """Пользователь отказался от выбора в интерактивном режиме."""


def log(msg: str) -> None:
    print(f"[{datetime.now():%H:%M:%S}] {msg}", flush=True)


# --------------------------------------------------------------------------- конфиг


def load_config(path: Path) -> dict:
    if not path.exists():
        raise ConfigError(
            f"конфиг не найден: {path}\n"
            f"скопируйте config.example.json в {path.name} и заполните"
        )
    try:
        cfg = json.loads(path.read_text(encoding="utf-8-sig"))
    except json.JSONDecodeError as exc:
        raise ConfigError(f"{path}: не разбирается как JSON — {exc}") from exc

    for section in ("source", "target"):
        node = cfg.get(section)
        if not isinstance(node, dict):
            raise ConfigError(f'{path}: нет секции "{section}"')
        for key in ("host", "user"):
            if not node.get(key):
                raise ConfigError(f"{path}: {section}.{key} не задан")
    return cfg


def resolve_tools(cfg: dict) -> dict:
    bin_dir = cfg.get("pg_bin_dir") or ""
    tools = {}
    for name in ("pg_dump", "pg_restore", "psql"):
        if bin_dir:
            candidate = Path(bin_dir) / (name + (".exe" if os.name == "nt" else ""))
            if not candidate.exists():
                raise ConfigError(f"не найден {name} в pg_bin_dir: {bin_dir}")
            tools[name] = str(candidate)
        else:
            found = shutil.which(name)
            if not found:
                raise ConfigError(f"{name} не найден в PATH, задайте pg_bin_dir в конфиге")
            tools[name] = found
    return tools


_password_cache: dict = {}


def resolve_password(node: dict, label: str) -> str | None:
    """Пароль: из конфига, из переменной окружения, иначе спросить один раз за запуск."""
    if node.get("password"):
        return str(node["password"])
    var = node.get("password_env")
    if var:
        value = os.environ.get(var)
        if value:
            return value
    key = (label, node.get("host"), node.get("port", 5432), node.get("user"))
    if key in _password_cache:
        return _password_cache[key]
    if not sys.stdin.isatty():
        if var:
            raise ConfigError(f"переменная окружения {var} пуста, а спросить пароль негде")
        return None
    value = getpass.getpass(
        f"пароль {node['user']}@{node['host']} ({label}), Enter — без пароля: ")
    _password_cache[key] = value
    return value


def conn_env(node: dict, label: str = "", dbname: str | None = None) -> dict:
    """Параметры подключения уходят в окружение, чтобы пароль не светился в командной строке."""
    env = dict(os.environ)
    env["PGHOST"] = str(node["host"])
    env["PGPORT"] = str(node.get("port", 5432))
    env["PGUSER"] = str(node["user"])
    database = dbname or node.get("database") or node.get("maintenance_database") or "postgres"
    env["PGDATABASE"] = str(database)
    password = resolve_password(node, label or node["host"])
    if password:
        env["PGPASSWORD"] = password
    else:
        env.pop("PGPASSWORD", None)
    if node.get("sslmode"):
        env["PGSSLMODE"] = str(node["sslmode"])
    env.setdefault("PGCLIENTENCODING", "UTF8")
    env["PGCONNECT_TIMEOUT"] = str(node.get("connect_timeout", 15))
    return env


# --------------------------------------------------------------------------- запуск утилит


def run(cmd: list, env: dict, what: str, dry_run: bool) -> None:
    shown = " ".join(f'"{c}"' if " " in c else c for c in cmd)
    log(f"{what}: {shown}")
    if dry_run:
        return
    started = time.monotonic()
    proc = subprocess.run(cmd, env=env, text=True, encoding="utf-8", errors="replace",
                          stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    if proc.stdout:
        for line in proc.stdout.splitlines():
            print("   " + line, flush=True)
    if proc.returncode != 0:
        raise StepError(f"{what} завершился с кодом {proc.returncode}")
    log(f"{what}: готово за {time.monotonic() - started:.1f} c")


def no_prompt(env: dict) -> list:
    """Без пароля в окружении утилиты просят его с терминала и подвисают в скриптах."""
    return [] if env.get("PGPASSWORD") else ["--no-password"]


def capture(cmd: list, env: dict, what: str) -> str:
    proc = subprocess.run(cmd, env=env, text=True, encoding="utf-8", errors="replace",
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if proc.returncode != 0:
        detail = (proc.stderr or "").strip().splitlines()
        raise StepError(f"{what}: " + (detail[-1] if detail else f"код {proc.returncode}"))
    return proc.stdout


def psql_exec(tools: dict, node: dict, dbname: str, sql: str, what: str,
              dry_run: bool, label: str = "") -> None:
    env = conn_env(node, label, dbname)
    cmd = [tools["psql"], "--no-psqlrc", "-v", "ON_ERROR_STOP=1", *no_prompt(env), "-c", sql]
    run(cmd, env, what, dry_run)


def list_databases(tools: dict, node: dict, label: str) -> list:
    """Список баз сервера: (имя, человекочитаемый размер)."""
    admin_db = node.get("maintenance_database") or node.get("database") or "postgres"
    env = conn_env(node, label, admin_db)
    sql = (
        "SELECT datname, "
        "  CASE WHEN has_database_privilege(datname, 'CONNECT') "
        "       THEN pg_size_pretty(pg_database_size(datname)) ELSE '-' END "
        "FROM pg_database WHERE NOT datistemplate AND datallowconn ORDER BY datname;"
    )
    cmd = [tools["psql"], "--no-psqlrc", "-At", "-F", "\t", *no_prompt(env), "-c", sql]
    out = capture(cmd, env, f"список баз {node['host']}")
    rows = []
    for line in out.splitlines():
        if not line.strip():
            continue
        name, _, size = line.partition("\t")
        rows.append((name, size or ""))
    return rows


# --------------------------------------------------------------------------- интерактив


def ask(prompt: str) -> str:
    try:
        return input(prompt).strip()
    except EOFError:
        raise Cancelled from None


def pick(items: list, title: str, render, multi: bool = False) -> list:
    """Список с фильтрацией по вводу: текст сужает список, число(а) выбирают, пустой ввод — сброс."""
    if not items:
        raise StepError(f"{title}: список пуст")
    if not sys.stdin.isatty():
        raise StepError(f"{title}: нет терминала для выбора")

    flt = ""
    while True:
        shown = [it for it in items if flt.lower() in render(it).lower()] if flt else list(items)
        print()
        print(f"— {title} —" + (f"  фильтр: {flt!r}" if flt else ""))
        if not shown:
            print("  ничего не найдено")
        for i, it in enumerate(shown, 1):
            print(f"  {i:>3}. {render(it)}")
        hint = "номера через запятую" if multi else "номер"
        choice = ask(f"[{hint} | текст — фильтр | пусто — сброс | q — назад]: ")

        if choice.lower() in ("q", "выход", "назад"):
            raise Cancelled
        if not choice:
            flt = ""
            continue
        if re.fullmatch(r"[\d\s,]+", choice):
            nums = [int(n) for n in re.split(r"[\s,]+", choice) if n]
            bad = [n for n in nums if not 1 <= n <= len(shown)]
            if bad or not nums:
                print(f"  нет такого номера: {', '.join(map(str, bad)) or '-'}")
                continue
            if not multi and len(nums) > 1:
                print("  здесь выбирается один пункт")
                continue
            return [shown[n - 1] for n in nums]
        flt = choice


def confirm(question: str, default: bool = True) -> bool:
    suffix = "[Y/n]" if default else "[y/N]"
    answer = ask(f"{question} {suffix}: ").lower()
    if not answer:
        return default
    return answer in ("y", "yes", "д", "да")


# --------------------------------------------------------------------------- шаги


def dump(cfg: dict, tools: dict, dbname: str, out: Path, dry_run: bool) -> Path:
    src = cfg["source"]
    opts = cfg.get("dump", {})
    cmd = [tools["pg_dump"], "--format=custom"]
    if opts.get("verbose"):
        cmd.append("--verbose")
    if opts.get("compress") is not None:
        cmd.append(f"--compress={opts['compress']}")
    if opts.get("schema_only"):
        cmd.append("--schema-only")
    if opts.get("no_owner", True):
        cmd.append("--no-owner")
    if opts.get("no_privileges", True):
        cmd.append("--no-privileges")
    for schema in opts.get("schemas", []):
        cmd += ["--schema", schema]
    for schema in opts.get("exclude_schemas", []):
        cmd += ["--exclude-schema", schema]
    for table in opts.get("exclude_tables", []):
        cmd += ["--exclude-table", table]
    for table in opts.get("exclude_table_data", []):
        cmd += ["--exclude-table-data", table]
    cmd += ["--file", str(out)]

    env = conn_env(src, "стенд", dbname)
    cmd += no_prompt(env)
    out.parent.mkdir(parents=True, exist_ok=True)
    run(cmd, env, f"дамп {dbname}@{src['host']}", dry_run)
    if not dry_run:
        size = out.stat().st_size / 1024 / 1024
        log(f"дамп: {out} ({size:.1f} МБ)")
    return out


def recreate_target(cfg: dict, tools: dict, dbname: str, dry_run: bool) -> None:
    tgt = cfg["target"]
    maintenance = tgt.get("maintenance_database", "postgres")
    quoted = '"' + dbname.replace('"', '""') + '"'
    literal = "'" + dbname.replace("'", "''") + "'"

    log(f"пересоздаю локальную базу {dbname}")
    psql_exec(tools, tgt, maintenance,
              f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
              f"WHERE datname = {literal} AND pid <> pg_backend_pid();",
              "обрыв активных сессий", dry_run, "локальная")
    psql_exec(tools, tgt, maintenance, f"DROP DATABASE IF EXISTS {quoted};",
              "DROP DATABASE", dry_run, "локальная")
    create = f"CREATE DATABASE {quoted}"
    if tgt.get("owner"):
        create += ' OWNER "' + tgt["owner"].replace('"', '""') + '"'
    if tgt.get("template"):
        create += f" TEMPLATE {tgt['template']}"
    psql_exec(tools, tgt, maintenance, create + ";", "CREATE DATABASE", dry_run, "локальная")


def restore(cfg: dict, tools: dict, dump_file: Path, dbname: str,
            recreated: bool, dry_run: bool) -> None:
    tgt = cfg["target"]
    opts = cfg.get("restore", {})
    if not dry_run and not dump_file.exists():
        raise StepError(f"файл дампа не найден: {dump_file}")

    cmd = [tools["pg_restore"], "--dbname", dbname]
    jobs = int(opts.get("jobs", 4))
    if jobs > 1:
        cmd += ["--jobs", str(jobs)]
    if opts.get("no_owner", True):
        cmd.append("--no-owner")
    if opts.get("no_privileges", True):
        cmd.append("--no-privileges")
    if opts.get("clean", True) and not recreated:
        cmd += ["--clean", "--if-exists"]
    if opts.get("exit_on_error", False):
        cmd.append("--exit-on-error")
    if opts.get("verbose"):
        cmd.append("--verbose")
    env = conn_env(tgt, "локальная", dbname)
    cmd += no_prompt(env)
    cmd.append(str(dump_file))

    try:
        run(cmd, env, f"restore в {dbname}@{tgt['host']}", dry_run)
    except StepError as exc:
        # без --exit-on-error pg_restore возвращает 1 на любых игнорируемых ошибках
        if opts.get("exit_on_error", False):
            raise
        log(f"внимание: {exc} — часть объектов могла не примениться, см. вывод выше")


def post_sql(cfg: dict, tools: dict, dbname: str, dry_run: bool) -> None:
    statements = cfg.get("post_restore_sql") or []
    for i, sql in enumerate(statements, 1):
        psql_exec(tools, cfg["target"], dbname, sql,
                  f"post-SQL {i}/{len(statements)}", dry_run, "локальная")


def prune(dump_dir: Path, keep: int) -> None:
    if keep <= 0 or not dump_dir.exists():
        return
    files = sorted(dump_dir.glob("*.dump"), key=lambda p: p.stat().st_mtime, reverse=True)
    for old in files[keep:]:
        log(f"удаляю старый дамп {old.name}")
        old.unlink()


# --------------------------------------------------------------------------- сценарии


def dump_dir_of(cfg: dict) -> Path:
    dump_dir = Path(cfg.get("dump_dir") or "dumps")
    if not dump_dir.is_absolute():
        dump_dir = Path(__file__).parent / dump_dir
    return dump_dir


def dump_path_for(cfg: dict, dbname: str) -> Path:
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    return dump_dir_of(cfg) / f"{dbname}-{stamp}.dump"


def db_from_dump_name(path: Path) -> str:
    return STAMP_RE.sub("", path.stem)


def local_dumps(cfg: dict) -> list:
    dump_dir = dump_dir_of(cfg)
    if not dump_dir.exists():
        return []
    return sorted(dump_dir.glob("*.dump"), key=lambda p: p.stat().st_mtime, reverse=True)


def do_dump(cfg: dict, tools: dict, dbnames: list, dry_run: bool) -> list:
    made = []
    for dbname in dbnames:
        made.append(dump(cfg, tools, dbname, dump_path_for(cfg, dbname), dry_run))
    prune(dump_dir_of(cfg), int(cfg.get("keep_dumps", 5)))
    return made


def do_restore(cfg: dict, tools: dict, dump_file: Path, target_db: str,
               no_recreate: bool, dry_run: bool) -> None:
    recreate = bool(cfg["target"].get("recreate", True)) and not no_recreate
    if recreate:
        recreate_target(cfg, tools, target_db, dry_run)
    restore(cfg, tools, dump_file, target_db, recreate, dry_run)
    post_sql(cfg, tools, target_db, dry_run)


def do_sync(cfg: dict, tools: dict, dbnames: list, target_db: str | None,
            no_recreate: bool, dry_run: bool) -> None:
    for dbname in dbnames:
        dump_file = dump(cfg, tools, dbname, dump_path_for(cfg, dbname), dry_run)
        local = target_db or cfg["target"].get("database") or dbname
        do_restore(cfg, tools, dump_file, local, no_recreate, dry_run)
    prune(dump_dir_of(cfg), int(cfg.get("keep_dumps", 5)))


# --------------------------------------------------------------------------- меню


def choose_source_dbs(cfg: dict, tools: dict, multi: bool) -> list:
    log(f"читаю список баз со стенда {cfg['source']['host']}")
    rows = list_databases(tools, cfg["source"], "стенд")
    chosen = pick(rows, f"базы на {cfg['source']['host']}",
                  lambda r: f"{r[0]:<32} {r[1]}", multi=multi)
    return [name for name, _ in chosen]


def choose_dump_file(cfg: dict) -> Path:
    files = local_dumps(cfg)
    if not files:
        raise StepError(f"в {dump_dir_of(cfg)} нет файлов *.dump")

    def render(p: Path) -> str:
        st = p.stat()
        return (f"{p.name:<44} {st.st_size / 1024 / 1024:>8.1f} МБ  "
                f"{datetime.fromtimestamp(st.st_mtime):%d.%m %H:%M}")

    return pick(files, f"дампы в {dump_dir_of(cfg)}", render)[0]


def ask_target_db(cfg: dict, default: str) -> str:
    answer = ask(f"локальная база [{default}] (Enter — оставить): ")
    return answer or default


def confirm_target(cfg: dict, what: str, target_db: str) -> bool:
    recreate = bool(cfg["target"].get("recreate", True))
    action = (f"пересоздать (DROP+CREATE) локальную {target_db}" if recreate
              else f"накатить поверх локальной {target_db}")
    return confirm(f"{what} → {action}?", default=not recreate)


def interactive(cfg: dict, tools: dict, dry_run: bool) -> int:
    src, tgt = cfg["source"], cfg["target"]
    while True:
        print()
        print("=" * 62)
        print(f"  стенд:    {src['user']}@{src['host']}:{src.get('port', 5432)}")
        print(f"  локально: {tgt['user']}@{tgt['host']}:{tgt.get('port', 5432)}"
              + ("   [dry-run]" if dry_run else ""))
        print("=" * 62)
        print("  1. Dump + Restore — снять со стенда и залить локально")
        print("  2. Dump           — только снять дамп со стенда")
        print("  3. Restore        — залить локально готовый дамп")
        print("  0. Выход")
        try:
            choice = ask("выбор: ")
            if choice in ("0", "q", "выход", ""):
                return 0

            if choice == "1":
                dbs = choose_source_dbs(cfg, tools, multi=False)
                target_db = ask_target_db(cfg, cfg["target"].get("database") or dbs[0])
                if not confirm_target(cfg, dbs[0], target_db):
                    continue
                do_sync(cfg, tools, dbs, target_db, no_recreate=False, dry_run=dry_run)
                log("перенос завершён")

            elif choice == "2":
                dbs = choose_source_dbs(cfg, tools, multi=True)
                do_dump(cfg, tools, dbs, dry_run)
                log("дамп завершён")

            elif choice == "3":
                dump_file = choose_dump_file(cfg)
                default_db = cfg["target"].get("database") or db_from_dump_name(dump_file)
                target_db = ask_target_db(cfg, default_db)
                if not confirm_target(cfg, dump_file.name, target_db):
                    continue
                do_restore(cfg, tools, dump_file, target_db, no_recreate=False, dry_run=dry_run)
                log("восстановление завершено")
            else:
                print("нет такого пункта")
        except Cancelled:
            continue
        except StepError as exc:
            print(f"ОШИБКА: {exc}", file=sys.stderr)


# --------------------------------------------------------------------------- CLI


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        description="Перенос БД PostgreSQL со стенда на локальный сервер. "
                    "Без команды запускается интерактивное меню.")
    ap.add_argument("-c", "--config", type=Path, default=DEFAULT_CONFIG, help="путь к config.json")
    ap.add_argument("--dry-run", action="store_true", help="показать команды, ничего не выполнять")
    sub = ap.add_subparsers(dest="command")

    p_sync = sub.add_parser("sync", help="dump + restore")
    p_sync.add_argument("db", nargs="*", help="базы на стенде (без имени — выбор из списка)")
    p_sync.add_argument("--target-db", help="имя локальной базы (по умолчанию как на стенде)")
    p_sync.add_argument("--no-recreate", action="store_true",
                        help="не пересоздавать локальную базу перед restore")

    p_dump = sub.add_parser("dump", help="только снять дамп со стенда")
    p_dump.add_argument("db", nargs="*", help="базы на стенде (без имени — выбор из списка)")

    p_restore = sub.add_parser("restore", help="только залить дамп локально")
    p_restore.add_argument("file", nargs="?", type=Path,
                           help="файл .dump (без имени — выбор из списка)")
    p_restore.add_argument("--target-db", help="имя локальной базы")
    p_restore.add_argument("--no-recreate", action="store_true",
                           help="не пересоздавать локальную базу перед restore")

    sub.add_parser("list", help="показать базы на стенде")
    sub.add_parser("dumps", help="показать локальные дампы")
    sub.add_parser("menu", help="интерактивное меню (то же, что запуск без команды)")
    return ap


def main() -> int:
    args = build_parser().parse_args()

    try:
        cfg = load_config(args.config)
        tools = resolve_tools(cfg)
        command = args.command or "menu"

        if command == "menu":
            return interactive(cfg, tools, args.dry_run)

        if command == "list":
            for name, size in list_databases(tools, cfg["source"], "стенд"):
                print(f"{name:<40} {size}")
            return 0

        if command == "dumps":
            for path in local_dumps(cfg):
                st = path.stat()
                print(f"{path.name:<48} {st.st_size / 1024 / 1024:>8.1f} МБ  "
                      f"{datetime.fromtimestamp(st.st_mtime):%d.%m.%Y %H:%M}")
            return 0

        if command in ("sync", "dump"):
            dbs = args.db or choose_source_dbs(cfg, tools, multi=(command == "dump"))
            if command == "dump":
                do_dump(cfg, tools, dbs, args.dry_run)
                log("дамп завершён")
            else:
                do_sync(cfg, tools, dbs, args.target_db, args.no_recreate, args.dry_run)
                log("перенос завершён")
            return 0

        if command == "restore":
            dump_file = args.file or choose_dump_file(cfg)
            target_db = (args.target_db or cfg["target"].get("database")
                         or db_from_dump_name(dump_file))
            do_restore(cfg, tools, dump_file, target_db, args.no_recreate, args.dry_run)
            log("восстановление завершено")
            return 0

        return 0
    except (ConfigError, StepError) as exc:
        print(f"ОШИБКА: {exc}", file=sys.stderr)
        return 1
    except Cancelled:
        print("отменено", file=sys.stderr)
        return 130
    except KeyboardInterrupt:
        print("прервано пользователем", file=sys.stderr)
        return 130


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