#!/usr/bin/env python3
"""splice-ai — SPLICE AI engine CLI.

Subcommands:
  suggest "<full cmdline>"   print corrected cmdline (rc 0) or nothing (rc 1)
  ask "<natural language>"   propose ONE shell command (context-aware)
  chat "<question>"          concise coding/help answer (context-aware)
  explain "<cmdline>"        one-paragraph explanation of a command

Token thrift: `suggest` is local-first — cache, then a curated typo table,
then a zero-token difflib fuzzy match over PATH executables + bash builtins/
keywords + the shell's aliases/functions (env SPLICE_SHELL_CMDS). Only then,
and only if allowed (settings.ai_enabled and SPLICE_NO_AI unset), does it make
a cheap haiku call. Outcomes (including negatives) are cached so repeat typos
cost zero tokens.
"""
import argparse
import difflib
import json
import os
import re
import shlex
import shutil
import subprocess
import sys


def _bootstrap_share():
    env = os.environ.get("SPLICE_SHARE")
    if env and os.path.isdir(env):
        return os.path.abspath(env)
    here = os.path.dirname(os.path.abspath(sys.argv[0]))
    dev = os.path.normpath(os.path.join(here, "..", "share", "splice"))
    if os.path.isfile(os.path.join(dev, "splice_common.py")):
        return dev
    return "/usr/share/splice"


SHARE = _bootstrap_share()
if SHARE not in sys.path:
    sys.path.insert(0, SHARE)
import splice_common as sc  # noqa: E402

WRAPPERS = {
    "sudo", "doas", "env", "nohup", "time", "command", "builtin", "exec",
    "stdbuf", "nice", "ionice", "setsid", "watch",
}

# Curated typo table: bad token -> replacement (checked against installed
# commands before use, so we never suggest something that can't run).
TYPO_TABLE = {
    "ap-get": "apt-get",
    "apt-gte": "apt-get",
    "atp-get": "apt-get",
    "apt-egt": "apt-get",
    "sl": "ls",
    "gti": "git",
    "igt": "git",
    "gt": "git",
    "grpe": "grep",
    "gerp": "grep",
    "greo": "grep",
    "dokcer": "docker",
    "docekr": "docker",
    "dcoker": "docker",
    "pyhton": "python3",
    "pyhton3": "python3",
    "pytohn": "python3",
    "pytohn3": "python3",
    "pythno": "python3",
    "sudp": "sudo",
    "suod": "sudo",
    "sduo": "sudo",
    "sudoo": "sudo",
    "cd..": "cd ..",
    "claer": "clear",
    "celar": "clear",
    "clera": "clear",
    "mkae": "make",
    "amke": "make",
    "maek": "make",
    "vmi": "vim",
    "ivm": "vim",
    "systemclt": "systemctl",
    "systmectl": "systemctl",
    "sytemctl": "systemctl",
    "systemctll": "systemctl",
    "chmdo": "chmod",
    "chomd": "chmod",
    "chwon": "chown",
    "tial": "tail",
    "tali": "tail",
    "haed": "head",
    "crul": "curl",
    "curll": "curl",
    "cutl": "curl",
    "wgte": "wget",
    "wegt": "wget",
    "pign": "ping",
    "pnig": "ping",
    "exti": "exit",
    "eixt": "exit",
    "kilall": "killall",
    "sssh": "ssh",
    "shh": "ssh",
    "mikdir": "mkdir",
    "mkdri": "mkdir",
    "touhc": "touch",
    "ehco": "echo",
    "ecoh": "echo",
    "whcih": "which",
    "hsitory": "history",
    "jounralctl": "journalctl",
    "journlactl": "journalctl",
}

CACHE_MAX = 500
ASSIGN_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=")
REFUSAL_PREFIXES = (
    "i ", "i'", "i can", "sorry", "unable", "cannot", "can't",
    "as an ai", "there is no", "unfortunately",
)

# ---------------------------------------------------------------------------
# suggest machinery
# ---------------------------------------------------------------------------
_BK = None  # lazy builtins+keywords set


def _builtins_keywords():
    global _BK
    if _BK is None:
        try:
            r = subprocess.run(
                ["bash", "-c", "compgen -b; compgen -k"],
                capture_output=True, text=True, timeout=5,
            )
            _BK = set(r.stdout.split())
        except (OSError, subprocess.TimeoutExpired):
            _BK = set()
    return _BK


def _shell_cmds():
    raw = os.environ.get("SPLICE_SHELL_CMDS", "")
    return {w for w in raw.split(":") if w}


def _path_commands():
    """All executable names on PATH. Built lazily, only on the fuzzy path."""
    names = set()
    for d in os.environ.get("PATH", "").split(os.pathsep):
        if not d:
            continue
        try:
            with os.scandir(d) as it:
                for e in it:
                    try:
                        if e.is_file(follow_symlinks=True) and os.access(e.path, os.X_OK):
                            names.add(e.name)
                    except OSError:
                        continue
        except OSError:
            continue
    return names


def _known(word, shell_cmds):
    return bool(shutil.which(word)) or word in _builtins_keywords() or word in shell_cmds


def _candidate(line):
    """The word that should name a command: skip leading VAR= assignments and
    wrapper words (plus the wrapper's options / env assignments)."""
    try:
        toks = shlex.split(line)
    except ValueError:
        toks = line.split()
    i, n = 0, len(toks)
    while i < n:
        t = toks[i]
        if ASSIGN_RE.match(t):
            i += 1
            continue
        if t in WRAPPERS:
            i += 1
            while i < n and (toks[i].startswith("-") or ASSIGN_RE.match(toks[i])):
                i += 1
            continue
        return t
    return None


def _replace_token(line, bad, repl):
    """Replace the first standalone occurrence of `bad` in `line`."""
    pat = r"(?<!\S)" + re.escape(bad) + r"(?!\S)"
    new, n = re.subn(pat, lambda m: repl, line, count=1)
    return new if n else None


def _cache_path():
    return os.path.join(sc.cache_dir(), "suggest.json")


def _load_cache():
    try:
        with open(_cache_path()) as f:
            d = json.load(f)
        if isinstance(d, dict) and isinstance(d.get("entries"), dict):
            return d
    except (OSError, ValueError):
        pass
    return {"version": 1, "entries": {}}


def _save_cache(cache):
    ents = cache["entries"]
    if len(ents) > CACHE_MAX:
        for k in list(ents)[: len(ents) - CACHE_MAX]:
            del ents[k]
    try:
        sc.atomic_write_json(_cache_path(), cache)
    except OSError:
        pass


def _clean_model_line(out):
    """Normalize a command-only model reply: strip fences/backticks, first line."""
    out = out.strip()
    if out.startswith("```"):
        lines = [l for l in out.splitlines() if not l.strip().startswith("```")]
        out = "\n".join(lines).strip()
    if len(out) >= 2 and out[0] == "`" and out[-1] == "`":
        out = out[1:-1].strip()
    for line in out.splitlines():
        if line.strip():
            return line.strip()
    return ""


def _ai_suggest(line, settings):
    prompt = (
        "Fix the typo in this shell command. Reply with ONLY the corrected "
        "command, or NONE.\nCommand: " + line
    )
    out = sc.run_claude(
        prompt,
        model=settings.get("fix_model") or "haiku",
        timeout=settings.get("suggest_timeout_s", 12),
    )
    if not out:
        return None
    ans = _clean_model_line(out)
    if not ans or ans.upper() == "NONE":
        return None
    if ans == line or len(ans) > 400:
        return None
    if ans.lower().startswith(REFUSAL_PREFIXES):
        return None
    return ans


def cmd_suggest(line):
    line = line.strip()
    if not line:
        return 1
    settings = sc.load_settings()
    ai_allowed = bool(settings.get("ai_enabled", True)) and not os.environ.get("SPLICE_NO_AI")

    # (1) cache
    cache = _load_cache()
    ent = cache["entries"].get(line)
    if isinstance(ent, dict):
        s = ent.get("s")
        if s:
            print(s)
            return 0
        # negative entry: honor it unless it was local-only and AI is now allowed
        if ent.get("ai") or not ai_allowed:
            return 1

    # (2) local: candidate word
    cand = _candidate(line)
    if not cand:
        return 1
    shell_cmds = _shell_cmds()
    if _known(cand, shell_cmds):
        return 1  # nothing to fix

    fixed = None
    # (2a) curated typo table (verify the replacement actually resolves)
    repl = TYPO_TABLE.get(cand)
    if repl and _known(repl.split()[0], shell_cmds):
        fixed = _replace_token(line, cand, repl)

    # (2b) difflib fuzzy over PATH + builtins + keywords + shell cmds
    if fixed is None:
        names = _path_commands() | _builtins_keywords() | shell_cmds
        names.discard(cand)
        cutoff = 0.75 if len(cand) >= 5 else 0.8
        m = difflib.get_close_matches(cand, names, n=1, cutoff=cutoff)
        if m:
            fixed = _replace_token(line, cand, m[0])

    # (3) model fallback, gated
    used_ai = False
    if fixed is None and ai_allowed:
        used_ai = True
        fixed = _ai_suggest(line, settings)

    # (4) cache the outcome (negatives too, so misses don't re-bill)
    if fixed and fixed != line:
        cache["entries"].pop(line, None)
        cache["entries"][line] = {"s": fixed, "ai": used_ai}
        _save_cache(cache)
        print(fixed)
        return 0
    cache["entries"].pop(line, None)
    cache["entries"][line] = {"s": None, "ai": used_ai}
    _save_cache(cache)
    return 1


# ---------------------------------------------------------------------------
# context-aware commands
# ---------------------------------------------------------------------------
def _ctx_blob():
    """Enabled-context blob via `splice-ctx blob` (local, token-budgeted)."""
    here = os.path.dirname(os.path.abspath(sys.argv[0]))
    exe = os.path.join(here, "splice-ctx")
    if not (os.path.isfile(exe) and os.access(exe, os.X_OK)):
        exe = shutil.which("splice-ctx")
    if not exe:
        return ""
    try:
        r = subprocess.run([exe, "blob"], capture_output=True, text=True, timeout=20)
    except (OSError, subprocess.TimeoutExpired):
        return ""
    if r.returncode != 0:
        return ""
    return r.stdout.strip()


def _fail_model():
    print(
        "splice-ai: model call failed (is the `claude` CLI installed and logged in?)",
        file=sys.stderr,
    )
    return 1


def cmd_ask(query):
    settings = sc.load_settings()
    blob = _ctx_blob()
    prompt = (blob + "\n\n") if blob else ""
    prompt += (
        "Task: " + query
        + "\nReply with ONLY the shell command, no backticks, no prose."
    )
    out = sc.run_claude(prompt, model=settings.get("ask_model") or None, timeout=90)
    if not out:
        return _fail_model()
    line = _clean_model_line(out)
    if not line:
        return _fail_model()
    print(line)
    return 0


def cmd_chat(question):
    settings = sc.load_settings()
    blob = _ctx_blob()
    prompt = (blob + "\n\n") if blob else ""
    prompt += question + "\n\nBe concise. No preamble."
    out = sc.run_claude(prompt, model=settings.get("ask_model") or None, timeout=120)
    if not out:
        return _fail_model()
    print(out)
    return 0


def cmd_explain(line):
    settings = sc.load_settings()
    prompt = (
        "Explain this shell command in one short paragraph. Be concise, no "
        "preamble.\nCommand: " + line
    )
    out = sc.run_claude(prompt, model=settings.get("fix_model") or "haiku", timeout=60)
    if not out:
        return _fail_model()
    print(out)
    return 0


def main():
    p = argparse.ArgumentParser(prog="splice-ai", description="SPLICE AI engine")
    sub = p.add_subparsers(dest="cmd", required=True)
    sub.add_parser("suggest", help="fix a mistyped command line").add_argument("cmdline")
    sub.add_parser("ask", help="propose one shell command").add_argument("query")
    sub.add_parser("chat", help="concise context-aware answer").add_argument("question")
    sub.add_parser("explain", help="explain a command").add_argument("cmdline")
    a = p.parse_args()
    if a.cmd == "suggest":
        return cmd_suggest(a.cmdline)
    if a.cmd == "ask":
        return cmd_ask(a.query)
    if a.cmd == "chat":
        return cmd_chat(a.question)
    return cmd_explain(a.cmdline)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(130)
