profile_switch.py

"""
Switch QueryForge runtime profile in .env.

Examples:
    python profile_switch.py --profile local --conversation on
    python profile_switch.py --profile rcs --api-key <key> --conversation off
    python profile_switch.py --profile local --dry-run
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path
from typing import Dict, List, Tuple


ENV_PATTERN = re.compile(r"^\s*([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(.*)\s*$")


def parse_env_lines(lines: List[str]) -> Dict[str, str]:
    """Parse KEY=VALUE lines from .env text, last key wins."""
    values: Dict[str, str] = {}
    for line in lines:
        match = ENV_PATTERN.match(line)
        if match:
            values[match.group(1)] = match.group(2).strip()
    return values


def apply_updates(lines: List[str], updates: Dict[str, str]) -> List[str]:
    """
    Apply updates to KEY=VALUE lines while preserving comments/order where possible.
    Missing keys are appended at the end.
    """
    output = list(lines)
    key_to_idx: Dict[str, int] = {}

    for idx, line in enumerate(output):
        match = ENV_PATTERN.match(line)
        if match:
            key_to_idx[match.group(1)] = idx

    for key, value in updates.items():
        new_line = f"{key}={value}"
        if key in key_to_idx:
            output[key_to_idx[key]] = new_line
        else:
            if output and output[-1].strip():
                output.append("")
            output.append(new_line)
            key_to_idx[key] = len(output) - 1

    return output


def mask_value(key: str, value: str) -> str:
    """Mask secrets when printing planned changes."""
    if "KEY" in key.upper() and value:
        return f"<set:{len(value)} chars>"
    return value if value else "<empty>"


def build_updates(args, current: Dict[str, str]) -> Tuple[Dict[str, str], List[str]]:
    """Build .env updates and warnings based on selected profile."""
    updates: Dict[str, str] = {}
    warnings: List[str] = []

    if args.profile == "local":
        updates["OLLAMA_HOST"] = args.local_host
        updates["OLLAMA_MODEL"] = args.model or "llama3.2"
        if args.clear_key_on_local:
            updates["OLLAMA_API_KEY"] = ""
    else:
        updates["OLLAMA_HOST"] = args.rcs_host
        updates["OLLAMA_MODEL"] = args.model or "llama3.3"
        api_key = args.api_key or current.get("OLLAMA_API_KEY") or current.get("RCS_API_KEY") or ""
        if api_key:
            updates["OLLAMA_API_KEY"] = api_key
        else:
            warnings.append(
                "No API key found for RCS profile. Set --api-key or add OLLAMA_API_KEY in .env."
            )

    if args.conversation != "keep":
        updates["ENABLE_CONVERSATION"] = "true" if args.conversation == "on" else "false"

    return updates, warnings


def main() -> int:
    parser = argparse.ArgumentParser(description="Switch .env profile for local vs Carleton RCS usage.")
    parser.add_argument("--profile", choices=["local", "rcs"], required=True, help="Target runtime profile")
    parser.add_argument("--env-file", default=".env", help="Path to .env file")
    parser.add_argument("--model", help="Override model for selected profile")
    parser.add_argument("--api-key", help="API key for RCS profile (writes OLLAMA_API_KEY)")
    parser.add_argument("--rcs-host", default="https://rcsllm.carleton.ca/rcsapi", help="RCS host URL")
    parser.add_argument("--local-host", default="http://localhost:11434", help="Local Ollama host URL")
    parser.add_argument(
        "--conversation",
        choices=["on", "off", "keep"],
        default="keep",
        help="Toggle ENABLE_CONVERSATION flag",
    )
    parser.add_argument(
        "--clear-key-on-local",
        action="store_true",
        help="When switching to local, clear OLLAMA_API_KEY in .env",
    )
    parser.add_argument("--dry-run", action="store_true", help="Show changes without writing file")
    args = parser.parse_args()

    env_path = Path(args.env_file)
    if not env_path.exists():
        print(f"Error: {env_path} does not exist.")
        return 1

    lines = env_path.read_text(encoding="utf-8").splitlines()
    current = parse_env_lines(lines)
    updates, warnings = build_updates(args, current)

    if warnings:
        for warning in warnings:
            print(f"Warning: {warning}")

    if not updates:
        print("No changes needed.")
        return 0

    print("Planned updates:")
    for key in sorted(updates.keys()):
        before = current.get(key, "")
        after = updates[key]
        print(f"  {key}: {mask_value(key, before)} -> {mask_value(key, after)}")

    if args.dry_run:
        print("\nDry-run only. No file changes were written.")
        return 0

    new_lines = apply_updates(lines, updates)
    env_path.write_text("\n".join(new_lines) + "\n", encoding="utf-8")
    print(f"\nUpdated {env_path}.")
    return 0


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