import Foundation

/// Names the worktree files a raw shell command writes (`echo … > f`, `cat > f`, `sed -i`, `tee`,
/// `dd of=`, …), so those writes can go through the **same lock/arbitration path a structured edit
/// takes** instead of being refused.
///
/// Why this changed. Structured edits (Edit/Write/MultiEdit, Codex `apply_patch`, the ACP edit) are
/// synthesized into `fileChange` events and landed as per-edit trunk commits, with conflict
/// arbitration on the touched paths. A file mutated by raw shell emits no `fileChange`, so it used
/// to skip conflict locking and never get its own trunk commit — and the approval layer answered
/// that by **denying** the command outright and steering the agent to the structured tool. Denial
/// was a blunt answer: `echo … >> CHANGELOG.md` and `sed -i` are ordinary, legitimate shell, and
/// refusing them costs the agent a round trip on work it had every right to do.
///
/// Nucleic now has the stack to allow them *safely* instead (docs/NASH.md, docs/LOCKING.md §4):
///
/// 1. **Before the command runs**, the approval gate asks this type which worktree files the
///    command writes and arbitrates exactly those paths — the same `arbitrate` seam Edit/Write use.
///    A contended file queues the Bash call just as it would queue an edit; a granted one lets the
///    command through holding real locks.
/// 2. **After the command runs**, nash — the agent's own shell — reports the redirection it
///    actually performed (NASH §5.2) with a resolved path. The host claims the lock for anything
///    this static parse missed and lands the bytes into the trunk, so a shell write ends up tracked
///    exactly like a structured edit.
///
/// So this type's question changed from "is this command forbidden?" to "which worktree files does
/// this command write?". One denial survives: a mutation whose target **cannot be named
/// statically** — `patch` applying a diff, a glob, an unexpanded `$VAR`. We cannot lock what we
/// cannot name, so those still steer to the structured edit tool.
///
/// The parse stays deliberately conservative about what counts as *authoring* a file, so ordinary
/// build/tooling output is never dragged into the lock system: a redirection is a write only when
/// its producer is a plain content emitter (`make > build.log` and a compiler's `-o` are not), and
/// clearly-ephemeral destinations (`/tmp`, `/dev/null`, `$TMPDIR`, `.git/`) are never claimed.
public enum WorktreeMutationGuard {

    /// What a shell command does to worktree files.
    public struct ShellWriteAnalysis: Sendable, Equatable {
        /// Files the command writes, exactly as the command spells them — relative to the command's
        /// cwd, or absolute. Nameable, therefore lockable: these are what the approval gate
        /// arbitrates before letting the command run.
        public let writes: [String]
        /// Mutations whose target cannot be named statically, as the offending fragment
        /// (`patch`, `sed -i` over a glob, `echo > $OUT`). These are still denied — a lock needs a
        /// path.
        public let unnameable: [String]

        public init(writes: [String], unnameable: [String]) {
            self.writes = writes
            self.unnameable = unnameable
        }

        public var isEmpty: Bool { writes.isEmpty && unnameable.isEmpty }
    }

    // MARK: - Public surface

    /// Every worktree file `command` writes, plus the mutations whose target can't be named.
    ///
    /// A relative target is resolved against the directory the segment runs in, which a `cd`
    /// earlier in the same command line moves: in `cd Sources && echo x > main.swift` the file is
    /// `Sources/main.swift`, and locking a bare `main.swift` would guard the wrong path. The walk
    /// therefore carries a running directory and gives up — reporting the target as unnameable
    /// rather than guessing — the moment it can't follow (`cd -`, `popd`, a subshell, a `cd $VAR`).
    /// Giving up is never a regression: an unnameable target is denied, which is what *every* one
    /// of these commands got before locks reached them at all.
    public static func analyze(command: String) -> ShellWriteAnalysis {
        let (flattened, _) = ShellLexer.stripHeredocs(command)
        var writes: [String] = []
        var unnameable: [String] = []
        var cwd: String? = ""  // relative to the shell's own cwd; nil once we've lost track
        for segment in ShellLexer.splitSegments(flattened) {
            let found = mutations(in: segment)
            for target in found.writes {
                if target.hasPrefix("/") {
                    writes.append(target)
                } else if let cwd {
                    writes.append(cwd.isEmpty ? target : joinPath(cwd, target))
                } else {
                    unnameable.append(target)
                }
            }
            unnameable.append(contentsOf: found.unnameable)
            cwd = advanced(cwd, over: segment)
        }
        return ShellWriteAnalysis(
            writes: orderedUnique(writes), unnameable: orderedUnique(unnameable))
    }

    /// The paths the approval gate should acquire locks on before letting `command` run: every
    /// nameable write target resolved to an absolute path against `workdir` (the directory the
    /// agent's shell runs in) and filtered to those that actually live inside it.
    ///
    /// Filtering to `workdir` is what keeps the lock system honest: `sudo tee /etc/hosts` writes a
    /// file, but not a *worktree* file, so it is none of the lock system's business (the risk
    /// classifier and the approval prompt still see it). A nil `workdir` — no run context — leaves
    /// the targets as written and lets the arbiter normalize them.
    public static func lockableWrites(command: String, workdir: String?) -> [String] {
        let writes = analyze(command: command).writes
        guard let workdir, !workdir.isEmpty else { return writes }
        let root = normalized(workdir)
        return orderedUnique(writes.compactMap { target -> String? in
            let clean = target.hasPrefix("/") ? normalized(target) : joinPath(root, target)
            guard clean.hasPrefix(root + "/") else { return nil }
            return clean
        })
    }

    /// A human-readable denial when `command` mutates a worktree file in a way that **can't** be
    /// locked, because its target isn't statically nameable; nil otherwise (including for writes
    /// that are nameable — those are arbitrated and allowed, not denied).
    public static func denialReason(command: String) -> String? {
        let unnameable = analyze(command: command).unnameable
        guard let culprit = unnameable.first else { return nil }
        return denialMessage(culprit: culprit)
    }

    // MARK: - Segment analysis

    /// The worktree writes one command segment performs, split into nameable targets and
    /// mutations whose target can't be resolved without running the shell.
    private static func mutations(in segment: String) -> (writes: [String], unnameable: [String]) {
        let (commandText, redirectTargets) = splitRedirections(segment)
        let rawTokens = ShellLexer.tokenize(commandText)
        guard !rawTokens.isEmpty else { return ([], []) }
        let tokens = ShellLexer.stripLeadingPrefixes(
            ShellLexer.stripBodyIntroducers(rawTokens))
        guard let first = tokens.first else { return ([], []) }
        let cmd = commandBasename(first)
        let args = Array(tokens.dropFirst()).filter { !$0.isEmpty }

        var writes: [String] = []
        var unnameable: [String] = []

        /// Record `paths` as this command's targets, sorting each into nameable/not and dropping
        /// throwaway destinations. `culprit` names the mutation when a path can't be resolved.
        func claim(_ paths: [String], culprit: String, requireTarget: Bool) {
            let interesting = paths.filter { !isEphemeral($0) }
            if interesting.isEmpty {
                // Either everything went somewhere throwaway (fine) or the command named no file at
                // all — for a mutator whose target is implicit (`patch` reading a diff) that second
                // case is precisely the one we can't lock.
                if requireTarget && paths.isEmpty { unnameable.append(culprit) }
                return
            }
            for path in interesting {
                if isNameable(path) { writes.append(path) } else { unnameable.append("\(culprit) \(path)") }
            }
        }

        // 1. In-place editors / file-writing utilities — the file operand IS the mutation target,
        //    regardless of any redirection.
        switch cmd {
        case "sed":
            if hasInPlaceFlag(args, letter: "i") {
                claim(scriptToolOperands(args, scriptFlags: ["-e", "--expression", "-f", "--file"]),
                      culprit: "sed -i", requireTarget: true)
            }
        case "perl":
            if hasInPlaceFlag(args, letter: "i") {
                claim(perlOperands(args), culprit: "perl -i", requireTarget: true)
            }
        case "awk", "gawk", "mawk":
            // GNU awk in-place is `-i inplace` / `-iinplace`.
            if args.contains(where: { $0 == "inplace" || $0.hasSuffix("=inplace") || $0 == "-iinplace" }) {
                claim(scriptToolOperands(args, scriptFlags: ["-f", "--file", "--source"],
                                         valueFlags: ["-v", "-i", "--assign"]),
                      culprit: "awk -i inplace", requireTarget: true)
            }
        case "ed":
            // `ed` edits its file operand in place from a script on stdin.
            claim(ShellLexer.positionals(args, valueFlags: []), culprit: "ed", requireTarget: true)
        case "patch":
            // Without an explicit operand the target lives inside the diff, which we can't read.
            var targets = ShellLexer.positionals(args, valueFlags: patchValueFlags)
            if let out = ShellLexer.flagValue(args, short: "-o", long: "--output") { targets.append(out) }
            claim(targets, culprit: "patch", requireTarget: true)
        case "tee":
            // tee's non-flag operands are the files it writes; `-a` (append) takes no value.
            claim(ShellLexer.positionals(args, valueFlags: []), culprit: "tee", requireTarget: false)
        case "dd":
            if let of = args.first(where: { $0.hasPrefix("of=") }) {
                claim([String(of.dropFirst(3))], culprit: "dd of=", requireTarget: false)
            }
        case "truncate":
            claim(ShellLexer.positionals(args, valueFlags: ["-s", "--size", "-r", "--reference"]),
                  culprit: "truncate", requireTarget: false)
        default:
            break
        }

        // 2. Output redirection (`>`, `>>`) whose producer is a plain content emitter — the
        //    "hand-write a file's contents" shape (`echo … > f`, `cat > f`, `sed … > f`). A build
        //    tool's redirected output (`make > build.log`) is an artifact, not an edit, and is
        //    deliberately left alone.
        if contentEmitters.contains(cmd) {
            claim(redirectTargets, culprit: "\(cmd) >", requireTarget: false)
        }
        return (writes, unnameable)
    }

    /// Commands that emit text to stdout, so redirecting them into a file authors that file's
    /// contents by hand — as opposed to build/compiler tools whose file output is a legitimate
    /// artifact and must not be pulled into the lock system.
    private static let contentEmitters: Set<String> = [
        "echo", "printf", "cat", "tee", "yes", "jq", "sed", "awk", "gawk",
    ]

    /// `patch` flags that consume the following argument, so its file operand is read correctly.
    private static let patchValueFlags: Set<String> = [
        "-p", "--strip", "-i", "--input", "-o", "--output", "-d", "--directory",
        "-r", "--reject-file", "-b", "--suffix", "-D", "--ifdef", "-F", "--fuzz",
        "-z", "--backup-suffix", "-V", "--version-control",
    ]

    // MARK: - Following `cd` across a command line

    /// The directory the *next* segment runs in, given the directory this one ran in. Returns nil
    /// ("lost track") for anything whose destination we can't compute statically — a bare `cd`
    /// (home), `cd -`, `popd`, a `cd "$DIR"`, or a subshell, whose directory change is scoped in a
    /// way this flat segment walk doesn't model.
    private static func advanced(_ cwd: String?, over segment: String) -> String? {
        guard let cwd else { return nil }
        let (commandText, _) = splitRedirections(segment)
        var tokens = ShellLexer.tokenize(commandText)
        // A subshell (`(cd x && …)`) scopes its own cd; rather than model that, stop following.
        if tokens.first?.hasPrefix("(") == true || tokens.contains(")") { return nil }
        tokens = ShellLexer.stripLeadingPrefixes(ShellLexer.stripBodyIntroducers(tokens))
        guard let head = tokens.first else { return cwd }
        switch commandBasename(head) {
        case "cd", "pushd":
            guard let dest = ShellLexer.positionals(Array(tokens.dropFirst()), valueFlags: [])
                .first(where: { !$0.isEmpty }), isNameable(dest), dest != "-"
            else { return nil }
            return dest.hasPrefix("/") ? normalized(dest) : joinPath(cwd, dest)
        case "popd":
            return nil
        default:
            return cwd
        }
    }

    /// Lexically join two path fragments and collapse `.`/`..`, preserving whether the result is
    /// absolute. Purely textual — these paths name files that may not exist yet, so nothing here
    /// may touch the filesystem.
    private static func joinPath(_ base: String, _ relative: String) -> String {
        normalized(base.isEmpty ? relative : "\(base)/\(relative)")
    }

    private static func normalized(_ path: String) -> String {
        let absolute = path.hasPrefix("/")
        var stack: [String] = []
        for component in path.split(separator: "/") {
            switch component {
            case ".": continue
            case "..":
                if let last = stack.last, last != ".." { stack.removeLast() }
                else if !absolute { stack.append("..") }
            default: stack.append(String(component))
            }
        }
        return (absolute ? "/" : "") + stack.joined(separator: "/")
    }

    // MARK: - Operand extraction

    /// File operands of a "script tool" — `sed`/`awk`, whose first positional is the *program*
    /// unless the program was supplied by a flag (`-e`, `-f`). Drops that leading program so only
    /// the files remain.
    private static func scriptToolOperands(
        _ args: [String], scriptFlags: Set<String>, valueFlags: Set<String> = []
    ) -> [String] {
        let consuming = valueFlags.union(scriptFlags)
        let positionals = ShellLexer.positionals(args, valueFlags: consuming).filter { !$0.isEmpty }
        let scriptGivenByFlag = args.contains { arg in
            scriptFlags.contains(arg) || scriptFlags.contains(where: { arg.hasPrefix("\($0)=") })
        }
        return scriptGivenByFlag ? positionals : Array(positionals.dropFirst())
    }

    /// File operands of a `perl` one-liner. Perl's program arrives in `-e`/`-E`, which may be
    /// bundled (`-pi -e`, `-pe`), so a value-flag set can't find it — scan for any short cluster
    /// carrying `e` and skip the argument it consumes. With no `-e` the first positional is a
    /// script file and is dropped, as with `sed`.
    private static func perlOperands(_ args: [String]) -> [String] {
        var files: [String] = []
        var sawProgramFlag = false
        var skipNext = false
        for arg in args {
            if skipNext { skipNext = false; continue }
            guard arg.hasPrefix("-"), arg != "-" else { files.append(arg); continue }
            if arg.hasPrefix("--") { continue }
            let letters = arg.dropFirst().prefix(while: { $0 != "." })
            if letters.contains("e") || letters.contains("E") { sawProgramFlag = true; skipNext = true }
        }
        return sawProgramFlag ? files : Array(files.dropFirst())
    }

    // MARK: - Redirection parsing

    /// Splits a segment into its command text (every redirection removed) and the targets of its
    /// **output** redirections (`>`, `>>`, `>|`, `&>`, `2>`). Input redirections are dropped from
    /// the command text but not reported — they write nothing. Fd duplication (`>&1`, `2>&1`) and
    /// process substitution (`>(…)`) are not plain file writes and are skipped. Handles both the
    /// spaced (`> file`) and attached (`>file`, `1>file`) forms; heredocs are already stripped by
    /// the caller.
    ///
    /// Removing redirections from the command text is what lets the operand readers above work:
    /// `patch -p1 < changes.diff` must tokenize as `patch -p1`, not as a `patch` with two
    /// positional operands, or the diff would be mistaken for the file being patched.
    static func splitRedirections(_ segment: String) -> (command: String, writeTargets: [String]) {
        var command = ""
        var targets: [String] = []
        let chars = Array(segment)
        var i = 0
        var quote: Character? = nil
        while i < chars.count {
            let c = chars[i]
            if let q = quote { command.append(c); if c == q { quote = nil }; i += 1; continue }
            if c == "'" || c == "\"" { quote = c; command.append(c); i += 1; continue }
            guard c == ">" || c == "<" else { command.append(c); i += 1; continue }

            // A leading fd digit (`2>`, `1>`) was already appended to `command`; drop it back off,
            // along with an `&` from the `&>` form.
            var fdDigits = ""
            while let last = command.last, last.isNumber { fdDigits.insert(last, at: fdDigits.startIndex); command.removeLast() }
            if fdDigits.isEmpty, command.last == "&" { command.removeLast() }

            let isWrite = c == ">"
            var j = i + 1
            if j < chars.count && chars[j] == c { j += 1 }               // ">>" / "<<"
            if j < chars.count && chars[j] == "|" { j += 1 }             // ">|"
            if j < chars.count && chars[j] == "&" {                      // ">&2", "<&-": fd dup
                j += 1
                while j < chars.count && !chars[j].isWhitespace { j += 1 }
                i = j
                continue
            }
            if j < chars.count && chars[j] == "(" {                      // ">(…)": process substitution
                command.append(fdDigits)
                command.append(c)
                i += 1
                continue
            }
            while j < chars.count && (chars[j] == " " || chars[j] == "\t") { j += 1 }

            var target = ""
            var tq: Character? = nil
            while j < chars.count {
                let d = chars[j]
                if let q = tq { if d == q { tq = nil } else { target.append(d) }; j += 1; continue }
                if d == "'" || d == "\"" { tq = d; j += 1; continue }
                if d == " " || d == "\t" || d == ">" || d == "<" || d == "|" { break }
                target.append(d)
                j += 1
            }
            if isWrite && !target.isEmpty { targets.append(target) }
            command.append(" ")
            i = j
        }
        return (command, targets)
    }

    // MARK: - Helpers

    /// Whether a write target is a throwaway/out-of-tree destination the lock system should never
    /// claim (`/dev/null`, `/tmp/…`, `$TMPDIR/…`, macOS per-user temp, fd targets, git plumbing).
    static func isEphemeral(_ rawTarget: String) -> Bool {
        let t = rawTarget.trimmingCharacters(in: CharacterSet(charactersIn: " \t'\""))
        if t.isEmpty || t == "-" { return true }
        if t.hasPrefix("&") { return true }              // fd dup (>&1)
        let ephemeralPrefixes = [
            "/dev/", "/tmp/", "/private/tmp/", "/var/tmp/", "/private/var/tmp/", "/var/folders/",
            "$TMPDIR", "${TMPDIR", "/dev/null",
        ]
        if ephemeralPrefixes.contains(where: { t.hasPrefix($0) }) { return true }
        // git plumbing / internal metadata is tooling, not source authoring.
        if t.contains("/.git/") || t.hasPrefix(".git/") { return true }
        return false
    }

    /// Whether a write target names one concrete file — the precondition for locking it. Anything
    /// the shell would still expand (a `$VAR`, a command substitution, a glob, `~`) names a set we
    /// can't resolve without running it, so it can't be arbitrated and is denied instead.
    static func isNameable(_ rawTarget: String) -> Bool {
        let t = rawTarget.trimmingCharacters(in: CharacterSet(charactersIn: " \t'\""))
        guard !t.isEmpty else { return false }
        if t.hasPrefix("~") { return false }
        return !t.contains(where: { "$`*?[".contains($0) })
    }

    /// Whether a tokenized command carries an in-place flag whose letter is `letter` — matches a
    /// bundled short flag (`-i`, `-pi`, `-i.bak`) as well as the long `--in-place` form.
    private static func hasInPlaceFlag(_ tokens: [String], letter: Character) -> Bool {
        for t in tokens {
            if t == "--in-place" || t.hasPrefix("--in-place=") { return true }
            guard t.hasPrefix("-"), !t.hasPrefix("--") else { continue }
            // A short-flag cluster like `-i`, `-i.bak`, `-pi`, `-ne` — the in-place letter counts
            // even when bundled with other single-letter flags before an optional `.suffix`.
            let body = t.dropFirst()
            let flagLetters = body.prefix(while: { $0 != "." })
            if flagLetters.contains(letter) { return true }
        }
        return false
    }

    /// The basename of a command word (`/usr/bin/sed` → `sed`), lowercased.
    private static func commandBasename(_ word: String) -> String {
        (word as NSString).lastPathComponent.lowercased()
    }

    private static func orderedUnique(_ items: [String]) -> [String] {
        var seen = Set<String>()
        var out: [String] = []
        for item in items where seen.insert(item).inserted { out.append(item) }
        return out
    }

    private static func denialMessage(culprit: String) -> String {
        "Blocked: `\(culprit)` mutates a file in the workspace, but its target isn't a path Nucleic "
            + "can name before the command runs, so the file can't be locked against other agents. "
            + "Either name the file explicitly in the command, or use the structured edit tool "
            + "(Edit/Write, apply_patch, or the ACP edit). Raw-shell writes to a *named* file are "
            + "fine — they take the same locks an edit does. For scratch output, redirect to /tmp "
            + "or /dev/null."
    }
}
