#!/usr/bin/env python3
"""
nlsh — Natural Language Shell
Usage:
  nlsh "search for text in this folder and write to new file"
  nlsh -v "list all python files"   # verbose flat output (no TUI)
  nlsh --dry-run "delete tmp files" # show commands without running
"""

import argparse
import concurrent.futures
import html.parser
import json
import os
import subprocess
import sys
import urllib.error
import urllib.request

import anthropic

# ── optional clipboard support ──────────────────────────────────────────────
try:
    import pyperclip
    HAS_CLIPBOARD = True
except ImportError:
    HAS_CLIPBOARD = False

# ── optional curses TUI ──────────────────────────────────────────────────────
try:
    import curses
    HAS_CURSES = True
except ImportError:
    HAS_CURSES = False


# ── Extra context resolution ────────────────────────────────────────────────

_URL_FETCH_LIMIT = 8_000  # chars to pass to the model from a fetched page


class _HTMLTextExtractor(html.parser.HTMLParser):
    """Strip tags, collapse whitespace."""
    _SKIP = {"script", "style", "head", "nav", "footer"}

    def __init__(self):
        super().__init__()
        self._skip_depth = 0
        self.parts: list[str] = []

    def handle_starttag(self, tag, attrs):
        if tag in self._SKIP:
            self._skip_depth += 1

    def handle_endtag(self, tag):
        if tag in self._SKIP and self._skip_depth:
            self._skip_depth -= 1

    def handle_data(self, data):
        if not self._skip_depth:
            stripped = data.strip()
            if stripped:
                self.parts.append(stripped)


_PLAIN_TEXT_LIMIT = 8_000  # chars — same budget as URLs


def _fetch_url(url: str) -> tuple[str, str | None]:
    """Return (text, warning_or_None). Never raises."""
    req = urllib.request.Request(url, headers={"User-Agent": "nlsh/1.0"})
    try:
        with urllib.request.urlopen(req, timeout=8) as resp:
            content_type = resp.headers.get_content_type()
            raw = resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as e:
        return "", f"HTTP {e.code} fetching {url}"
    except urllib.error.URLError as e:
        return "", f"Could not reach {url}: {e.reason}"
    except Exception as e:
        return "", f"Unexpected error fetching {url}: {e}"

    if "html" in content_type:
        p = _HTMLTextExtractor()
        p.feed(raw)
        text = "\n".join(p.parts)
    else:
        text = raw

    if not text.strip():
        return "", f"Got empty content from {url} (login wall or access denied?)"

    warning = None
    if len(text) > _URL_FETCH_LIMIT:
        text = text[:_URL_FETCH_LIMIT] + "\n[… truncated]"
        warning = f"Fetched content from {url} was large; using first {_URL_FETCH_LIMIT} chars"

    return text, warning


def resolve_extra_context(items: list[str]) -> str:
    """Turn -c values into labelled context strings. Warns on bad input, never aborts."""
    urls, texts = [], []
    for item in items:
        item = item.strip()
        if not item:
            continue
        if item.startswith("http://") or item.startswith("https://"):
            urls.append(item)
        else:
            if len(item) > _PLAIN_TEXT_LIMIT:
                item = item[:_PLAIN_TEXT_LIMIT] + " [… truncated]"
                print(f"Warning: -c text was very long; truncated to {_PLAIN_TEXT_LIMIT} chars",
                      file=sys.stderr)
            texts.append(f"[note] {item}")

    parts = list(texts)

    if urls:
        print(f"Fetching {len(urls)} URL(s)…", end="\r", flush=True)
        with concurrent.futures.ThreadPoolExecutor(max_workers=len(urls)) as ex:
            futures = {ex.submit(_fetch_url, url): url for url in urls}
            for future in concurrent.futures.as_completed(futures):
                url = futures[future]
                content, warning = future.result()
                if warning:
                    print(f"Warning: {warning}", file=sys.stderr)
                if content:
                    parts.append(f"[URL: {url}]\n{content}")
        print(" " * 70, end="\r", flush=True)

    return "\n\n".join(parts)


# ── Claude call ─────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """\
You are a shell command expert. Given a natural language description, return \
a JSON array of 3–5 command suggestions. Each item must have:
  "command":     the full shell command (string)
  "one_liner":   a short (≤8 words) description shown inline
  "explanation": a 2–4 line explanation of exactly what the command does,
                 broken into logical chunks with → separators

Return ONLY the JSON array, no markdown, no code fences."""

FIX_SYSTEM_PROMPT = """\
You are a shell debugging expert. Given a failed command and its error output, \
return a JSON array of 3–5 fix suggestions. Each item must have:
  "command":     the corrected shell command (string)
  "one_liner":   a short (≤8 words) description of what this fix does
  "explanation": what went wrong and why this fixes it, \
                 broken into logical chunks with → separators

Return ONLY the JSON array, no markdown, no code fences."""


def get_suggestions(query: str, context: str, extra: str = "",
                    fix_error: str = "") -> list[dict]:
    client = anthropic.Anthropic()

    if fix_error:
        system = FIX_SYSTEM_PROMPT
        user_msg = f"Context:\n{context}\n\nFailed command / error output:\n{fix_error}"
        if query:
            user_msg += f"\n\nAdditional description: {query}"
    else:
        system = SYSTEM_PROMPT
        user_msg = f"Context:\n{context}"
        if extra:
            user_msg += f"\n\nExtra context provided by user:\n{extra}"
        user_msg += f"\n\nRequest: {query}"

    response = client.messages.create(
        model="claude-haiku-4-5",
        max_tokens=1024,
        system=[{"type": "text", "text": system,
                 "cache_control": {"type": "ephemeral"}}],
        messages=[{"role": "user", "content": user_msg}],
    )

    if not response.content:
        raise ValueError(f"Empty response from API (stop_reason: {response.stop_reason})")

    raw = response.content[0].text.strip()

    # Strip markdown code fences if the model wrapped the JSON anyway
    if raw.startswith("```"):
        raw = raw.split("\n", 1)[-1]          # drop opening fence line
        raw = raw.rsplit("```", 1)[0].strip()  # drop closing fence

    try:
        return json.loads(raw)
    except json.JSONDecodeError as e:
        raise ValueError(
            f"Could not parse model response as JSON ({e})\n\nRaw response:\n{raw}"
        ) from e


def build_context() -> str:
    lines = [f"cwd: {os.getcwd()}"]

    def _git():
        try:
            r = subprocess.run(["git", "status", "--short"],
                               capture_output=True, text=True, timeout=2)
            if r.returncode == 0 and r.stdout.strip():
                return f"git status:\n{r.stdout.strip()}"
        except Exception:
            pass
        return ""

    def _ls():
        try:
            r = subprocess.run(["ls", "-1"],
                               capture_output=True, text=True, timeout=2)
            if r.returncode == 0:
                entries = r.stdout.strip().split("\n")[:20]
                return "files: " + ", ".join(entries)
        except Exception:
            pass
        return ""

    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex:
        git_f, ls_f = ex.submit(_git), ex.submit(_ls)
        for result in (git_f.result(), ls_f.result()):
            if result:
                lines.append(result)

    return "\n".join(lines)


# ── Flat / verbose output ────────────────────────────────────────────────────

def run_flat(suggestions: list[dict], dry_run: bool) -> None:
    for i, s in enumerate(suggestions, 1):
        print(f"\n[{i}] {s['command']}")
        print(f"    {s['one_liner']}")
        print(f"    {s['explanation']}")

    if dry_run:
        return

    print()
    choice = input("Run which? (number, or q to quit): ").strip()
    if choice.lower() == "q" or not choice.isdigit():
        return
    idx = int(choice) - 1
    if 0 <= idx < len(suggestions):
        _run_command(suggestions[idx]["command"])


# ── Interactive TUI (curses) ─────────────────────────────────────────────────

# ── Widget installer ─────────────────────────────────────────────────────────

_WIDGET_CONTENT = """\
# nlsh zsh widget — added by `nlsh --install-widget`
_nlsh_widget() {
  local query="$BUFFER"
  local tmpfile
  tmpfile=$(mktemp)

  BUFFER=""
  zle reset-prompt

  zle -I
  nlsh ${query:+"$query"} --print > "$tmpfile"
  local exit_code=$?

  if [[ $exit_code -eq 0 ]] && [[ -s "$tmpfile" ]]; then
    BUFFER=$(cat "$tmpfile")
    rm -f "$tmpfile"
    CURSOR=${#BUFFER}
    zle accept-line
  else
    BUFFER="$query"
    rm -f "$tmpfile"
    CURSOR=${#BUFFER}
    zle reset-prompt
  fi
}

zle -N _nlsh_widget
bindkey '^F' _nlsh_widget   # ctrl+f — change to taste
"""

_ZSHRC_MARKER = "# nlsh widget"


def install_widget() -> None:
    widget_dir = os.path.expanduser("~/.config/nlsh")
    widget_path = os.path.join(widget_dir, "widget.zsh")
    zshrc_path = os.path.expanduser("~/.zshrc")
    source_line = f'source "{widget_path}"  {_ZSHRC_MARKER}'

    os.makedirs(widget_dir, exist_ok=True)
    with open(widget_path, "w") as f:
        f.write(_WIDGET_CONTENT)
    print(f"Widget written to {widget_path}")

    # Check if already sourced
    try:
        zshrc = open(zshrc_path).read()
    except FileNotFoundError:
        zshrc = ""

    if _ZSHRC_MARKER in zshrc:
        print("~/.zshrc already sources the widget — nothing to do.")
    else:
        with open(zshrc_path, "a") as f:
            f.write(f"\n{source_line}\n")
        print(f"Added source line to {zshrc_path}")

    print("\nDone! Run:  source ~/.zshrc")
    print("Then press ctrl+f at any prompt to invoke nlsh.")
    print("(Change the keybinding by editing bindkey in the widget file.)")


# ── TUI drawing ───────────────────────────────────────────────────────────────

def _draw(stdscr, suggestions: list[dict], selected: int) -> None:
    stdscr.clear()
    h, w = stdscr.getmaxyx()

    # ── command list (commands only — no inline hints) ───────────────────────
    for i, s in enumerate(suggestions):
        prefix = "❯ " if i == selected else "  "
        cmd_col = curses.color_pair(2) | curses.A_BOLD if i == selected else curses.color_pair(1)

        line = f"{prefix}{s['command']}"
        if len(line) > w - 1:
            line = line[: w - 2] + "…"

        row = i
        if row >= h - 6:
            break

        try:
            stdscr.addstr(row, 0, line, cmd_col)
        except curses.error:
            pass

    # ── divider ──────────────────────────────────────────────────────────────
    div_row = len(suggestions) + 1
    if div_row < h - 2:
        try:
            stdscr.addstr(div_row - 1, 0, "─" * w, curses.color_pair(3))
        except curses.error:
            pass

    # ── hover panel: one-liner + explanation for selected item only ───────────
    s = suggestions[selected]
    exp_row = div_row
    # one-liner in dim colour
    try:
        stdscr.addstr(exp_row, 2, s["one_liner"][: w - 4], curses.color_pair(3))
        exp_row += 1
    except curses.error:
        pass
    # explanation broken on → separators
    for chunk in s["explanation"].split("→"):
        chunk = chunk.strip()
        if not chunk:
            continue
        if exp_row >= h - 2:
            break
        try:
            stdscr.addstr(exp_row, 2, chunk[: w - 4], curses.color_pair(1))
        except curses.error:
            pass
        exp_row += 1

    # ── keybindings ──────────────────────────────────────────────────────────
    keys = "  [↑↓] navigate   [↵] run   [y] copy   [q] quit"
    try:
        stdscr.addstr(h - 1, 0, keys[: w], curses.color_pair(3))
    except curses.error:
        pass

    stdscr.refresh()


def run_tui(suggestions: list[dict], dry_run: bool, print_only: bool = False) -> None:
    def _main(stdscr):
        curses.curs_set(0)
        curses.start_color()
        curses.use_default_colors()
        curses.init_pair(1, -1, -1)                          # default
        curses.init_pair(2, curses.COLOR_GREEN, -1)          # selected cmd
        curses.init_pair(3, curses.COLOR_WHITE, -1)  # dim / hints

        selected = 0
        chosen = None

        while True:
            _draw(stdscr, suggestions, selected)
            key = stdscr.getch()

            if key in (curses.KEY_UP, ord("k")) and selected > 0:
                selected -= 1
            elif key in (curses.KEY_DOWN, ord("j")) and selected < len(suggestions) - 1:
                selected += 1
            elif key in (curses.KEY_ENTER, 10, 13):
                chosen = ("run", suggestions[selected]["command"])
                break
            elif key == ord("y"):
                chosen = ("copy", suggestions[selected]["command"])
                break
            elif key in (ord("q"), 27):  # q or Esc
                break

        return chosen

    if print_only:
        # stdout is captured and stdin may be closed by the widget —
        # redirect both to /dev/tty so curses can draw and read keys,
        # then restore stdout to emit the chosen command
        original_stdout_fd = os.dup(1)
        original_stdin_fd = os.dup(0)
        tty_fd = os.open("/dev/tty", os.O_RDWR)
        os.dup2(tty_fd, 0)
        os.dup2(tty_fd, 1)
        os.close(tty_fd)
        result = curses.wrapper(_main)
        os.dup2(original_stdin_fd, 0)
        os.dup2(original_stdout_fd, 1)
        os.close(original_stdin_fd)
        os.close(original_stdout_fd)
    else:
        result = curses.wrapper(_main)

    if result is None:
        return
    action, cmd = result

    if action == "copy":
        if HAS_CLIPBOARD:
            pyperclip.copy(cmd)
            print(f"Copied to clipboard: {cmd}")
        else:
            print(f"Command (install pyperclip for clipboard): {cmd}")
    elif action == "run":
        if dry_run or print_only:
            print(cmd)
        else:
            _run_command(cmd)


# ── Shell execution ──────────────────────────────────────────────────────────

def _run_command(cmd: str) -> None:
    """Run cmd in the user's shell and inject it into history."""
    shell = os.environ.get("SHELL", "/bin/sh")

    # Write to history so it appears as if the user typed it
    hist_file = os.environ.get("HISTFILE", "")
    if not hist_file:
        if "zsh" in shell:
            hist_file = os.path.expanduser("~/.zsh_history")
        else:
            hist_file = os.path.expanduser("~/.bash_history")

    try:
        with open(hist_file, "a") as f:
            f.write(cmd + "\n")
    except Exception:
        pass

    print(f"$ {cmd}")
    os.execvp(shell, [shell, "-c", cmd])


# ── Entry point ──────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        prog="nlsh",
        description="Natural language → shell command suggestions",
    )
    parser.add_argument("query", nargs="?", help="Natural language description")
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Flat output, no interactive TUI")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show commands without executing")
    parser.add_argument("-c", "--context", action="append", default=[],
                        metavar="URL_OR_TEXT",
                        help="Extra context: a URL (fetched) or plain text. Repeatable.")
    parser.add_argument("--fix", action="store_true",
                        help="Fix mode: pipe in error output or describe what broke.")
    parser.add_argument("--print", action="store_true",
                        help="Print selected command to stdout instead of running it (for shell widgets).")
    parser.add_argument("--install-widget", action="store_true",
                        help="Install the zsh ctrl+f widget into ~/.config/nlsh/ and add it to ~/.zshrc.")
    args = parser.parse_args()

    if args.install_widget:
        install_widget()
        sys.exit(0)

    # ── API key check ─────────────────────────────────────────────────────────
    if not os.environ.get("ANTHROPIC_API_KEY"):
        print("Error: ANTHROPIC_API_KEY is not set.", file=sys.stderr)
        print("Get a key at https://console.anthropic.com and add to your shell:", file=sys.stderr)
        print('  export ANTHROPIC_API_KEY="sk-ant-..."', file=sys.stderr)
        sys.exit(1)

    # ── fix mode ─────────────────────────────────────────────────────────────
    fix_error = ""
    if args.fix:
        if not sys.stdin.isatty():
            # e.g. npm run build 2>&1 | nlsh --fix
            piped = sys.stdin.read().strip()
            if len(piped) > 6000:
                piped = piped[-6000:]  # keep the tail — errors are usually last
                print("(stdin truncated to last 6000 chars)", file=sys.stderr)
            fix_error = piped
        elif args.query:
            # nlsh --fix "error description"  (no pipe)
            fix_error = args.query
            args.query = ""  # don't also treat it as a normal query
        else:
            fix_error = input("Paste the error (ctrl+D when done):\n").strip()

        if not fix_error:
            print("No error provided.", file=sys.stderr)
            sys.exit(1)

    # ── normal mode ───────────────────────────────────────────────────────────
    query = args.query or ""
    if not args.fix and not query:
        with open("/dev/tty", "r") as tty:
            sys.stderr.write("What do you want to do? ")
            sys.stderr.flush()
            query = tty.readline().strip()
        if not query:
            sys.exit(0)

    print("Thinking…", end="\r", flush=True, file=sys.stderr)
    context = build_context()

    try:
        extra = resolve_extra_context(args.context) if args.context else ""
        suggestions = get_suggestions(query, context, extra, fix_error)
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    except anthropic.APIError as e:
        print(f"API error ({e.status_code}): {e.message}", file=sys.stderr)
        sys.exit(1)

    print("          ", end="\r", file=sys.stderr)  # clear "Thinking…"

    if args.verbose or not HAS_CURSES:
        run_flat(suggestions, args.dry_run or args.print)
    else:
        run_tui(suggestions, args.dry_run, print_only=args.print)


if __name__ == "__main__":
    main()
