"""Terminals behind Warpgate, reached through the server-access relay client:

    agent-access.py --config <identity> run --target <target> --timeout N -- <argv...>

The client shlex-joins argv and the target account's login shell runs it (cmd.exe on Windows); stdin is
/dev/null, so everything travels in argv. `relay()` is the single place that builds that call.

- os linux / darwin -> PosixWarpgateTransport: every operation is one `bash -c '<script>'` relay call. Files go up
  as base64 chunks (`printf %s '<b64>' | base64 -d`), the runner starts detached with `setsid nohup bash run.sh`
  (its own session: pid == pgid == sid), `poll` reads exit.txt, `kill` signals the process group and session.
- os windows -> WindowsWarpgateTransport: still drives the helper scripts win10-ps.sh (PowerShell
  -EncodedCommand) and win10-put.sh (chunked upload, SHA-256 verified), which hard-code `--target win10`;
  runs are one-shot scheduled tasks (schtasks) so they outlive the SSH session.
  TODO: move Windows onto relay() (`powershell -EncodedCommand` as argv) and drop the helper scripts."""
import base64
import hashlib
import os
import posixpath
import re
import subprocess
import sys
import tempfile

from . import TransportError
from .local import RUNNER as POSIX_RUNNER, sh_quote

DEFAULT_CLIENT = "~/.local/share/server-access-client/agent-access.py"


def relay_argv(cfg, target, argv, timeout):
    """argv for one `agent-access.py run` call (the vault client and identity come from the `vault` config)."""
    v = cfg.data.get("vault") or {}
    cmd = [sys.executable or "python3", "-I", os.path.expanduser(v.get("client") or DEFAULT_CLIENT)]
    if v.get("config"):
        cmd += ["--config", v["config"]]
    return cmd + ["run", "--target", target, "--timeout", str(int(timeout)), "--"] + list(argv)


def call(cfg, argv, timeout, extra_env=None):
    """Run a local relay process (client or helper script) with Plane secrets scrubbed from its environment.
    -> (returncode, stdout, stderr); raises TransportError on timeout / launch failure."""
    env = dict(cfg.scrubbed_env(), **(extra_env or {}))
    try:
        out = subprocess.run(argv, capture_output=True, timeout=timeout + 60, env=env, stdin=subprocess.DEVNULL)
    except subprocess.TimeoutExpired:
        raise TransportError("relay timeout after %ss" % timeout) from None
    except OSError as exc:
        raise TransportError("relay launch failed: %s" % exc) from None
    dec = lambda b: (b or b"").decode("utf-8", "replace").replace("\r", "")  # noqa: E731
    return out.returncode, dec(out.stdout), dec(out.stderr)


def relay(cfg, target, argv, timeout=150):
    """Run argv on a Warpgate target. -> (returncode, stdout, stderr)."""
    return call(cfg, relay_argv(cfg, target, argv, timeout), timeout)


RUNNER = r"""$ErrorActionPreference = 'Continue'
[Console]::OutputEncoding = [Text.Encoding]::UTF8
$OutputEncoding = [Text.Encoding]::UTF8
try {{ Set-Location -LiteralPath {cwd} -ErrorAction Stop }} catch {{
  Set-Content -Path '{d}\exit.txt' -Value ('exit=97 at ' + (Get-Date).ToString('s') + ' (cd failed)'); exit 97 }}
Set-Content -Path '{d}\pid.txt' -Value $PID
Set-Content -Path '{d}\started.txt' -Value (Get-Date).ToString('s')
$prompt = Get-Content -Raw -Encoding UTF8 '{d}\prompt.md'
$prompt | {command} 2> '{d}\stderr.txt' | ForEach-Object {{ Add-Content -Path '{d}\{stdout}' -Value $_ -Encoding UTF8 }}
Set-Content -Path '{d}\exit.txt' -Value ("exit=" + $LASTEXITCODE + " at " + (Get-Date).ToString('s'))
"""


def _b64_payload(out, marker):
    m = re.search(marker + r"=([A-Za-z0-9+/=]*)", out)
    if not m:
        return None
    return base64.b64decode(m.group(1)).decode("utf-8", "replace") if m.group(1) else ""


class WindowsWarpgateTransport:
    os = "windows"

    def __init__(self, cfg, terminal):
        self.cfg, self.terminal = cfg, terminal
        t = cfg.terminals[terminal]
        self.target = t.get("target", terminal)
        helpers = cfg.helpers_dir(terminal)
        self.ps_script, self.put_script = str(helpers / "win10-ps.sh"), str(helpers / "win10-put.sh")
        self.repo = t["repo"]

    def default_run_root(self):
        return self.join(self.repo, "runtime", "agentd")

    def join(self, *parts):
        return "\\".join(p.rstrip("\\") for p in parts)

    def quote(self, s):
        return "'" + str(s).replace("'", "''") + "'"

    # ---------------------------------------------------------------- relay primitives
    def ps(self, script, timeout=150):
        rc, out, err = call(self.cfg, [self.ps_script, script], timeout, {"WIN10_TIMEOUT": str(timeout)})
        if rc != 0 and not out.strip():
            raise TransportError("relay exit %s: %s" % (rc, err.strip()[-300:]))
        return out

    def put(self, local, remote):
        with open(local, "rb") as fh:
            sha = hashlib.sha256(fh.read()).hexdigest()
        rc, out, err = call(self.cfg, [self.put_script, local, remote], 540)
        if ("remote sha256 " + sha) not in out:
            raise TransportError("upload of %s failed: %s" % (remote, (out + err).strip()[-300:]))

    # ---------------------------------------------------------------- file ops
    def mkdir(self, path, clean=False):
        q = self.quote(path)
        script = "New-Item -ItemType Directory -Force -Path %s | Out-Null\n" % q
        if clean:
            script += "Get-ChildItem -LiteralPath %s -File | Remove-Item -Force\n" % q
        self.ps(script + "'MKDIR=ok'")

    def write_text(self, path, text):
        fd, tmp = tempfile.mkstemp(prefix="agentd-", suffix=".txt")
        try:
            with os.fdopen(fd, "w", encoding="utf-8", newline="") as fh:
                fh.write(text)
            self.put(tmp, path)
        finally:
            os.unlink(tmp)

    def read_text(self, path):
        q = self.quote(path)
        out = self.ps("if (Test-Path -LiteralPath %s) { 'B64=' + [Convert]::ToBase64String([IO.File]::ReadAllBytes(%s)) }"
                      " else { 'MISSING' }" % (q, q), timeout=200)
        if "MISSING" in out and "B64=" not in out:
            return None
        text = _b64_payload(out, "B64")
        if text is None:
            raise TransportError("read %s: unexpected relay output" % path)
        return text.lstrip("")

    def grep_first(self, path, needle):
        q = self.quote(path)
        out = self.ps("$m = if (Test-Path -LiteralPath %s) { Select-String -LiteralPath %s -SimpleMatch -Pattern %s | Select-Object -First 1 }\n"
                      "'B64=' + $(if ($m) { [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($m.Line)) } else { '' })"
                      % (q, q, self.quote(needle)))
        return _b64_payload(out, "B64") or None

    def tail(self, path, n):
        q = self.quote(path)
        out = self.ps("$t = if (Test-Path -LiteralPath %s) { (Get-Content -LiteralPath %s -Encoding UTF8 -Tail %d) -join \"`n\" } else { '' }\n"
                      "'B64=' + [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes([string]$t))" % (q, q, n))
        return _b64_payload(out, "B64") or ""

    # ---------------------------------------------------------------- processes
    def runner(self, run_dir, cwd, command, stdout_name):
        return "run.ps1", RUNNER.format(cwd=self.quote(cwd), d=run_dir.replace("'", "''"), command=command, stdout=stdout_name)

    def launch_script(self, task, script_path):
        """PowerShell that (re)creates and runs the one-shot task. Values sit in '...' literals (`'` doubled). The
        script path is wrapped in \\" so that Windows PowerShell 5.1 (legacy native quoting) hands schtasks
        /TR `... -File "<path>"` intact even when the path contains spaces."""
        return ("schtasks /Delete /F /TN '%(t)s' 2>$null | Out-Null\n"
                "$tr = 'powershell.exe -NoProfile -NonInteractive -ExecutionPolicy Bypass -WindowStyle Hidden -File \\\"%(s)s\\\"'\n"
                "schtasks /Create /F /TN '%(t)s' /SC ONCE /ST 00:00 /TR $tr | Out-Null\n"
                "schtasks /Run /TN '%(t)s' | Out-Null\n"
                "'LAUNCHED=' + $LASTEXITCODE" % dict(t=task.replace("'", "''"), s=script_path.replace("'", "''")))

    def launch(self, task, script_path, run_dir):
        out = self.ps(self.launch_script(task, script_path))
        if "LAUNCHED=0" not in out:
            raise TransportError("schtasks launch of %s failed: %s" % (task, out.strip()[-300:]))

    def kill(self, task, run_dir):
        self.ps("$p = Join-Path %s 'pid.txt'\n"
                "if (Test-Path $p) { taskkill /T /F /PID (Get-Content $p | Select-Object -First 1) 2>&1 | Out-Null }\n"
                "schtasks /End /TN '%s' 2>&1 | Out-Null\n"
                "schtasks /Delete /F /TN '%s' 2>&1 | Out-Null\n'KILLED=ok'" % (self.quote(run_dir), task, task))

    def forget(self, task):
        self.ps("schtasks /Delete /F /TN '%s' 2>&1 | Out-Null; 'FORGOT=ok'" % task)

    def process_count(self, name):
        try:
            out = self.ps("'COUNT=' + @(Get-Process -Name %s -ErrorAction SilentlyContinue).Count" % self.quote(name))
        except TransportError:
            return None
        m = re.search(r"COUNT=(\d+)", out)
        return int(m.group(1)) if m else None


# ==================================================================== Linux / macOS
DETACH = {
    # bash puts no background job in its own process group (no job control), so setsid(1) does not fork:
    # $! is the runner itself and becomes pid == pgid == sid.
    "linux": "setsid nohup bash %s",
    # macOS has no setsid(1); perl's POSIX::setsid does the same in-process before exec.
    "darwin": "nohup perl -MPOSIX -e 'POSIX::setsid() or die \"setsid: $!\"; exec @ARGV' bash %s",
}


class PosixWarpgateTransport:
    """Linux/macOS terminal through the relay. Each primitive is one `bash -c` round trip (~10 s through
    Vault + Warpgate), so `start_run` stages the run directory and launches it in a single call."""
    CHUNK = 60000    # base64 chars per call: one argv element must stay below 128 KiB (MAX_ARG_STRLEN) on both ends

    def __init__(self, cfg, terminal):
        self.cfg, self.terminal = cfg, terminal
        t = cfg.terminals[terminal]
        self.os = t["os"]
        self.target = t.get("target", terminal)
        self.repo = t["repo"]

    def default_run_root(self):
        return posixpath.join(self.repo, "runtime", "agentd")

    def join(self, *parts):
        return posixpath.join(*parts)

    def quote(self, s):
        return sh_quote(s)

    # ---------------------------------------------------------------- relay primitive
    def sh(self, script, timeout=150, marker=None):
        """Run a bash script on the target. With `marker`, its absence in stdout is an error."""
        rc, out, err = relay(self.cfg, self.target, ["bash", "-c", script], timeout)
        if (marker and marker not in out) or (not marker and rc != 0 and not out.strip()):
            raise TransportError("relay %s exit %s: %s" % (self.target, rc, (out + err).strip()[-300:]))
        return out

    # ---------------------------------------------------------------- file ops
    def mkdir_script(self, path, clean=False):
        q = self.quote(path)
        s = "mkdir -p %s || exit 1\n" % q
        if clean:
            s += "find %s -maxdepth 1 -type f -exec rm -f {} + || exit 1\n" % q
        return s

    def mkdir(self, path, clean=False):
        self.sh(self.mkdir_script(path, clean) + "echo MKDIR=ok\n", marker="MKDIR=ok")

    @staticmethod
    def b64(text):
        return base64.b64encode(text.encode("utf-8")).decode("ascii")

    def put_script(self, path, b64):
        """Write a base64 payload (alphabet has no quote characters) to path atomically, echo its byte size."""
        q, part = self.quote(path), self.quote(path + ".part")
        return ("printf %%s '%s' | base64 -d > %s && mv -f %s %s || exit 1\n"
                "echo \"WROTE=$(wc -c < %s | tr -d ' ')\"\n" % (b64, part, part, q, q))

    def write_text(self, path, text):
        data = self.b64(text)
        size = len(text.encode("utf-8"))
        if len(data) <= self.CHUNK:
            out = self.sh(self.put_script(path, data), marker="WROTE=")
        else:  # append chunks to <path>.b64, decode on the last call
            q, qb = self.quote(path), self.quote(path + ".b64")
            chunks = [data[i:i + self.CHUNK] for i in range(0, len(data), self.CHUNK)]
            for i, c in enumerate(chunks):
                s = "printf %%s '%s' %s %s || exit 1\n" % (c, ">" if i == 0 else ">>", qb)
                if i < len(chunks) - 1:
                    self.sh(s + "echo CHUNK=ok\n", marker="CHUNK=ok")
                else:
                    out = self.sh(s + "base64 -d < %s > %s && rm -f %s || exit 1\necho \"WROTE=$(wc -c < %s | tr -d ' ')\"\n"
                                  % (qb, q, qb, q), marker="WROTE=")
        if re.findall(r"WROTE=(\d+)", out) != [str(size)]:
            raise TransportError("upload of %s: size mismatch (%s)" % (path, out.strip()[-200:]))

    def read_text(self, path):
        q = self.quote(path)
        out = self.sh("if [ -f %s ]; then printf B64=; base64 < %s | tr -d '\\n'; echo; else echo MISSING; fi" % (q, q),
                      timeout=200)
        if "B64=" not in out:
            if "MISSING" in out:
                return None
            raise TransportError("read %s: unexpected relay output" % path)
        return _b64_payload(out, "B64")

    def _b64_cmd(self, pipeline):
        out = self.sh("printf B64=; { %s; } 2>/dev/null | base64 | tr -d '\\n'; echo" % pipeline)
        text = _b64_payload(out, "B64")
        if text is None:
            raise TransportError("relay %s: unexpected output" % self.target)
        return text

    def grep_first(self, path, needle):
        line = self._b64_cmd("grep -F -m1 -e %s %s" % (self.quote(needle), self.quote(path)))
        return line.rstrip("\n") or None

    def tail(self, path, n):
        return self._b64_cmd("tail -n %d %s" % (int(n), self.quote(path))).rstrip("\n")

    # ---------------------------------------------------------------- processes
    def runner(self, run_dir, cwd, command, stdout_name):
        return "run.sh", POSIX_RUNNER.format(cwd=self.quote(cwd), d=self.quote(run_dir), command=command,
                                             stdout=stdout_name)

    def launch_script(self, run_dir, script_name):
        """Detach the runner from the SSH session; stdio closed so the channel ends at once. $! -> launch.pid
        (kill falls back to it until run.sh has written pid.txt)."""
        d = self.quote(run_dir)
        return ("cd %s || exit 1\n"
                "%s </dev/null >/dev/null 2>&1 &\n"
                "echo $! > launch.pid\n"
                "echo LAUNCHED=$!\n" % (d, DETACH[self.os] % self.quote(script_name)))

    def launch(self, task, script_path, run_dir):
        self.sh(self.launch_script(run_dir, posixpath.basename(script_path)), marker="LAUNCHED=")

    def start_run(self, task, run_dir, files, script_name):
        """mkdir + clean, write `files` {name: text}, launch `script_name` -- one relay call unless a file is
        larger than CHUNK (those go up first through write_text)."""
        big = {n: t for n, t in files.items() if len(self.b64(t)) > self.CHUNK}
        script = self.mkdir_script(run_dir, clean=True)
        if big:
            self.mkdir(run_dir, clean=True)
            for name, text in big.items():
                self.write_text(self.join(run_dir, name), text)
            script = ""
        sizes = {}
        for name, text in files.items():
            if name not in big:
                script += self.put_script(self.join(run_dir, name), self.b64(text))
                sizes[name] = len(text.encode("utf-8"))
        out = self.sh(script + self.launch_script(run_dir, script_name), marker="LAUNCHED=")
        wrote = re.findall(r"WROTE=(\d+)", out)
        if wrote != [str(v) for v in sizes.values()]:
            raise TransportError("staging %s: size mismatch %s != %s" % (run_dir, wrote, list(sizes.values())))

    def kill_script(self, run_dir):
        """TERM then KILL the runner's process group and session. Skipped once exit.txt exists (the runner is
        done and its pid may be reused)."""
        d = self.quote(run_dir)
        return ("d=%s\n"
                "if [ -f \"$d/exit.txt\" ]; then echo KILLED=finished; exit 0; fi\n"
                "p=$(head -n1 \"$d/pid.txt\" 2>/dev/null); [ -n \"$p\" ] || p=$(head -n1 \"$d/launch.pid\" 2>/dev/null)\n"
                "case \"$p\" in ''|*[!0-9]*) echo KILLED=nopid; exit 0;; esac\n"
                "kill -TERM -- \"-$p\" 2>/dev/null; pkill -TERM -s \"$p\" 2>/dev/null\n"
                "for i in 1 2 3 4 5; do pgrep -s \"$p\" >/dev/null 2>&1 || break; sleep 1; done\n"
                "kill -KILL -- \"-$p\" 2>/dev/null; pkill -KILL -s \"$p\" 2>/dev/null\n"
                "sleep 1; echo \"KILLED=ok left=$(pgrep -s \"$p\" 2>/dev/null | wc -l | tr -d ' ')\"\n" % d)

    def kill(self, task, run_dir):
        self.sh(self.kill_script(run_dir), marker="KILLED=")

    def forget(self, task):
        pass

    def process_count(self, name):
        try:
            out = self.sh("echo COUNT=$(pgrep -x %s | wc -l | tr -d ' ')" % self.quote(name), timeout=60)
        except TransportError:
            return None
        m = re.search(r"COUNT=(\d+)", out)
        return int(m.group(1)) if m else None


def WarpgateSshTransport(cfg, terminal):
    """Factory: the transport class for a warpgate-ssh terminal follows its `os`."""
    if cfg.terminals[terminal]["os"] == "windows":
        return WindowsWarpgateTransport(cfg, terminal)
    return PosixWarpgateTransport(cfg, terminal)