# tools/shell.py - Shell execution, confirmation IPC, environment setup, and app prefs loader.
from __future__ import annotations

import json
import os
import selectors
import shlex
import subprocess
import time
from datetime import datetime, timedelta, timezone
from threading import Event, Lock

from langchain_core.tools import tool

from persistence import MUGI_LOGS_DIR
from tools._logging import _debug_log, _emit_input_status
from tools.job_stream_schema import (
    close_job_stream,
    error_job_stream,
    make_capabilities,
    metric_def,
    open_job_stream,
    update_job_stream,
)
from tools.live_view_visuals import emit_live_view_preview
from tools.plugin_stream_controls import clear_stream_state, get_action_count
from tools._logging import register_cancellable_subprocess, should_stop_tool, unregister_cancellable_subprocess

_SHELL_ALLOW_CACHE_PATH = os.path.join(os.path.expanduser("~"), ".mugi", "shell_allow_cache.json")

_MUGI_DIR_IPC = os.path.abspath(os.path.expanduser(os.environ.get("MUGI_DATA_DIR", "~/.mugi")))
if os.path.isdir(_MUGI_DIR_IPC):
    try:
        _ipc_stat = os.stat(_MUGI_DIR_IPC)
        if _ipc_stat.st_uid != os.getuid():
            _MUGI_DIR_IPC = os.path.expanduser("~/.mugi")
    except OSError:
        _MUGI_DIR_IPC = os.path.expanduser("~/.mugi")
_SHELL_CONFIRM_REQUEST_PATH = os.path.join(_MUGI_DIR_IPC, ".shell_confirm_request.json")
_SHELL_CONFIRM_RESPONSE_PATH = os.path.join(_MUGI_DIR_IPC, ".shell_confirm_response.json")
_YTDL_COOKIES_REQUEST_PATH = os.path.join(_MUGI_DIR_IPC, ".ytdl_cookies_request.json")
_YTDL_COOKIES_RESPONSE_PATH = os.path.join(_MUGI_DIR_IPC, ".ytdl_cookies_response.json")
_shell_allow_cache_lock = Lock()
_pending_shell_confirm_lock = Lock()
_pending_shell_confirm = None

_ENV_WHITELIST = frozenset({"HOME", "USER", "PATH", "LANG", "LC_ALL", "TERM", "TMPDIR", "SHELL", "LOGNAME", "XPC_FLAGS", "XPC_SERVICE_NAME"})


def _load_app_prefs_tools():
    """Load app prefs from ~/.mugi/app_prefs.json via shared cache."""
    try:
        from utils.app_prefs_store import load_app_prefs_cached
        return load_app_prefs_cached()
    except Exception:
        return {}


def get_pending_shell_confirm():
    """Called from app main thread: return (command, response_holder, event) and clear pending, or None."""
    with _pending_shell_confirm_lock:
        p = _pending_shell_confirm
        globals()["_pending_shell_confirm"] = None
        return p


def complete_shell_confirm(response_holder, event, result):
    """Called from app after user responds; result is one of 'deny'|'allow_once'|'allow_1'|'allow_5'|'allow_10'|'allow_20'."""
    if response_holder is not None:
        response_holder[0] = result
    if event is not None:
        event.set()


def get_pending_shell_confirm_from_file():
    """Check for a file-based confirmation request (agent in another process). Returns (command, request_id) or None. Deletes request file."""
    if not os.path.isfile(_SHELL_CONFIRM_REQUEST_PATH):
        return None
    try:
        with open(_SHELL_CONFIRM_REQUEST_PATH, "r", encoding="utf-8") as f:
            data = json.load(f)
        command = (data.get("command") or "").strip()
        request_id = (data.get("request_id") or "").strip()
        if not request_id:
            return None
        try:
            os.unlink(_SHELL_CONFIRM_REQUEST_PATH)
        except OSError:
            pass
        return (command, request_id)
    except Exception:
        try:
            os.unlink(_SHELL_CONFIRM_REQUEST_PATH)
        except OSError:
            pass
        return None


def complete_shell_confirm_file(request_id: str, choice: str):
    """Write the user's choice to the response file so the agent process can see it. Uses atomic write with 0o600."""
    try:
        from persistence import secure_atomic_write
        secure_atomic_write(_SHELL_CONFIRM_RESPONSE_PATH, json.dumps({"request_id": request_id, "choice": choice}))
    except Exception:
        pass


def get_pending_ytdl_cookies_from_file():
    """Legacy shim for ytdl cookies pending request, backed by generic user-decision IPC."""
    try:
        from user_decision_ipc import get_pending_user_decision_from_file
        pending = get_pending_user_decision_from_file()
        if not pending or (pending.get("kind") or "") != "browser_cookie_source":
            return None
        ctx = pending.get("context") if isinstance(pending.get("context"), dict) else {}
        return {
            "request_id": pending.get("request_id") or "",
            "url": ctx.get("url") or "",
            "out_template": ctx.get("out_template") or "",
            "format_preference": ctx.get("format_preference") or "",
            "audio_only": bool(ctx.get("audio_only")),
        }
    except Exception:
        return None


def complete_ytdl_cookies_file(request_id: str, choice: str):
    """Legacy shim for writing browser choice through generic user-decision IPC."""
    try:
        from user_decision_ipc import complete_user_decision_file
        complete_user_decision_file(request_id, choice, kind="browser_cookie_source")
    except Exception:
        pass


def wait_for_ytdl_cookies_choice(url: str, out_template: str, format_preference: str, audio_only: bool) -> str | None:
    """Request user choice to retry yt-dlp with browser cookies."""
    try:
        from user_decision_ipc import request_user_decision
        return request_user_decision(
            kind="browser_cookie_source",
            context={
                "title": "Retry with browser cookies?",
                "message": (
                    "Download failed due to cookie/auth restrictions. Pick the browser where you are signed in. "
                    "Prefer Chrome or Brave if you use them—Safari’s cookie file is often blocked unless Mugi has "
                    "Full Disk Access (System Settings → Privacy & Security → Full Disk Access)."
                ),
                "url": url,
                "out_template": out_template,
                "format_preference": format_preference,
                "audio_only": bool(audio_only),
            },
            source="core",
            timeout_s=300.0,
        )
    except Exception:
        return None


def wait_for_shell_confirm(command: str, use_allow_cache: bool = True) -> str | None:
    """Request user confirmation to run the given command. Returns choice string or None on timeout."""
    now_utc = datetime.now(timezone.utc)
    end_utc = None
    if use_allow_cache:
        with _shell_allow_cache_lock:
            cache = {}
            if os.path.isfile(_SHELL_ALLOW_CACHE_PATH):
                try:
                    with open(_SHELL_ALLOW_CACHE_PATH, "r", encoding="utf-8") as f:
                        cache = json.load(f)
                except Exception:
                    pass
            end_time_str = (cache or {}).get("end_time")
            if end_time_str:
                try:
                    s = end_time_str.replace("Z", "+00:00")
                    end_utc = datetime.fromisoformat(s)
                    if end_utc.tzinfo is None:
                        end_utc = end_utc.replace(tzinfo=timezone.utc)
                    if now_utc >= end_utc:
                        end_utc = None
                except Exception:
                    end_utc = None
            if end_utc is not None:
                count = cache.get("count", 0)
                if count >= 10:
                    end_utc = None
                else:
                    cache["count"] = count + 1
                    try:
                        from persistence import secure_atomic_write
                        secure_atomic_write(_SHELL_ALLOW_CACHE_PATH, json.dumps(cache, ensure_ascii=False))
                    except Exception:
                        pass
        if end_utc is not None:
            return "allow_once"
    response_holder = [None]
    event = Event()
    from persistence import generate_secure_token, secure_atomic_write
    request_id = generate_secure_token(32)
    with _pending_shell_confirm_lock:
        globals()["_pending_shell_confirm"] = (command, response_holder, event)
    try:
        secure_atomic_write(_SHELL_CONFIRM_REQUEST_PATH, json.dumps({"command": command, "request_id": request_id}))
    except Exception:
        pass
    deadline = time.time() + 300
    choice = None
    while time.time() < deadline:
        if event.is_set() and response_holder[0] is not None:
            choice = response_holder[0]
            break
        if os.path.isfile(_SHELL_CONFIRM_RESPONSE_PATH):
            try:
                with open(_SHELL_CONFIRM_RESPONSE_PATH, "r", encoding="utf-8") as f:
                    resp = json.load(f)
                if (resp.get("request_id") or "").strip() == request_id:
                    choice = (resp.get("choice") or "deny").strip() or "deny"
                    try:
                        os.unlink(_SHELL_CONFIRM_RESPONSE_PATH)
                    except OSError:
                        pass
                    break
            except Exception:
                pass
        event.wait(timeout=0.2)
    with _pending_shell_confirm_lock:
        globals()["_pending_shell_confirm"] = None
    try:
        if os.path.isfile(_SHELL_CONFIRM_REQUEST_PATH):
            with open(_SHELL_CONFIRM_REQUEST_PATH, "r", encoding="utf-8") as f:
                if json.load(f).get("request_id") == request_id:
                    os.unlink(_SHELL_CONFIRM_REQUEST_PATH)
    except Exception:
        pass
    if choice is None:
        return None
    if choice in ("allow_1", "allow_5", "allow_10", "allow_20"):
        minutes = {"allow_1": 1, "allow_5": 5, "allow_10": 10, "allow_20": 20}[choice]
        end_utc = now_utc + timedelta(minutes=minutes)
        with _shell_allow_cache_lock:
            try:
                from persistence import secure_atomic_write
                secure_atomic_write(_SHELL_ALLOW_CACHE_PATH, json.dumps({"end_time": end_utc.isoformat(), "count": 0}, ensure_ascii=False))
            except Exception:
                pass
    return choice


def _mugi_rich_path() -> str:
    """Build a PATH that includes ~/.mugi/bin, ~/.mugi/homebrew/bin, and any CLI venv bins first."""
    from tools.mugi_brew_bootstrap import get_mugi_homebrew_bin_dir
    from persistence import CLI_VENVS_DIR
    mugi_bin = os.path.join(os.path.expanduser("~"), ".mugi", "bin")
    parts = [mugi_bin, get_mugi_homebrew_bin_dir()]
    if os.path.isdir(CLI_VENVS_DIR):
        try:
            for name in sorted(os.listdir(CLI_VENVS_DIR)):
                venv_bin = os.path.join(CLI_VENVS_DIR, name, "bin")
                if os.path.isdir(venv_bin):
                    parts.append(venv_bin)
        except OSError:
            pass
    parts.extend(["/opt/homebrew/bin", "/usr/local/bin"])
    existing = os.environ.get("PATH", "")
    if existing:
        parts.append(existing)
    return os.pathsep.join(parts)


def _safe_subprocess_env(**extra):
    """Return a minimal environment for subprocesses, keeping only safe keys."""
    env = {k: v for k, v in os.environ.items() if k in _ENV_WHITELIST}
    env["PATH"] = _mugi_rich_path()
    env.update(extra)
    return env


def _contains_shell_substitution(command: str) -> bool:
    """Detect backticks or $() outside single-quoted literals.

    This intentionally flags substitutions in unquoted and double-quoted contexts.
    """
    in_single = False
    in_double = False
    escaped = False
    i = 0
    while i < len(command):
        ch = command[i]
        if escaped:
            escaped = False
            i += 1
            continue
        if ch == "\\":
            escaped = True
            i += 1
            continue
        if ch == "'" and not in_double:
            in_single = not in_single
            i += 1
            continue
        if ch == '"' and not in_single:
            in_double = not in_double
            i += 1
            continue
        if not in_single:
            if ch == "`":
                return True
            if ch == "$" and i + 1 < len(command) and command[i + 1] == "(":
                return True
        i += 1
    return False


@tool
def shell_execute(command: str) -> str:
    """Execute a shell command on the laptop.
    DANGEROUS: Requires user confirmation each time unless temporarily allowed. Only available if enabled in Preferences.
    Input: the shell command to run (string). Example: "date", "ls -la ~/Downloads", "pwd"
    """
    _debug_log(f"[DEBUG shell_execute] Called with command: {command!r}")
    try:
        command = (command or "").strip()
        if not command:
            return "No command provided."
        prefs = _load_app_prefs_tools()
        if not prefs.get("allow_shell_execute", False):
            return "Shell commands are disabled in app preferences."
        _BLOCKED_BINARIES = frozenset({
            "mkfs", "shred", "wipefs", "fdisk", "diskutil",
            "dd", "sudo", "launchctl", "csrutil", "kill", "mkfifo",
            "chmod", "chown", "srm",
            "nohup", "at", "crontab", "nc", "ncat",
        })
        _DANGEROUS_RM_TARGETS = frozenset({"/", "/etc", "/usr", "/var", "/System", "/Applications", "/Library"})
        try:
            argv = shlex.split(command)
        except ValueError:
            return "Blocked: could not parse command safely."
        if not argv:
            return "No command provided."
        binary = os.path.basename(argv[0])
        if binary in _BLOCKED_BINARIES:
            return f"Blocked: '{binary}' is a potentially destructive or privileged command."
        if binary == "rm" and any(t.rstrip("/") in _DANGEROUS_RM_TARGETS for t in argv[1:] if not t.startswith("-")):
            return "Blocked: command targets a protected system path."
        _cmd_lower = command.lower()
        if "curl" in _cmd_lower and ("| sh" in _cmd_lower or "| bash" in _cmd_lower or "|sh" in _cmd_lower or "|bash" in _cmd_lower):
            return "Blocked: piping curl to shell is not allowed."
        if "dd" in _cmd_lower and "if=/dev/zero" in _cmd_lower:
            return "Blocked: dd with /dev/zero is not allowed."
        if _contains_shell_substitution(command):
            return "Blocked: command contains shell substitution patterns (backtick or $()) that are not permitted."
        if "rm" in _cmd_lower and ("homebrew" in _cmd_lower or ".mugi/homebrew" in _cmd_lower):
            return "Blocked: do not remove ~/.mugi/homebrew. Use the app or ensure_mugi_homebrew to reinstall if needed."
        _plugins_abs = os.path.expanduser("~/.mugi/plugins")
        if ".mugi/plugins" in command or _plugins_abs in command:
            return "Blocked: do not modify the plugins directory via shell. Use create_plugin (evolution mode) or manage plugins in Settings."
        choice = wait_for_shell_confirm(command, use_allow_cache=True)
        if choice is None:
            return "Command confirmation timed out (no response from user)."
        if choice == "deny":
            return "Command denied by user."
        _debug_log(f"shell: running — {command[:100]}{'...' if len(command) > 100 else ''}", "INFO")
        stream_id = f"shell_job:{int(time.time() * 1000)}"
        open_job_stream(
            plugin_id="shell_job",
            plugin_name="Shell Job",
            title="Shell Command",
            job_type="shell_execute",
            stream_id=stream_id,
            controls=[
                {
                    "id": "stop_task",
                    "type": "button",
                    "label": "Stop Command",
                    "action_id": "stop_task",
                    "style": "destructive",
                }
            ],
            capabilities=make_capabilities(
                cancel=True,
                progress=False,
                metrics=True,
                phases=True,
                stream_logs=True,
            ),
            metric_schema=[
                metric_def("stdout_lines", "Stdout Lines", unit="lines", format="number", priority=2),
                metric_def("stderr_lines", "Stderr Lines", unit="lines", format="number", priority=2),
                metric_def("elapsed_seconds", "Elapsed", unit="s", format="duration_s", priority=3),
            ],
            metadata={"command": command[:500]},
        )
        safe_cwd = os.path.expanduser("~/Desktop")
        if not os.path.isdir(safe_cwd):
            safe_cwd = os.path.expanduser("~")
        started = time.time()
        try:
            update_job_stream(
                plugin_id="shell_job",
                plugin_name="Shell Job",
                stream_id=stream_id,
                phase="run",
                status="running",
                message="Running shell command...",
            )
            proc = subprocess.Popen(
                command,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                cwd=safe_cwd,
                env=_safe_subprocess_env(),
                bufsize=1,
            )
            register_cancellable_subprocess(proc)
            try:
                sel = selectors.DefaultSelector()
                if proc.stdout:
                    sel.register(proc.stdout, selectors.EVENT_READ, data="stdout")
                if proc.stderr:
                    sel.register(proc.stderr, selectors.EVENT_READ, data="stderr")
                stdout_lines: list[str] = []
                stderr_lines: list[str] = []
                last_emit = 0.0
                while proc.poll() is None:
                    if should_stop_tool() or get_action_count(stream_id, "stop_task") > 0:
                        try:
                            proc.terminate()
                        except Exception:
                            pass
                        close_job_stream(
                            plugin_id="shell_job",
                            plugin_name="Shell Job",
                            stream_id=stream_id,
                            status="stopped",
                            message="Shell command stopped by user.",
                        )
                        return "Shell command stopped by user."
                    events = sel.select(timeout=0.2)
                    for key, _ in events:
                        line = key.fileobj.readline()
                        if not line:
                            continue
                        txt = line.rstrip("\n")
                        if key.data == "stdout":
                            stdout_lines.append(txt)
                        else:
                            stderr_lines.append(txt)
                        now = time.time()
                        if now - last_emit >= 0.8:
                            update_job_stream(
                                plugin_id="shell_job",
                                plugin_name="Shell Job",
                                stream_id=stream_id,
                                phase="run",
                                status="running",
                                message=f"{key.data}: {txt[:120]}",
                                metrics={
                                    "stdout_lines": float(len(stdout_lines)),
                                    "stderr_lines": float(len(stderr_lines)),
                                    "elapsed_seconds": float(round(now - started, 2)),
                                },
                                extra={"log_line": txt[-400:], "log_stream": key.data},
                            )
                            last_emit = now
                if proc.stdout:
                    rest = proc.stdout.read() or ""
                    if rest.strip():
                        stdout_lines.extend([ln for ln in rest.splitlines() if ln.strip()])
                if proc.stderr:
                    rest = proc.stderr.read() or ""
                    if rest.strip():
                        stderr_lines.extend([ln for ln in rest.splitlines() if ln.strip()])
                class _Result:
                    pass
                result = _Result()
                result.returncode = proc.returncode if proc.returncode is not None else 1
                result.stdout = "\n".join(stdout_lines)
                result.stderr = "\n".join(stderr_lines)
                _debug_log(f"shell: finished exit {result.returncode}", "INFO")
            finally:
                unregister_cancellable_subprocess(proc)
            output = f"Command: {command}\n"
            if result.stdout and result.stdout.strip():
                output += f"stdout:\n{result.stdout.strip()}\n"
            if result.stderr and result.stderr.strip():
                output += f"stderr:\n{result.stderr.strip()}\n"
            output += f"Return code: {result.returncode}"
            if result.returncode != 0:
                combined = (result.stdout + "\n" + result.stderr).lower()
                clt_indicators = ("xcode-select", "command line tools", "no developer tools", "xcrun: error")
                if any(ind in combined for ind in clt_indicators):
                    output += (
                        "\n\nHint: This failure may be caused by missing Apple Command Line Tools."
                        " You can install them by running: xcode-select --install"
                    )
            if int(result.returncode) == 0:
                close_job_stream(
                    plugin_id="shell_job",
                    plugin_name="Shell Job",
                    stream_id=stream_id,
                    status="completed",
                    message="Shell command completed.",
                    metrics={
                        "stdout_lines": float(len(stdout_lines)),
                        "stderr_lines": float(len(stderr_lines)),
                        "elapsed_seconds": float(round(time.time() - started, 2)),
                    },
                    extra={"return_code": int(result.returncode)},
                )
            else:
                error_job_stream(
                    plugin_id="shell_job",
                    plugin_name="Shell Job",
                    stream_id=stream_id,
                    phase="run",
                    message=f"Shell command exited with code {result.returncode}.",
                    extra={"return_code": int(result.returncode)},
                )
            return output
        except subprocess.TimeoutExpired:
            _debug_log("shell: timed out after 60s", "INFO")
            error_job_stream(
                plugin_id="shell_job",
                plugin_name="Shell Job",
                stream_id=stream_id,
                phase="run",
                message=f"Command timed out after 60 seconds: {command}",
            )
            return f"Command timed out after 60 seconds: {command}"
        except Exception as e:
            _debug_log(f"shell: failed — {e}", "INFO")
            error_job_stream(
                plugin_id="shell_job",
                plugin_name="Shell Job",
                stream_id=stream_id,
                phase="run",
                message=f"Shell execution failed: {str(e)}",
            )
            return f"Shell execution failed: {str(e)}"
        finally:
            clear_stream_state(stream_id)
    except Exception as e:
        return f"Shell execution failed: {str(e)}"
