alias_gen.py

#!/usr/bin/env python3
import os
import sys
import argparse
import datetime
import shutil
import json
import re
import shlex
from collections import Counter, deque
from openai import OpenAI

# module‐level client for OpenAI
client = None

MARKER_START = "# alias-gen start"
MARKER_END   = "# alias-gen end"
CONFIG_FILE  = os.path.expanduser("~/.alias_gen_config.json")

TRIVIAL_COMMANDS  = {"ls", "pwd", "echo", "clear", "exit", "git"}
IGNORE_PATTERNS   = {"brew install", "curl -fsSL", "alias_gen.py config", "alias_gen.py scan"}
SCRIPT_EXT        = re.compile(r'^[^/]+\.[^/]+$')
MIN_SINGLE_TOKENS = 2
MIN_SINGLE_LENGTH = 30
SPECIAL_WORKFLOWS = [("cd", ["source", "."])]

def detect_rc():
    home = os.path.expanduser("~")
    # Windows PowerShell profile
    if os.name == 'nt':
        profile = os.environ.get('PROFILE')
        if profile and os.path.exists(profile):
            return profile
        return os.path.join(home, "Documents", "PowerShell", "Microsoft.PowerShell_profile.ps1")
    # Unix-like shells
    shell = os.environ.get("SHELL", "").lower()
    if "zsh" in shell:
        return os.path.join(home, ".zshrc")
    if "fish" in shell:
        return os.path.join(home, ".config", "fish", "config.fish")
    bash_profile = os.path.join(home, ".bash_profile")
    if os.path.exists(bash_profile):
        return bash_profile
    return os.path.join(home, ".bashrc")

def detect_history_file():
    home = os.path.expanduser("~")
    # Windows PowerShell PSReadLine history
    if os.name == 'nt':
        appdata = os.environ.get('APPDATA')
        if appdata:
            return os.path.join(appdata, "Microsoft", "Windows", "PowerShell", "PSReadLine", "ConsoleHost_history.txt")
        return os.path.join(home, "AppData", "Roaming", "Microsoft", "Windows", "PowerShell", "PSReadLine", "ConsoleHost_history.txt")
    shell = os.environ.get("SHELL", "").lower()
    if "zsh" in shell:
        return os.path.join(home, ".zsh_history")
    if "fish" in shell:
        return os.path.join(home, ".local", "share", "fish", "fish_history")
    return os.path.join(home, ".bash_history")

def safe_split(cmd):
    try:
        return shlex.split(cmd)
    except:
        return cmd.strip().split()

def load_existing_alias_cmds():
    path = detect_rc()
    cmds = set()
    if not os.path.exists(path):
        return cmds
    # match Bash/Zsh alias and PowerShell Set-Alias
    alias_re = re.compile(r"^(?:alias|Set-Alias)\s+(\w+)[= ]+['\"]?(.+?)['\"]?$", re.IGNORECASE)
    func_re  = re.compile(r"^function\s+\w+\(\)\s*{\s*(.+?)\s*}")
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            line = line.strip()
            m = alias_re.match(line)
            if m:
                parts = safe_split(m.group(2))
                if parts:
                    cmds.add(parts[0])
                continue
            m2 = func_re.match(line)
            if m2:
                first = m2.group(1).split("&&")[0].strip()
                parts = safe_split(first)
                if parts:
                    cmds.add(parts[0])
    return cmds

def normalize_cmd(cmd):
    parts = safe_split(cmd)
    if parts and parts[0] == "cd" and len(parts) > 1:
        return f"cd {os.path.basename(parts[1])}"
    if parts and parts[0] in ("source", ".") and len(parts) > 1:
        return f"{parts[0]} {os.path.basename(parts[1])}"
    return cmd

def normalize_seq(seq):
    return tuple(normalize_cmd(c) for c in seq)

def load_config():
    if os.path.exists(CONFIG_FILE):
        try:
            return json.load(open(CONFIG_FILE, "r"))
        except:
            pass
    return {}

def save_config(config):
    with open(CONFIG_FILE, "w") as f:
        json.dump(config, f)
    print(f"Configuration saved to {CONFIG_FILE}")

def get_api_key(config):
    if os.environ.get("OPENAI_API_KEY"):
        return os.environ["OPENAI_API_KEY"]
    if "openai_api_key" in config:
        return config["openai_api_key"]
    key = input("Enter your OpenAI API key: ").strip()
    if not key:
        print("No API key; exiting.", file=sys.stderr)
        sys.exit(1)
    if input(f"Save to {CONFIG_FILE}? [y/N] ").strip().lower() == "y":
        config["openai_api_key"] = key
        save_config(config)
    return key

def choose_ui_mode(config):
    choice = input("Choose UI mode (terminal/tui/web) [terminal]: ").strip().lower()
    config["ui_mode"] = choice if choice in ("terminal","tui","web") else "terminal"
    save_config(config)

def load_history(path, max_lines=10000):
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            return [l.strip() for l in f if l.strip()][-max_lines:]
    except Exception as e:
        print(f"Error reading history: {e}", file=sys.stderr)
        sys.exit(1)

def find_patterns(cmds, window_sizes=(2,3), min_occurrence=3):
    existing = load_existing_alias_cmds()
    counts = Counter()
    for w in window_sizes:
        buf = deque(maxlen=w)
        for c in cmds:
            buf.append(c)
            if len(buf) == w:
                counts[tuple(buf)] += 1
    patterns = []
    for seq, cnt in counts.items():
        cmd0 = safe_split(seq[0])[0]
        if cmd0 in existing: continue
        text = " && ".join(seq)
        if any(ignore in text for ignore in IGNORE_PATTERNS): continue
        if len(seq) == 1:
            toks = safe_split(seq[0])
            if toks[0] in TRIVIAL_COMMANDS: continue
            if len(toks) >= MIN_SINGLE_TOKENS or len(seq[0]) > MIN_SINGLE_LENGTH:
                patterns.append(seq)
            continue
        if len(seq) == 2:
            t0, t1 = safe_split(seq[0])[0], safe_split(seq[1])[0]
            for s0, lst in SPECIAL_WORKFLOWS:
                if t0==s0 and t1 in lst:
                    patterns.append(seq); break
            if seq in patterns: continue
        if cnt < min_occurrence: continue
        bad = False
        for c in seq:
            for t in safe_split(c):
                if SCRIPT_EXT.match(t) and not t.startswith("/"):
                    bad = True; break
            if bad: break
        if bad: continue
        firsts = [safe_split(c)[0] for c in seq]
        if len(set(firsts)) < 2 or any(t in TRIVIAL_COMMANDS for t in firsts):
            continue
        patterns.append(seq)
    # dedupe normalized
    norm_map = {}
    for seq in patterns:
        n = normalize_seq(seq)
        if n not in norm_map:
            norm_map[n] = seq
    unique = list(norm_map.values())
    multi  = [p for p in unique if len(p)>1]
    single = [p for p in unique if len(p)==1]
    multi.sort(key=lambda s: counts[s], reverse=True)
    return multi[:5] + single

def prompt_for_aliases(client, patterns):
    prompt = (
        "You are an expert shell user. For each pattern, suggest one alias or function. "
        "- Multi-step (2+ cmds): function using \"$@\". "
        "- Single long cmd: short alias.\n"
        "Format:\nalias name='definition'  # desc\n"
        "or\nfunction name() { cmd1 && cmd2 \"$@\"; }  # desc\n\nPatterns:"
    )
    for seq in patterns:
        prompt += "\n- " + " && ".join(seq)
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role":"system","content":prompt}],
        temperature=0.2
    )
    return resp.choices[0].message.content.strip()

def backup_and_write(rc_path, new_block):
    ts = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    bak = f"{rc_path}.alias_gen.bak.{ts}"
    if os.path.exists(rc_path):
        shutil.copy(rc_path, bak)
        lines = open(rc_path, "r", encoding="utf-8", errors="ignore").readlines()
    else:
        lines = []
    before, old, after = [], [], []
    in_block = False; saw = False
    for l in lines:
        if l.strip() == MARKER_START:
            saw = True; in_block = True; continue
        if l.strip() == MARKER_END:
            in_block = False; continue
        if not saw:
            before.append(l)
        elif in_block:
            old.append(l.rstrip("\n"))
        else:
            after.append(l)
    merged = old + [l for l in new_block.splitlines() if l.strip()]
    seen = set(); unique = []
    for l in merged:
        if l not in seen:
            seen.add(l); unique.append(l)
    out = before + [MARKER_START+"\n"] + [l+"\n" for l in unique] + [MARKER_END+"\n"] + after
    with open(rc_path, "w", encoding="utf-8") as f:
        f.writelines(out)
    print(f"Updated {rc_path}, backup at {bak}")

def write_alias_block(rc_path, alias_lines):
    ts = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    bak = f"{rc_path}.alias_gen.bak.{ts}"
    if os.path.exists(rc_path):
        shutil.copy(rc_path, bak)
        lines = open(rc_path, "r", encoding="utf-8", errors="ignore").readlines()
    else:
        lines = []
    start = next((i for i,l in enumerate(lines) if l.strip()==MARKER_START), None)
    end   = next((i for i,l in enumerate(lines) if l.strip()==MARKER_END),   None)
    before = lines[:start+1] if start is not None else [MARKER_START+"\n"]
    after  = lines[end:]     if end   is not None else [MARKER_END+"\n"]
    out = before + [l+"\n" for l in alias_lines] + after
    with open(rc_path, "w", encoding="utf-8") as f:
        f.writelines(out)
    print(f"Edited {rc_path}, backup at {bak}")

def parse_alias(line):
    m1 = re.match(r"^(alias)\s+(\w+)=['\"](.+?)['\"]", line)
    if m1:
        return ("alias", m1.group(2), m1.group(3), "")
    m2 = re.match(r"^(function)\s+(\w+)\s*\(\)\s*{\s*(.+?)\s*}", line)
    if m2:
        return ("function", m2.group(2), m2.group(3), "")
    return (None, None, None, None)

def choose_aliases(block, ui_mode, patterns):
    return []  # placeholder, implement UI-specific selection

def manual_alias(ui_mode):
    pass  # placeholder for manual entry

def edit_aliases(ui_mode):
    global client
    rc_path = detect_rc()
    lines = open(rc_path, "r", encoding="utf-8", errors="ignore").readlines()
    start = next((i for i,l in enumerate(lines) if l.strip()==MARKER_START), None)
    end   = next((i for i,l in enumerate(lines) if l.strip()==MARKER_END),   None)
    current = []
    if start is not None and end is not None:
        for l in lines[start+1:end]:
            parsed = parse_alias(l.strip())
            if parsed[0]:
                current.append(l.strip())
    parsed = [parse_alias(l) for l in current]

    if ui_mode == "terminal":
        for i,(k,n,b,_) in enumerate(parsed,1):
            print(f"{i}. {k} {n} -> {b}")
        idx = int(input("Number to edit/delete: ").strip())-1
        k,n,b,c = parsed[idx]
        if input("Delete? [y/N]: ").strip().lower()=="y":
            parsed.pop(idx)
        else:
            new_n = input(f"Name [{n}]: ").strip() or n
            new_b = input(f"Cmd [{b}]: ").strip() or b
            parsed[idx] = (k,new_n,new_b,c)
        new_block = [
            f"{'alias' if k=='alias' else 'function'} {n}" +
            (f"='{b}'" if k=='alias' else f"() {{ {b} \"$@\"; }}")
            for k,n,b,c in parsed
        ]
        write_alias_block(rc_path, new_block)
        return

    # TUI and web modes omitted for brevity; keep your existing implementations here
    print("Edit in other UI modes not shown here.")
    sys.exit(0)

def main():
    global client
    parser = argparse.ArgumentParser()
    parser.add_argument("action", choices=("config","scan","install","manual","edit"))
    parser.add_argument("--history-file", dest="history_file", default=None)
    parser.add_argument("--ui", choices=("terminal","tui","web","gui"))
    args = parser.parse_args()

    config = load_config()
    if args.action == "config":
        get_api_key(config); choose_ui_mode(config); return

    key = get_api_key(config)
    client = OpenAI(api_key=key)

    raw_ui = args.ui or config.get("ui_mode","terminal")
    ui_mode = "web" if raw_ui=="gui" else raw_ui

    if args.action in ("scan","install"):
        path = args.history_file or detect_history_file()
        cmds = load_history(path)
        pats = find_patterns(cmds)
        if not pats:
            print("No patterns found.", file=sys.stderr)
            return
        block  = prompt_for_aliases(client, pats)
        if args.action == "scan":
            print(block)
            return
        chosen = choose_aliases(block, ui_mode, pats)
        backup_and_write(detect_rc(), "\n".join(chosen))
        return

    if args.action == "manual":
        manual_alias(ui_mode)
        return

    if args.action == "edit":
        edit_aliases(ui_mode)
        return

if __name__ == "__main__":
    main()