#!/usr/bin/env python3
from pathlib import Path
from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedSeq
import sys

yaml = YAML()
yaml.preserve_quotes = True
yaml.indent(mapping=2, sequence=4, offset=2)

def ensure_env_file(service: dict, env_path: str = ".env") -> bool:
    """
    Ensure the service has:
      env_file:
        - .env

    Returns True if the service was modified, else False.
    """
    changed = False

    if "env_file" not in service:
        seq = CommentedSeq()
        seq.append(env_path)
        service["env_file"] = seq
        return True

    env_file = service["env_file"]

    if isinstance(env_file, str):
        if env_file != env_path:
            seq = CommentedSeq()
            seq.append(env_file)
            seq.append(env_path)
            service["env_file"] = seq
            changed = True

    elif isinstance(env_file, list):
        if env_path not in env_file:
            env_file.append(env_path)
            changed = True

    else:
        raise TypeError(f"Unsupported env_file type: {type(env_file).__name__}")

    return changed


def patch_file(path: Path, env_path: str = ".env") -> bool:
    with path.open("r") as f:
        data = yaml.load(f)

    if not isinstance(data, dict) or "services" not in data:
        print(f"[skip] {path}: no services section")
        return False

    services = data["services"]
    changed = False

    for service_name, service_def in services.items():
        if not isinstance(service_def, dict):
            print(f"[skip] {path}: service '{service_name}' is not a mapping")
            continue

        if ensure_env_file(service_def, env_path):
            changed = True
            print(f"[patched] {path}: service '{service_name}'")

    if changed:
        with path.open("w") as f:
            yaml.dump(data, f)

    return changed


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} FILE [FILE ...]")
        sys.exit(1)

    paths = [Path(arg) for arg in sys.argv[1:]]
    any_changed = False

    for path in paths:
        if not path.exists():
            print(f"[error] {path}: file does not exist")
            continue

        try:
            changed = patch_file(path)
            if changed:
                any_changed = True
            else:
                print(f"[ok] {path}: already patched or nothing to do")
        except Exception as e:
            print(f"[error] {path}: {e}")

    sys.exit(0 if any_changed else 0)


if __name__ == "__main__":
    main()