"""Claude Code style tools, every one scoped to a workspace (see workspaces.py).

Docstrings are the tool descriptions the model sees (Gradio builds the MCP schema from
them), so they are written for the model. Errors are returned as "error: ..." text.
"""

import fnmatch
import glob as _glob
import itertools
import json
import os
import re
import shlex
import shutil
import signal
import subprocess
import time
from collections import defaultdict
from pathlib import Path

import workspaces as W

MAX_OUTPUT = 30_000
# chat-ui aborts MCP calls after ~120s; stay below it and push long jobs to the background.
MAX_FG_TIMEOUT = 110
SKIP_DIRS = {".git", "node_modules", "__pycache__", ".venv", ".cache"}
_SECRET_HINTS = ("TOKEN", "SECRET", "KEY", "PASSWORD", "CREDENTIAL")


def _clip(text: str, limit: int = MAX_OUTPUT) -> str:
    if len(text) <= limit:
        return text
    half = limit // 2
    return f"{text[:half]}\n... [{len(text) - limit} characters truncated] ...\n{text[-half:]}"


def _ws(workspace: str) -> W.Workspace:
    return W.get(workspace)


def _numbered(lines: list[str], start: int = 1) -> str:
    return "\n".join(f"{i:>6}\t{ln[:2000]}" for i, ln in enumerate(lines, start))


def _guard(fn):
    """Turn workspace/IO errors into 'error: ...' text instead of tracebacks."""
    import functools

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except W.WorkspaceError as exc:
            return f"error: {exc}"
        except Exception as exc:  # noqa: BLE001 - surface any failure to the model
            return f"error: {type(exc).__name__}: {exc}"

    return wrapper


# --------------------------------------------------------------------------- workspaces


@_guard
def workspace_new(title: str = "") -> str:
    """Create a new workspace (a private project folder on the remote computer). Call this ONCE at the very start of every new conversation, before any other tool. Then pass the returned id as `workspace` in EVERY later tool call of this conversation, and repeat the id ('workspace: ws-xxxxxx') in your replies so it stays in the chat.

    Args:
        title: Short label for the workspace, e.g. the task being worked on.

    Returns:
        The workspace id and its working directory.
    """
    ws = W.create(title)
    return (
        f"Workspace created: {ws.id}\n"
        f"Working directory: {ws.work}\n"
        f"Rules: pass workspace=\"{ws.id}\" in every tool call of this conversation, and mention "
        f"`workspace: {ws.id}` in your replies. Prefer absolute paths."
    )


@_guard
def workspace_list() -> str:
    """List existing workspaces, most recently used first. Use it to recover the id of a workspace when this conversation has lost it, or to resume older work.

    Returns:
        One line per workspace: id, title, last used time.
    """
    rows = []
    for ws in W.list_all()[:50]:
        m = ws.meta()
        last = time.strftime("%Y-%m-%d %H:%M", time.localtime(m.get("last_used", 0)))
        rows.append(f"{ws.id}  last_used={last}  title={m.get('title') or '-'}")
    return "\n".join(rows) or "(no workspaces)"


@_guard
def workspace_delete(workspace: str) -> str:
    """Permanently delete a workspace and all its files. Only do this when the user asks.

    Args:
        workspace: Workspace id to delete.

    Returns:
        Confirmation.
    """
    ws = _ws(workspace)
    _kill_all_shells(ws.id)
    W.delete(ws.id)
    return f"deleted {ws.id}"


# --------------------------------------------------------------------------- Bash


class _Shell:
    def __init__(self, proc: subprocess.Popen, out_path: Path, command: str):
        self.proc, self.out_path, self.command = proc, out_path, command
        self.offset = 0


_SHELLS: dict[tuple[str, str], _Shell] = {}
_BG_COUNTER: defaultdict[str, "itertools.count[int]"] = defaultdict(lambda: itertools.count(1))


def _env(ws: W.Workspace) -> dict:
    env = {k: v for k, v in os.environ.items() if not any(h in k.upper() for h in _SECRET_HINTS)}
    env.update(CC_WORKSPACE=ws.id, CC_WORKDIR=str(ws.work), TERM="dumb", PAGER="cat", GIT_PAGER="cat")
    local_bin = str(Path.home() / ".local" / "bin")
    env["PATH"] = f"{local_bin}:{env.get('PATH', '')}"
    return env


def _kill_group(proc: subprocess.Popen) -> None:
    try:
        os.killpg(proc.pid, signal.SIGKILL)
    except ProcessLookupError:
        pass


def _kill_all_shells(ws_id: str) -> None:
    for key in [k for k in _SHELLS if k[0] == ws_id]:
        _kill_group(_SHELLS.pop(key).proc)


def _script(ws: W.Workspace, command: str) -> str:
    # Persist the final directory so `cd` carries over to the next Bash call.
    return f'{command}\n__cc_rc=$?\npwd -P > {shlex.quote(str(ws.cwd_file))}\nexit $__cc_rc\n'


@_guard
def Bash(workspace: str, command: str, timeout: int = 60, run_in_background: bool = False) -> str:
    """Run a bash command on the remote computer, inside the workspace. The working directory persists between calls (cd carries over); environment variables do not. Use it for git, package installs, builds, tests, scripts and anything else a terminal does. Prefer the dedicated Read, Write, Edit, Glob and Grep tools for file work. Commands must be non-interactive. For anything longer than ~100 seconds (servers, long builds) set run_in_background=true and poll with BashOutput.

    Args:
        workspace: Workspace id of this conversation (from workspace_new).
        command: The bash command to run.
        timeout: Seconds to wait before killing a foreground command (1-110).
        run_in_background: Start the command detached and return a bash_id immediately.

    Returns:
        Combined stdout+stderr and the exit code, or the bash_id of the background shell.
    """
    ws = _ws(workspace)
    timeout = max(1, min(int(timeout), MAX_FG_TIMEOUT))
    popen = dict(
        cwd=ws.cwd(), env=_env(ws), stdin=subprocess.DEVNULL, start_new_session=True,
        stderr=subprocess.STDOUT,
    )
    # Output goes to a file, not a pipe: with a pipe, `server &` keeps it open and the call
    # would hang until the timeout even though the shell already exited.
    if run_in_background:
        bash_id = f"bg{next(_BG_COUNTER[ws.id])}"  # never reused, even after KillShell
        out_path = ws.shells_dir / f"{bash_id}.out"
    else:
        out_path = ws.shells_dir / f"fg-{os.getpid()}-{time.monotonic_ns()}.out"
    with out_path.open("wb") as out:
        proc = subprocess.Popen(["bash", "-c", _script(ws, command)], stdout=out, **popen)
    if run_in_background:
        _SHELLS[(ws.id, bash_id)] = _Shell(proc, out_path, command)
        return f"started background shell {bash_id} (pid {proc.pid}). Use BashOutput to read it, KillShell to stop it."
    timed_out = False
    try:
        proc.wait(timeout=timeout)
    except subprocess.TimeoutExpired:
        _kill_group(proc)
        proc.wait()
        timed_out = True
    out = out_path.read_bytes().decode(errors="replace")
    out_path.unlink(missing_ok=True)
    if timed_out:
        return _clip(f"{out}\n[killed after {timeout}s timeout. Use run_in_background=true for long jobs.]")
    return _clip(f"{out}\n[exit code: {proc.returncode}]" if out else f"[exit code: {proc.returncode}]")


@_guard
def BashOutput(workspace: str, bash_id: str, filter: str = "") -> str:
    """Read NEW output (since the last read) from a background shell started with Bash(run_in_background=true), plus whether it is still running.

    Args:
        workspace: Workspace id of this conversation.
        bash_id: Id returned when the background shell was started, e.g. "bg1".
        filter: Optional regex; only matching lines are returned.

    Returns:
        New output and the shell status.
    """
    ws = _ws(workspace)
    sh = _SHELLS.get((ws.id, bash_id))
    if not sh:
        return f"error: no background shell '{bash_id}' in {ws.id}"
    with sh.out_path.open("rb") as f:
        f.seek(sh.offset)
        chunk = f.read()
    sh.offset += len(chunk)
    text = chunk.decode(errors="replace")
    if filter:
        rx = re.compile(filter)
        text = "\n".join(ln for ln in text.splitlines() if rx.search(ln))
    code = sh.proc.poll()
    status = "running" if code is None else f"exited with code {code}"
    return _clip(f"{text}\n[status: {status}]")


@_guard
def KillShell(workspace: str, shell_id: str) -> str:
    """Kill a background shell (and everything it started).

    Args:
        workspace: Workspace id of this conversation.
        shell_id: Id of the background shell, e.g. "bg1".

    Returns:
        Confirmation.
    """
    ws = _ws(workspace)
    sh = _SHELLS.pop((ws.id, shell_id), None)
    if not sh:
        return f"error: no background shell '{shell_id}' in {ws.id}"
    _kill_group(sh.proc)
    return f"killed {shell_id}"


# --------------------------------------------------------------------------- files


@_guard
def Read(workspace: str, file_path: str, offset: int = 1, limit: int = 2000) -> str:
    """Read a text file. Returns lines prefixed with line numbers (cat -n style). Always Read a file before you Edit it.

    Args:
        workspace: Workspace id of this conversation.
        file_path: Absolute path (relative paths resolve from the current directory).
        offset: 1-based line number to start from.
        limit: Maximum number of lines to return.

    Returns:
        The numbered lines.
    """
    ws = _ws(workspace)
    p = ws.path(file_path)
    if p.is_dir():
        return f"error: {p} is a directory; use LS"
    if not p.is_file():
        return f"error: file not found: {p}"
    raw = p.read_bytes()
    if b"\0" in raw[:8192]:
        return f"error: {p} looks binary ({len(raw)} bytes)"
    lines = raw.decode(errors="replace").splitlines()
    if not lines:
        return "(file is empty)"
    start = max(1, int(offset))
    chunk = lines[start - 1 : start - 1 + max(1, int(limit))]
    if not chunk:
        return f"error: offset {start} is past the end of the file ({len(lines)} lines)"
    end = start + len(chunk) - 1
    note = f"\n[showing lines {start}-{end} of {len(lines)}]" if end < len(lines) or start > 1 else ""
    return _clip(_numbered(chunk, start) + note)


@_guard
def Write(workspace: str, file_path: str, content: str) -> str:
    """Create a file or completely overwrite an existing one (parent folders are created). To change part of an existing file use Edit instead.

    Args:
        workspace: Workspace id of this conversation.
        file_path: Absolute path (relative paths resolve from the current directory).
        content: The full file content.

    Returns:
        Confirmation.
    """
    ws = _ws(workspace)
    p = ws.path(file_path)
    if p.is_dir():
        return f"error: {p} is a directory"
    existed = p.exists()
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(content, encoding="utf-8")
    return f"{'Overwrote' if existed else 'Created'} {p} ({content.count(chr(10)) + 1} lines, {len(content.encode())} bytes)"


def _apply_edit(text: str, old: str, new: str, replace_all: bool) -> str:
    if old == new:
        raise ValueError("old_string and new_string are identical")
    if not old:
        raise ValueError("old_string is empty (use Write to create a file)")
    count = text.count(old)
    if count == 0:
        raise ValueError("old_string not found. It must match exactly, including whitespace and indentation")
    if count > 1 and not replace_all:
        raise ValueError(
            f"old_string appears {count} times. Add surrounding context to make it unique, or set replace_all=true"
        )
    return text.replace(old, new) if replace_all else text.replace(old, new, 1)


@_guard
def Edit(workspace: str, file_path: str, old_string: str, new_string: str, replace_all: bool = False) -> str:
    """Replace an exact string in an existing file. old_string must match the file exactly (whitespace included) and be unique unless replace_all is true. Read the file first.

    Args:
        workspace: Workspace id of this conversation.
        file_path: Absolute path of the file to modify.
        old_string: The exact text to replace.
        new_string: The replacement text.
        replace_all: Replace every occurrence instead of requiring a unique match.

    Returns:
        Confirmation with a numbered snippet around the change.
    """
    ws = _ws(workspace)
    p = ws.path(file_path)
    if not p.is_file():
        return f"error: file not found: {p}"
    text = p.read_text(encoding="utf-8", errors="replace")
    try:
        updated = _apply_edit(text, old_string, new_string, replace_all)
    except ValueError as exc:
        return f"error: {exc}"
    p.write_text(updated, encoding="utf-8")
    first = updated.find(new_string) if new_string else 0
    line_no = updated.count("\n", 0, max(first, 0)) + 1
    lines = updated.splitlines()
    lo = max(0, line_no - 4)
    return f"Edited {p}. Snippet:\n{_numbered(lines[lo : line_no + 3 + new_string.count(chr(10))], lo + 1)}"


@_guard
def MultiEdit(workspace: str, file_path: str, edits: list[dict]) -> str:
    """Apply several Edit operations to ONE file in order, atomically (if any edit fails nothing is written). Each edit sees the result of the previous one.

    Args:
        workspace: Workspace id of this conversation.
        file_path: Absolute path of the file to modify.
        edits: List of objects {"old_string": str, "new_string": str, "replace_all": bool (optional)}.

    Returns:
        Confirmation.
    """
    ws = _ws(workspace)
    p = ws.path(file_path)
    if not p.is_file():
        return f"error: file not found: {p}"
    text = p.read_text(encoding="utf-8", errors="replace")
    for i, e in enumerate(edits or [], 1):
        try:
            text = _apply_edit(text, e["old_string"], e["new_string"], bool(e.get("replace_all", False)))
        except (ValueError, KeyError) as exc:
            return f"error: edit #{i} failed ({exc}). Nothing was written."
    p.write_text(text, encoding="utf-8")
    return f"Applied {len(edits)} edits to {p}"


@_guard
def Glob(workspace: str, pattern: str, path: str = "") -> str:
    """Find files by glob pattern (e.g. "**/*.py", "src/**/*.ts"), newest first. Skips .git, node_modules and caches.

    Args:
        workspace: Workspace id of this conversation.
        pattern: Glob pattern; ** matches any depth.
        path: Directory to search in. Empty means the current directory.

    Returns:
        Absolute file paths, one per line.
    """
    ws = _ws(workspace)
    base = ws.path(path) if path else ws.cwd()
    if not base.is_dir():
        return f"error: not a directory: {base}"
    hits = []
    for rel in _glob.glob(pattern, root_dir=base, recursive=True, include_hidden=True):
        full = base / rel
        if full.is_file() and not (set(Path(rel).parts) & SKIP_DIRS):
            hits.append(full)
    hits.sort(key=lambda f: f.stat().st_mtime, reverse=True)
    out = "\n".join(str(h) for h in hits[:250])
    return (out + (f"\n[{len(hits) - 250} more not shown]" if len(hits) > 250 else "")) or "(no matches)"


def _py_grep(base: Path, pattern, glob_pat, flags, mode, ctx_b, ctx_a, line_numbers, multiline) -> list[str]:
    rx = re.compile(pattern, flags | (re.DOTALL | re.MULTILINE if multiline else 0))
    files: list[Path] = [base] if base.is_file() else []
    if base.is_dir():
        for root, dirs, names in os.walk(base):
            dirs[:] = sorted(d for d in dirs if d not in SKIP_DIRS)
            files.extend(Path(root) / n for n in sorted(names))
    out: list[str] = []
    for f in files:
        if glob_pat and not fnmatch.fnmatch(f.name, glob_pat) and not fnmatch.fnmatch(str(f), glob_pat):
            continue
        try:
            raw = f.read_bytes()
        except OSError:
            continue
        if b"\0" in raw[:8192]:
            continue
        text = raw.decode(errors="replace")
        if multiline:
            n = len(rx.findall(text))
            if n:
                out.append(f"{f}:{n}" if mode == "count" else str(f))
            continue
        lines = text.splitlines()
        hit = [i for i, ln in enumerate(lines) if rx.search(ln)]
        if not hit:
            continue
        if mode == "files_with_matches":
            out.append(str(f))
        elif mode == "count":
            out.append(f"{f}:{len(hit)}")
        else:
            shown: set[int] = set()
            for i in hit:
                shown.update(range(max(0, i - ctx_b), min(len(lines), i + ctx_a + 1)))
            for i in sorted(shown):
                sep = ":" if i in hit else "-"
                out.append(f"{f}{sep}{i + 1}{sep}{lines[i][:500]}" if line_numbers else f"{f}{sep}{lines[i][:500]}")
    return out


@_guard
def Grep(
    workspace: str,
    pattern: str,
    path: str = "",
    glob: str = "",
    output_mode: str = "files_with_matches",
    case_insensitive: bool = False,
    line_numbers: bool = True,
    context_before: int = 0,
    context_after: int = 0,
    multiline: bool = False,
    head_limit: int = 250,
) -> str:
    """Search file contents with a regular expression (ripgrep). Respects .gitignore, skips binary files.

    Args:
        workspace: Workspace id of this conversation.
        pattern: Regular expression to search for.
        path: File or directory to search. Empty means the current directory.
        glob: Only search files matching this glob, e.g. "*.py".
        output_mode: "files_with_matches" (paths only), "content" (matching lines) or "count".
        case_insensitive: Ignore case.
        line_numbers: Include line numbers in content mode.
        context_before: Lines of context before each match (content mode).
        context_after: Lines of context after each match (content mode).
        multiline: Let the pattern span lines (. matches newline).
        head_limit: Maximum number of output lines.

    Returns:
        Matches, one per line.
    """
    ws = _ws(workspace)
    if output_mode not in ("files_with_matches", "content", "count"):
        return "error: output_mode must be files_with_matches, content or count"
    base = ws.path(path) if path else ws.cwd()
    if not base.exists():
        return f"error: path not found: {base}"
    if shutil.which("rg"):
        cmd = ["rg", "--no-heading", "--with-filename", "--color=never", "--max-columns=500"]
        cmd += {"files_with_matches": ["-l"], "count": ["-c"], "content": []}[output_mode]
        if output_mode == "content":
            if line_numbers:
                cmd.append("-n")
            if context_before:
                cmd += ["-B", str(int(context_before))]
            if context_after:
                cmd += ["-A", str(int(context_after))]
        if case_insensitive:
            cmd.append("-i")
        if multiline:
            cmd += ["-U", "--multiline-dotall"]
        if glob:
            cmd += ["--glob", glob]
        cmd += ["-e", pattern, "--", str(base)]
        res = subprocess.run(cmd, capture_output=True, text=True, errors="replace", timeout=60, cwd=ws.cwd())
        if res.returncode == 2:
            return f"error: {res.stderr.strip()}"
        lines = res.stdout.splitlines()
    else:
        lines = _py_grep(base, pattern, glob, re.I if case_insensitive else 0, output_mode,
                         int(context_before), int(context_after), line_numbers, multiline)
    if not lines:
        return "(no matches)"
    cap = max(1, int(head_limit))
    return _clip("\n".join(lines[:cap]) + (f"\n[{len(lines) - cap} more lines not shown]" if len(lines) > cap else ""))


@_guard
def LS(workspace: str, path: str = "", depth: int = 2, ignore: str = "") -> str:
    """List files and folders as a tree. Prefer Glob/Grep when you know what you are looking for.

    Args:
        workspace: Workspace id of this conversation.
        path: Directory to list. Empty means the current directory.
        depth: How many levels to descend (1-5).
        ignore: Optional glob of names to hide, e.g. "*.pyc".

    Returns:
        An indented tree; folders end with "/".
    """
    ws = _ws(workspace)
    base = ws.path(path) if path else ws.cwd()
    if not base.is_dir():
        return f"error: not a directory: {base}"
    depth = max(1, min(int(depth), 5))
    out, budget = [f"{base}/"], [400]

    def walk(d: Path, level: int) -> None:
        try:
            entries = sorted(d.iterdir(), key=lambda e: (not e.is_dir(), e.name.lower()))
        except OSError:
            return
        for e in entries:
            if budget[0] <= 0:
                return
            if e.name in SKIP_DIRS or (ignore and fnmatch.fnmatch(e.name, ignore)):
                continue
            budget[0] -= 1
            out.append(f"{'  ' * level}- {e.name}{'/' if e.is_dir() else ''}")
            if e.is_dir() and not e.is_symlink() and level < depth:
                walk(e, level + 1)

    walk(base, 1)
    if budget[0] <= 0:
        out.append("[truncated at 400 entries]")
    return "\n".join(out)


# --------------------------------------------------------------------------- web


@_guard
def WebFetch(url: str, max_chars: int = 20000, offset: int = 0) -> str:
    """Download a web page and return it as markdown text (HTML is converted; other text is returned as is). Use offset to read further into long pages.

    Args:
        url: The http(s) URL to fetch.
        max_chars: Maximum characters to return.
        offset: Character offset to start from.

    Returns:
        The page content.
    """
    import httpx

    if not url.startswith(("http://", "https://")):
        return "error: url must start with http:// or https://"
    r = httpx.get(url, follow_redirects=True, timeout=25, headers={"User-Agent": "Mozilla/5.0 (claude-code-space)"})
    ctype = r.headers.get("content-type", "")
    body = r.text
    if "html" in ctype:
        import html2text

        conv = html2text.HTML2Text()
        conv.ignore_images, conv.body_width = True, 0
        body = conv.handle(body)
    body = body[max(0, int(offset)) :]
    cap = max(1, int(max_chars))
    more = f"\n[truncated: call again with offset={int(offset) + cap}]" if len(body) > cap else ""
    return f"[{r.status_code}] {r.url}\n\n{body[:cap]}{more}"


@_guard
def WebSearch(query: str, max_results: int = 8) -> str:
    """Search the web (DuckDuckGo) and return titles, URLs and snippets. Follow up with WebFetch to read a result.

    Args:
        query: The search query.
        max_results: Number of results (1-20).

    Returns:
        Numbered results.
    """
    from ddgs import DDGS

    rows = DDGS().text(query, max_results=max(1, min(int(max_results), 20)))
    return "\n\n".join(
        f"{i}. {r.get('title')}\n   {r.get('href')}\n   {r.get('body', '')}" for i, r in enumerate(rows, 1)
    ) or "(no results)"


# --------------------------------------------------------------------------- todos


@_guard
def TodoWrite(workspace: str, todos: list[dict]) -> str:
    """Create and update the task list for the current work. Send the COMPLETE list every time. Use it for tasks with 3+ steps; keep exactly one item in_progress.

    Args:
        workspace: Workspace id of this conversation.
        todos: List of {"content": str, "status": "pending" | "in_progress" | "completed"}.

    Returns:
        The rendered list.
    """
    ws = _ws(workspace)
    marks = {"pending": "[ ]", "in_progress": "[~]", "completed": "[x]"}
    for t in todos or []:
        if t.get("status") not in marks or not t.get("content"):
            return "error: each todo needs 'content' and a 'status' of pending, in_progress or completed"
    ws.todos_file.write_text(json.dumps(todos))
    return "Todos updated:\n" + "\n".join(f"{marks[t['status']]} {t['content']}" for t in todos)


ALL_TOOLS = [
    workspace_new, workspace_list, workspace_delete,
    Bash, BashOutput, KillShell,
    Read, Write, Edit, MultiEdit, Glob, Grep, LS,
    WebFetch, WebSearch, TodoWrite,
]