"""Windows terminal behind Warpgate, reached through the server-access relay helper scripts:
  win10-ps.sh  <script>          run PowerShell (sent as -EncodedCommand)
  win10-put.sh <local> <remote>  upload a file (gzip+base64 chunks, SHA-256 verified)
Runs are launched as one-shot scheduled tasks (schtasks) so they outlive the SSH session.
Note: the helper scripts currently hard-code `--target win10`; `target` is checked against that."""
import base64
import hashlib
import os
import re
import subprocess
import tempfile

from . import TransportError

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 WarpgateSshTransport:
    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):
        env = dict(self.cfg.scrubbed_env(), WIN10_TIMEOUT=str(timeout))
        try:
            out = subprocess.run([self.ps_script, script], capture_output=True, text=True, timeout=timeout + 60, env=env)
        except subprocess.TimeoutExpired:
            raise TransportError("relay timeout after %ss" % timeout) from None
        if out.returncode != 0 and not out.stdout.strip():
            raise TransportError("relay exit %s: %s" % (out.returncode, out.stderr.strip()[-300:]))
        return out.stdout.replace("\r", "")

    def put(self, local, remote):
        sha = hashlib.sha256(open(local, "rb").read()).hexdigest()
        out = subprocess.run([self.put_script, local, remote], capture_output=True, text=True, timeout=600,
                             env=self.cfg.scrubbed_env())
        if ("remote sha256 " + sha) not in out.stdout:
            raise TransportError("upload of %s failed: %s" % (remote, (out.stdout + out.stderr).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