#!/usr/bin/env python3
"""Universal idempotent merge: example configs → local working copies.

Discovers ``*.example*`` / ``env.example`` under the project root and fills
missing keys/items into the corresponding local files without overwriting
existing values (API keys, tokens, chosen models, etc.).

Rules
-----
* Scalars: local wins if key exists (even empty).
* Dicts: deep-fill — add keys present in example but absent locally.
* Lists of scalars: append example items not already present.
* Lists of maps: match by identity key (name/id/url/dir/version_id),
  deep-fill matched rows, append unmatched example rows.
* KEY=VALUE (``.env``, ``*.vars``): append keys never mentioned locally
  (commented counts as present). Secrets get empty values.
  Use ``--docs`` to also append commented-only example keys as ``# KEY=``.
* Missing local file: copied from example (``--no-create`` to skip).

Safe to re-run. Writes ``*.bak.<utc>`` before each change.

Usage
-----
  python scripts/migrate_local_configs.py
  python scripts/migrate_local_configs.py --dry-run
  python scripts/migrate_local_configs.py --docs
  python scripts/migrate_local_configs.py --root D:\\Github\\gpu-rent
"""

from __future__ import annotations

import argparse
import re
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

ROOT_DEFAULT = Path(__file__).resolve().parents[1]

_ENV_LINE = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*)=(.*)$")
_IDENTITY_KEYS = ("name", "id", "url", "dir", "version_id", "path", "key")

# Heuristic: never invent secret values when adding a missing key.
_SECRET_SUBSTR = (
    "PASSWORD",
    "SECRET",
    "TOKEN",
    "API_KEY",
    "ACCESS_KEY",
    "PRIVATE_KEY",
    "CREDENTIAL",
)

# Keys owned by *.vars (vars override .env at runtime) — skip when filling .env.
_VARS_PREFERRED = {
    "LLM_RUNTIME",
    "ENABLE_SWARMUI",
    "OLLAMA_LOCAL_PORT",
    "IDLE_MINUTES",
    "IDLE_GRACE_MINUTES",
    "GPU_RENT_DEFAULT_ARGS",
    "GPU_RENT_EXTRA_ARGS",
    "WORKLOAD",
    "UPDATE_GIT",
    "DEFAULT_SPOT",
    "SCAN_POOLS",
    "FLAVOR_PREFERENCE",
    "SWARMUI_LOCAL_PORT",
}


Log = Callable[[str], None]


def _ts() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def _backup(path: Path, *, dry: bool, log: Log) -> Path | None:
    if not path.is_file():
        return None
    bak = path.with_name(f"{path.name}.bak.{_ts()}")
    if not dry:
        shutil.copy2(path, bak)
    log(f"  backup -> {bak.name}{' (dry)' if dry else ''}")
    return bak


def _is_secret_key(key: str) -> bool:
    u = key.upper()
    return any(s in u for s in _SECRET_SUBSTR)


def example_to_target(example: Path) -> Path | None:
    """Map example filename → local working path. None = not a config example."""
    name = example.name
    if name.startswith(".") and name.endswith(".example"):
        # e.g. .env.example → .env
        return example.with_name(name[: -len(".example")])
    if name == "env.example":
        return example.with_name(".env")
    if ".example." in name:
        # foo.example.yaml → foo.yaml
        left, _, right = name.partition(".example.")
        if not left or not right:
            return None
        return example.with_name(f"{left}.{right}")
    if name.endswith(".example"):
        return example.with_name(name[: -len(".example")])
    return None


def discover_pairs(root: Path) -> list[tuple[Path, Path]]:
    """Find (example, target) pairs in root (non-recursive, top-level only)."""
    pairs: list[tuple[Path, Path]] = []
    seen_targets: set[Path] = set()
    for path in sorted(root.iterdir()):
        if not path.is_file():
            continue
        target = example_to_target(path)
        if target is None:
            continue
        if target.resolve() == path.resolve():
            continue
        if target in seen_targets:
            continue
        seen_targets.add(target)
        pairs.append((path, target))
    return pairs


def _read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8-sig")


def _write_text(path: Path, text: str, *, dry: bool) -> None:
    if dry:
        return
    path.write_text(text if text.endswith("\n") else text + "\n", encoding="utf-8")


# ---------------------------------------------------------------------------
# KEY=VALUE merge
# ---------------------------------------------------------------------------


def _parse_assignments(text: str, *, include_commented: bool = False) -> dict[str, str]:
    out: dict[str, str] = {}
    for line in text.splitlines():
        s = line.strip()
        commented = False
        if s.startswith("#"):
            if not include_commented:
                continue
            s = s.lstrip("#").strip()
            commented = True
        if not s:
            continue
        if s.lower().startswith("export "):
            s = s[7:].strip()
        m = _ENV_LINE.match(s)
        if not m:
            continue
        key, val = m.group(1), m.group(2)
        # Drop inline trailing comments: KEY=val  # note
        if " #" in val:
            val = val.split(" #", 1)[0].rstrip()
        if key in out and commented:
            continue
        out[key] = val
    return out


def _keys_mentioned(text: str) -> set[str]:
    keys: set[str] = set()
    for line in text.splitlines():
        s = line.strip().lstrip("#").strip()
        if s.lower().startswith("export "):
            s = s[7:].strip()
        m = _ENV_LINE.match(s)
        if m:
            keys.add(m.group(1))
    return keys


def _collect_vars_keys(root: Path) -> set[str]:
    keys: set[str] = set()
    for path in root.glob("*.vars"):
        if path.is_file():
            keys |= _keys_mentioned(_read_text(path))
    return keys


def merge_env_style(
    target: Path,
    example: Path,
    *,
    skip_keys: set[str] | None = None,
    include_docs: bool = False,
    dry: bool = False,
    log: Log = print,
) -> bool:
    skip = skip_keys or set()
    if not example.is_file():
        return False

    if not target.is_file():
        log(f"{target.name}: create from {example.name}")
        if not dry:
            shutil.copy2(example, target)
        return True

    cur = _read_text(target)
    ex = _read_text(example)
    present = _keys_mentioned(cur)

    active = _parse_assignments(ex, include_commented=False)
    commented = (
        _parse_assignments(ex, include_commented=True) if include_docs else {}
    )

    to_add: list[tuple[str, str, bool]] = []  # key, value, as_comment
    for key, val in active.items():
        if key in present or key in skip:
            continue
        if _is_secret_key(key):
            to_add.append((key, "", False))
        else:
            to_add.append((key, val, False))

    for key, val in commented.items():
        if key in present or key in skip or key in active:
            continue
        if _is_secret_key(key):
            to_add.append((key, "", True))
        else:
            to_add.append((key, val, True))

    if not to_add:
        log(f"{target.name}: ok")
        return False

    _backup(target, dry=dry, log=log)
    block = ["", f"# --- added by migrate_local_configs {_ts()} ---"]
    for key, val, as_comment in to_add:
        line = f"{key}={val}"
        if as_comment:
            line = f"# {line}"
        block.append(line)
        log(f"  + {line}")
    _write_text(target, cur.rstrip("\n") + "\n" + "\n".join(block) + "\n", dry=dry)
    log(f"updated {target.name} (+{len(to_add)})")
    return True


# ---------------------------------------------------------------------------
# YAML deep-fill
# ---------------------------------------------------------------------------


def _item_identity(item: Any) -> str | None:
    if isinstance(item, str):
        return item.strip() or None
    if not isinstance(item, dict):
        return None
    for key in _IDENTITY_KEYS:
        val = item.get(key)
        if val is not None and str(val).strip():
            return f"{key}={val}"
    return None


def deep_fill(local: Any, example: Any) -> tuple[Any, bool]:
    """Fill missing structure from example; never overwrite local scalars."""
    if example is None:
        return local, False

    if local is None:
        return example, True

    if isinstance(local, dict) and isinstance(example, dict):
        changed = False
        out = dict(local)
        for key, ex_val in example.items():
            if key not in out:
                out[key] = ex_val
                changed = True
            else:
                merged, c = deep_fill(out[key], ex_val)
                if c:
                    out[key] = merged
                    changed = True
        return out, changed

    if isinstance(local, list) and isinstance(example, list):
        return _merge_lists(local, example)

    # Type mismatch or scalar: keep local.
    return local, False


def _merge_lists(local: list, example: list) -> tuple[list, bool]:
    if not example:
        return local, False

    # Scalar / homogeneous simple lists
    if all(not isinstance(x, (dict, list)) for x in local + example):
        changed = False
        out = list(local)
        seen = {x for x in out}
        for item in example:
            if item not in seen:
                out.append(item)
                seen.add(item)
                changed = True
        return out, changed

    # Normalize string entries to maps when peers are maps (ollama name-only rows).
    changed = False
    out: list[Any] = []
    index: dict[str, int] = {}

    def _ensure(item: Any) -> Any:
        nonlocal changed
        if isinstance(item, str):
            # Keep as string until we know maps dominate; handled below.
            return item
        return item

    for item in local:
        item = _ensure(item)
        out.append(item)
        ident = _item_identity(item)
        if ident is not None:
            index[ident] = len(out) - 1

    for ex_item in example:
        ident = _item_identity(ex_item)
        if ident is not None and ident in index:
            i = index[ident]
            cur = out[i]
            if isinstance(cur, str) and isinstance(ex_item, dict):
                # Promote "name" string to map and fill from example.
                promoted = dict(ex_item)
                # Prefer local name string as name field.
                name_key = next((k for k in _IDENTITY_KEYS if k in ex_item), "name")
                promoted[name_key] = cur
                # Local scalar had no extra fields — take example fields.
                out[i] = promoted
                changed = True
                continue
            merged, c = deep_fill(cur, ex_item)
            if c:
                out[i] = merged
                changed = True
        elif ident is not None:
            out.append(ex_item)
            index[ident] = len(out) - 1
            changed = True
        else:
            # Unidentified example row: append if not deep-equal present.
            if ex_item not in out:
                out.append(ex_item)
                changed = True

    return out, changed


def _leading_comments(text: str) -> str:
    lines: list[str] = []
    for line in text.splitlines():
        if line.lstrip().startswith("#") or line.strip() == "":
            lines.append(line)
            if line.strip() == "" and lines:
                # stop after first blank following comments
                # keep collecting leading comment block only
                pass
            continue
        break
    # Trim trailing blanks in header
    while lines and lines[-1].strip() == "":
        lines.pop()
    return "\n".join(lines)


def _dump_yaml(data: Any) -> str:
    import yaml

    return yaml.safe_dump(
        data,
        allow_unicode=True,
        default_flow_style=False,
        sort_keys=False,
        indent=2,
    )


def merge_yaml(
    target: Path,
    example: Path,
    *,
    dry: bool = False,
    log: Log = print,
    post_hooks: list[Callable[[Any], tuple[Any, bool]]] | None = None,
) -> bool:
    import yaml

    if not example.is_file():
        return False

    if not target.is_file():
        log(f"{target.name}: create from {example.name}")
        if not dry:
            shutil.copy2(example, target)
        return True

    local_text = _read_text(target)
    ex_text = _read_text(example)
    local_data = yaml.safe_load(local_text)
    ex_data = yaml.safe_load(ex_text)

    if ex_data is None:
        log(f"{target.name}: example empty — skip")
        return False
    if local_data is None:
        local_data = type(ex_data)()

    merged, changed = deep_fill(local_data, ex_data)

    if post_hooks:
        for hook in post_hooks:
            merged, c = hook(merged)
            changed = changed or c

    if not changed:
        log(f"{target.name}: ok")
        return False

    _backup(target, dry=dry, log=log)
    header = _leading_comments(local_text)
    body = _dump_yaml(merged)
    text = (header + "\n" + body) if header else body
    _write_text(target, text, dry=dry)
    log(f"updated {target.name}")
    return True


# ---------------------------------------------------------------------------
# Domain hooks (optional, keyed by target name)
# ---------------------------------------------------------------------------


def _hook_ollama_models(data: Any) -> tuple[Any, bool]:
    """Ensure use: chat|memory and a memory embed when chat models exist."""
    if not isinstance(data, dict):
        return data, False
    items = data.get("models")
    if not isinstance(items, list) or not items:
        return data, False

    try:
        root = Path(__file__).resolve().parents[1]
        sys.path.insert(0, str(root / "src"))
        from gpu_rent.llm_runtime import MEMORY_EMBED_MODEL, _normalize_use
    except Exception:
        MEMORY_EMBED_MODEL = "nomic-embed-text"

        def _normalize_use(v: Any) -> str:  # type: ignore
            s = str(v or "chat").strip().lower()
            return "memory" if s in {"memory", "embed", "embedding"} else "chat"

    def looks_embed(name: str) -> bool:
        n = name.lower()
        return any(x in n for x in ("embed", "nomic", "bge-", "minilm", "e5-"))

    changed = False
    out: list[Any] = []
    has_memory = False
    has_chat = False

    for item in items:
        if isinstance(item, str):
            name = item.strip()
            use = "memory" if looks_embed(name) else "chat"
            out.append({"name": name, "use": use})
            changed = True
            has_memory = has_memory or use == "memory"
            has_chat = has_chat or use == "chat"
            continue
        if not isinstance(item, dict):
            out.append(item)
            continue
        entry = dict(item)
        name = str(entry.get("name") or "").strip()
        if "use" not in entry and "role" not in entry:
            entry["use"] = "memory" if looks_embed(name) else "chat"
            changed = True
        else:
            entry["use"] = _normalize_use(entry.get("use") or entry.get("role"))
        if entry["use"] == "memory" or looks_embed(name):
            entry["use"] = "memory"
            has_memory = True
            if entry.get("default"):
                entry["default"] = False
                changed = True
        else:
            has_chat = True
        out.append(entry)

    if has_chat and not has_memory:
        out.append({"name": MEMORY_EMBED_MODEL, "use": "memory"})
        changed = True

    if changed:
        data = dict(data)
        data["models"] = out
    return data, changed


_HOOKS: dict[str, list[Callable[[Any], tuple[Any, bool]]]] = {
    "ollama-models.yaml": [_hook_ollama_models],
}


def _is_env_style(path: Path) -> bool:
    name = path.name.lower()
    return (
        name == ".env"
        or name.endswith(".env")
        or name.endswith(".vars")
        or name == "env"
        or ".vars" in name
    )


def migrate_pair(
    example: Path,
    target: Path,
    *,
    root: Path,
    create_missing: bool,
    include_docs: bool,
    dry: bool,
    log: Log,
) -> bool:
    if not example.is_file():
        return False
    if not target.is_file() and not create_missing:
        log(f"skip {target.name}: missing (--no-create)")
        return False

    if _is_env_style(target) or _is_env_style(example) or example.name in {
        "env.example",
        "gpu-rent.vars.example",
    }:
        skip: set[str] = set()
        if target.name == ".env":
            vars_keys = _collect_vars_keys(root)
            skip = {k for k in _VARS_PREFERRED if k in vars_keys}
        return merge_env_style(
            target,
            example,
            skip_keys=skip,
            include_docs=include_docs,
            dry=dry,
            log=log,
        )

    if target.suffix.lower() in {".yaml", ".yml"} or example.suffix.lower() in {
        ".yaml",
        ".yml",
    }:
        hooks = _HOOKS.get(target.name, [])
        return merge_yaml(target, example, dry=dry, log=log, post_hooks=hooks)

    if not target.is_file():
        log(f"{target.name}: create from {example.name} (other)")
        if not dry:
            shutil.copy2(example, target)
        return True
    log(f"{target.name}: skip (unknown type, already exists)")
    return False


def run(
    root: Path,
    *,
    dry: bool = False,
    create_missing: bool = True,
    include_docs: bool = False,
    only: set[str] | None = None,
    log: Log = print,
) -> int:
    root = root.resolve()
    log(f"migrate_local_configs in {root}{' [dry-run]' if dry else ''}")
    pairs = discover_pairs(root)
    if only:
        pairs = [
            (e, t)
            for e, t in pairs
            if t.name in only or e.name in only or t.stem in only
        ]
    if not pairs:
        log("no example->target pairs found")
        return 0

    changed_n = 0
    for example, target in pairs:
        log(f"- {example.name} -> {target.name}")
        if migrate_pair(
            example,
            target,
            root=root,
            create_missing=create_missing,
            include_docs=include_docs,
            dry=dry,
            log=log,
        ):
            changed_n += 1

    log(f"done: changed={changed_n}/{len(pairs)} (re-run safe)")
    return changed_n


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument(
        "--root",
        type=Path,
        default=ROOT_DEFAULT,
        help="project root with *example* files (default: repo of this script)",
    )
    p.add_argument("--dry-run", action="store_true", help="report only, no writes")
    p.add_argument(
        "--no-create",
        action="store_true",
        help="do not create missing local files from examples",
    )
    p.add_argument(
        "--docs",
        action="store_true",
        help="also append commented-only KEY= from examples as # KEY= lines",
    )
    p.add_argument(
        "--only",
        action="append",
        default=[],
        help="limit to target/example basename (repeatable)",
    )
    args = p.parse_args(argv)
    run(
        args.root,
        dry=args.dry_run,
        create_missing=not args.no_create,
        include_docs=args.docs,
        only=set(args.only) if args.only else None,
    )
    return 0


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