alfred/scripts/bootstrap.py

import secrets
from pathlib import Path

import tomllib
from config_loader import load_build_config, write_env_make


def generate_secret(rule: str) -> str:
    """
    Generates a cryptographically secure secret based on a spec string.
    Example specs: '32:b64', '16:hex'.
    """
    chunks: list[str] = rule.split(":")
    size: int = int(chunks[0])
    tech: str = chunks[1]

    if tech == "b64":
        return secrets.token_urlsafe(size)
    elif tech == "hex":
        return secrets.token_hex(size)
    else:
        raise ValueError(f"Invalid security format: {tech}")


def extract_python_version(version_string: str) -> tuple[str, str]:
    """
    Extract Python version from poetry dependency string.
    Examples:
        "==3.14.2" -> ("3.14.2", "3.14")
        "^3.14.2" -> ("3.14.2", "3.14")
        "~3.14.2" -> ("3.14.2", "3.14")
        "3.14.2" -> ("3.14.2", "3.14")
    """
    import re  # noqa: PLC0415

    # Remove poetry version operators (==, ^, ~, >=, etc.)
    clean_version = re.sub(r"^[=^~><]+", "", version_string.strip())

    # Extract version parts
    parts = clean_version.split(".")

    if len(parts) >= 2:
        full_version = clean_version
        short_version = f"{parts[0]}.{parts[1]}"
        return full_version, short_version
    else:
        raise ValueError(f"Invalid Python version format: {version_string}")


# TODO: Refactor
def bootstrap():  # noqa: PLR0912, PLR0915
    """
    Initializes the .env file by merging .env.example with generated secrets
    and build variables from pyproject.toml.
    Also generates .env.make for Makefile.

    ALWAYS preserves existing secrets!
    """
    base_dir = Path(__file__).resolve().parent.parent
    env_path = base_dir / ".env"

    example_path = base_dir / ".env.example"
    if not example_path.exists():
        print(f"❌ {example_path.name} not found.")
        return

    toml_path = base_dir / "pyproject.toml"
    if not toml_path.exists():
        print(f"❌ {toml_path.name} not found.")
        return

    # ALWAYS load existing .env if it exists
    existing_env = {}
    if env_path.exists():
        print("🔄 Reading existing .env...")
        with open(env_path) as f:
            for line in f:
                if "=" in line and not line.strip().startswith("#"):
                    key, value = line.split("=", 1)
                    existing_env[key.strip()] = value.strip()
        print(f"   Found {len(existing_env)} existing keys")
        print("🔧 Updating .env file (keeping secrets)...")
    else:
        print("🔧 Initializing: Creating secure .env file...")

    # Load data from pyproject.toml
    with open(toml_path, "rb") as f:
        data = tomllib.load(f)
        security_keys = data["tool"]["alfred"]["security"]
        settings_keys = data["tool"]["alfred"]["settings"]
        dependencies = data["tool"]["poetry"]["dependencies"]
        alfred_version = data["tool"]["poetry"]["version"]

    # Normalize TOML keys to UPPER_CASE for .env format (done once)
    security_keys_upper = {k.upper(): v for k, v in security_keys.items()}
    settings_keys_upper = {k.upper(): v for k, v in settings_keys.items()}

    # Extract Python version
    python_version_full, python_version_short = extract_python_version(
        dependencies["python"]
    )

    # Read .env.example
    with open(example_path) as f:
        example_lines = f.readlines()

    new_lines = []
    # Process each line from .env.example
    for raw_line in example_lines:
        line = raw_line.strip()

        if line and not line.startswith("#") and "=" in line:
            key, value = line.split("=", 1)
            key = key.strip()

            # Check if key exists in current .env (update mode)
            if key in existing_env:
                # Keep existing value for secrets
                if key in security_keys_upper:
                    new_lines.append(f"{key}={existing_env[key]}\n")
                    print(f"  ↻ Kept existing {key}")
                # Update build vars from pyproject.toml
                elif key in settings_keys_upper:
                    new_value = settings_keys_upper[key]
                    if existing_env[key] != new_value:
                        new_lines.append(f"{key}={new_value}\n")
                        print(f"  ↻ Updated {key}: {existing_env[key]} → {new_value}")
                    else:
                        new_lines.append(f"{key}={existing_env[key]}\n")
                        print(f"  ↻ Kept {key}={existing_env[key]}")
                # Update Python versions
                elif key == "PYTHON_VERSION":
                    if existing_env[key] != python_version_full:
                        new_lines.append(f"{key}={python_version_full}\n")
                        print(
                            f"  ↻ Updated Python: {existing_env[key]} → {python_version_full}"
                        )
                    else:
                        new_lines.append(f"{key}={existing_env[key]}\n")
                        print(f"  ↻ Kept Python: {existing_env[key]}")
                elif key == "PYTHON_VERSION_SHORT":
                    if existing_env[key] != python_version_short:
                        new_lines.append(f"{key}={python_version_short}\n")
                        print(
                            f"  ↻ Updated Python (short): {existing_env[key]} → {python_version_short}"
                        )
                    else:
                        new_lines.append(f"{key}={existing_env[key]}\n")
                        print(f"  ↻ Kept Python (short): {existing_env[key]}")
                elif key == "ALFRED_VERSION":
                    if existing_env.get(key) != alfred_version:
                        new_lines.append(f"{key}={alfred_version}\n")
                        print(
                            f"  ↻ Updated Alfred version: {existing_env.get(key, 'N/A')} → {alfred_version}"
                        )
                    else:
                        new_lines.append(f"{key}={alfred_version}\n")
                        print(f"  ↻ Kept Alfred version: {alfred_version}")
                # Keep other existing values
                else:
                    new_lines.append(f"{key}={existing_env[key]}\n")
            # Key doesn't exist, generate/add it
            elif key in security_keys_upper:
                rule = security_keys_upper[key]
                secret = generate_secret(rule)
                new_lines.append(f"{key}={secret}\n")
                print(f"  + Secret generated for {key} ({rule})")
            elif key in settings_keys_upper:
                value = settings_keys_upper[key]
                new_lines.append(f"{key}={value}\n")
                print(f"  + Setting added: {key}={value}")
            elif key == "PYTHON_VERSION":
                new_lines.append(f"{key}={python_version_full}\n")
                print(f"  + Python version: {python_version_full}")
            elif key == "PYTHON_VERSION_SHORT":
                new_lines.append(f"{key}={python_version_short}\n")
                print(f"  + Python version (short): {python_version_short}")
            elif key == "ALFRED_VERSION":
                new_lines.append(f"{key}={alfred_version}\n")
                print(f"  + Alfred version: {alfred_version}")
            else:
                new_lines.append(raw_line)
        else:
            # Keep comments and empty lines
            new_lines.append(raw_line)

    # Compute database URIs from the generated values
    final_env = {}
    for line in new_lines:
        if "=" in line and not line.strip().startswith("#"):
            key, value = line.split("=", 1)
            final_env[key.strip()] = value.strip()

    # Compute MONGO_URI
    if "MONGO_USER" in final_env and "MONGO_PASSWORD" in final_env:
        mongo_uri = (
            f"mongodb://{final_env.get('MONGO_USER', 'alfred')}:"
            f"{final_env.get('MONGO_PASSWORD', '')}@"
            f"{final_env.get('MONGO_HOST', 'mongodb')}:"
            f"{final_env.get('MONGO_PORT', '27017')}/"
            f"{final_env.get('MONGO_DB_NAME', 'alfred')}?authSource=admin"
        )
        # Update MONGO_URI in new_lines
        for i, line in enumerate(new_lines):
            if line.startswith("MONGO_URI="):
                new_lines[i] = f"MONGO_URI={mongo_uri}\n"
                print("  ✓ Computed MONGO_URI")
                break

    # Compute POSTGRES_URI
    if "POSTGRES_USER" in final_env and "POSTGRES_PASSWORD" in final_env:
        postgres_uri = (
            f"postgresql://{final_env.get('POSTGRES_USER', 'alfred')}:"
            f"{final_env.get('POSTGRES_PASSWORD', '')}@"
            f"{final_env.get('POSTGRES_HOST', 'vectordb')}:"
            f"{final_env.get('POSTGRES_PORT', '5432')}/"
            f"{final_env.get('POSTGRES_DB_NAME', 'alfred')}"
        )
        # Update POSTGRES_URI in new_lines
        for i, line in enumerate(new_lines):
            if line.startswith("POSTGRES_URI="):
                new_lines[i] = f"POSTGRES_URI={postgres_uri}\n"
                print("  ✓ Computed POSTGRES_URI")
                break

    # Write .env file
    with open(env_path, "w", encoding="utf-8") as f:
        f.writelines(new_lines)
    print(f"\n✅ {env_path.name} generated successfully.")

    # Generate .env.make for Makefile using shared config loader
    config = load_build_config(base_dir)
    write_env_make(config, base_dir)
    print("✅ .env.make generated for Makefile.")
    print("\n⚠️  Reminder: Please manually add your API keys to the .env file.")


if __name__ == "__main__":
    bootstrap()