#!/usr/bin/env python3
"""Local CI orchestrator for pyaxml-ng.

Runs jobs from .gitlab-ci.yml on this host, automatically routing jobs whose
`tags:` match a configured runner over SSH to a Windows or macOS VM. Linux
jobs (and any job without a configured matching tag) are delegated to
`gitlab-ci-local` which executes them in Docker.

The .gitlab-ci.yml is unchanged so the same file keeps working on gitlab.com.

SSH-routed execution:
- Source tree synced with tar-pipe over ssh (tar is built-in on macOS and
  Windows 10+; rsync is not required on the remote).
- Job's resolved `before_script + script` is rendered to a single bash or
  PowerShell script, dropped on the remote via `cat >`, and executed.
- GitLab predefined CI variables (CI_PROJECT_DIR, CI_COMMIT_TAG, ...) are
  exported in both env and shell-local scope so YAML expressions like
  `${CI_PROJECT_DIR}` work in PowerShell as they do on real GitLab runners.
- Files declared under `artifacts.paths` are pulled back with rsync.
"""

from __future__ import annotations

import argparse
import os
import shlex
import subprocess
import sys
from pathlib import Path
from typing import Any

import yaml


REPO_ROOT = Path(__file__).resolve().parent.parent
HOSTS_FILE = REPO_ROOT / ".ci-local-hosts.yml"
LOCAL_VARS_FILE = REPO_ROOT / ".gitlab-ci-local-variables.yml"
HOME_VARS_FILE = Path.home() / ".gitlab-ci-local" / "variables.yml"

SYNC_EXCLUDES = (
    ".git",
    ".gitlab-ci-local",
    "target",
    ".venv",
    "venv",
    "node_modules",
    "build",
    "dist",
    "__pycache__",
    "*.pyc",
    ".mypy_cache",
    ".pytest_cache",
    "test_skype.apk",
)


# ---------------------------------------------------------------------------
# Pretty output


def _supports_colour() -> bool:
    return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None


def _c(text: str, code: str) -> str:
    return f"\033[{code}m{text}\033[0m" if _supports_colour() else text


def info(msg: str) -> None:
    print(_c(f"[ci-local] {msg}", "1;36"), flush=True)


def warn(msg: str) -> None:
    print(_c(f"[ci-local] {msg}", "1;33"), flush=True)


def fail(msg: str) -> None:
    print(_c(f"[ci-local] ERROR: {msg}", "1;31"), file=sys.stderr, flush=True)


# ---------------------------------------------------------------------------
# Pipeline / hosts loading


def load_pipeline(extra_vars: dict[str, str]) -> dict[str, Any]:
    """Run gitlab-ci-local --preview to get a fully expanded pipeline."""
    env = os.environ.copy()
    env.update(extra_vars)
    cmd = ["gitlab-ci-local", "--preview"]
    res = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True)
    if res.returncode != 0:
        fail("gitlab-ci-local --preview failed:\n" + res.stderr)
        raise SystemExit(res.returncode)
    return yaml.safe_load(res.stdout) or {}


def list_jobs(pipeline: dict[str, Any]) -> dict[str, dict[str, Any]]:
    reserved = {
        "stages", "variables", "default", "include", "workflow", "image",
        "services", "cache",
    }
    return {
        name: body
        for name, body in pipeline.items()
        if isinstance(body, dict) and name not in reserved and not name.startswith(".")
    }


def load_hosts() -> list[dict[str, Any]]:
    if not HOSTS_FILE.exists():
        return []
    data = yaml.safe_load(HOSTS_FILE.read_text()) or {}
    runners = data.get("runners") or []
    if not isinstance(runners, list):
        fail(f"{HOSTS_FILE.name}: 'runners' must be a list")
        raise SystemExit(2)
    return runners


def match_runner(tags: list[str], runners: list[dict[str, Any]]) -> dict[str, Any] | None:
    tagset = {str(t) for t in tags or []}
    for r in runners:
        if tagset & {str(t) for t in r.get("tags") or []}:
            return r
    return None


# ---------------------------------------------------------------------------
# Script extraction


def script_lines(job: dict[str, Any]) -> list[str]:
    out: list[str] = []
    for key in ("before_script", "script"):
        v = job.get(key)
        if not v:
            continue
        if isinstance(v, str):
            out.append(v)
            continue
        for item in v:
            if isinstance(item, list):
                out.extend(str(x) for x in item)
            else:
                out.append(str(item))
    return out


def _git_remote_info() -> dict[str, str]:
    """Best-effort derivation of CI_PROJECT_URL / CI_SERVER_URL from git remote."""
    try:
        url = subprocess.check_output(
            ["git", "-C", str(REPO_ROOT), "remote", "get-url", "origin"],
            text=True, stderr=subprocess.DEVNULL,
        ).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return {}
    # Normalise scp-style git@gitlab.com:group/proj.git → https://gitlab.com/group/proj
    if url.startswith("git@"):
        host, path = url[4:].split(":", 1)
        url = f"https://{host}/{path}"
    if url.endswith(".git"):
        url = url[:-4]
    # url like https://gitlab.com/group/proj
    try:
        scheme, rest = url.split("://", 1)
        host, path = rest.split("/", 1)
    except ValueError:
        return {}
    namespace, _, name = path.rpartition("/")
    return {
        "CI_PROJECT_URL": url,
        "CI_PROJECT_PATH": path,
        "CI_PROJECT_PATH_SLUG": path.replace("/", "-").lower(),
        "CI_PROJECT_NAMESPACE": namespace,
        "CI_PROJECT_NAME": name or "pyaxml",
        "CI_SERVER_URL": f"{scheme}://{host}",
        "CI_SERVER_HOST": host,
        "CI_REGISTRY": f"registry.{host}",
        "CI_REGISTRY_IMAGE": f"registry.{host}/{path}",
    }


def load_local_variables() -> dict[str, str]:
    """Read secrets / overrides from .gitlab-ci-local-variables.yml.

    Same files gitlab-ci-local itself consumes for Docker jobs; we read them
    too so SSH-routed jobs see the same values.
    """
    out: dict[str, str] = {}
    for path in (HOME_VARS_FILE, LOCAL_VARS_FILE):
        if not path.exists():
            continue
        data = yaml.safe_load(path.read_text()) or {}
        if not isinstance(data, dict):
            warn(f"{path} is not a mapping; skipped")
            continue
        for k, v in data.items():
            out[str(k)] = str(v)
    return out


def ci_variables(job_name: str, tag: str, workdir: str) -> dict[str, str]:
    ref_name = tag or "local"
    base = {
        "CI": "true",
        "GITLAB_CI": "true",
        "CI_PROJECT_NAME": "pyaxml",
        "CI_PROJECT_DIR": workdir,
        "CI_JOB_NAME": job_name,
        "CI_COMMIT_TAG": tag or "",
        "CI_COMMIT_REF_NAME": ref_name,
        "CI_COMMIT_REF_SLUG": ref_name.replace("/", "-").lower()[:63],
        "CI_PIPELINE_ID": "0",
        "CI_COMMIT_SHA": "0" * 40,
    }
    base.update(_git_remote_info())
    return base


# ---------------------------------------------------------------------------
# Local (Docker) execution via gitlab-ci-local


def run_local(
    job_name: str, tag: str | None, extra_vars: dict[str, str], manual: bool,
) -> int:
    info(f"running {job_name} locally via gitlab-ci-local (Docker)")
    cmd = ["gitlab-ci-local"]
    if tag:
        cmd += ["--variable", f"CI_COMMIT_TAG={tag}"]
    # Forward CLI-supplied --var KEY=VAL overrides. gitlab-ci-local already
    # auto-loads .gitlab-ci-local-variables.yml so we don't re-pass those.
    for k, v in extra_vars.items():
        cmd += ["--variable", f"{k}={v}"]
    # gitlab-ci-local needs the job name BOTH as a positional (to restrict
    # the pipeline to just this job, instead of running the whole upstream
    # chain) AND inside `--manual` (to bypass the `when: manual` gate).
    # `--manual` is an array option, so put it last with the job name as
    # its value — otherwise yargs reports "Not enough arguments following".
    cmd.append(job_name)
    if manual:
        cmd += ["--manual", job_name]
    return subprocess.call(cmd, cwd=REPO_ROOT)


# ---------------------------------------------------------------------------
# SSH execution


def _ssh_target(runner: dict[str, Any]) -> str:
    user = runner.get("user")
    host = runner["host"]
    return f"{user}@{host}" if user else str(host)


def _ssh_extra_args(runner: dict[str, Any]) -> list[str]:
    args: list[str] = []
    port = runner.get("port")
    if port:
        args += ["-p", str(port)]
    key = runner.get("identity_file")
    if key:
        args += ["-i", str(key)]
    return args


def _ssh(runner: dict[str, Any], *cmd: str, **kwargs: Any) -> subprocess.CompletedProcess[bytes]:
    return subprocess.run(["ssh", *_ssh_extra_args(runner), _ssh_target(runner), *cmd], **kwargs)


def _ensure_remote_workdir(runner: dict[str, Any]) -> None:
    workdir = runner["workdir"]
    if runner.get("shell", "bash") == "pwsh":
        cmd = (
            f"if (!(Test-Path -LiteralPath '{workdir}')) "
            f"{{ New-Item -ItemType Directory -Force -Path '{workdir}' | Out-Null }}"
        )
        res = _ssh(runner, "pwsh", "-NoProfile", "-Command", cmd)
    else:
        res = _ssh(runner, f"mkdir -p {shlex.quote(workdir)}")
    if res.returncode != 0:
        fail(f"failed to create remote workdir {workdir}")
        raise SystemExit(res.returncode)


def sync_to_remote(runner: dict[str, Any]) -> None:
    """tar-pipe the working tree to the remote workdir."""
    _ensure_remote_workdir(runner)
    workdir = runner["workdir"]
    info(f"syncing source tree → {runner['host']}:{workdir}")

    excludes: list[str] = []
    for ex in SYNC_EXCLUDES:
        excludes += [f"--exclude={ex}"]

    tar_cmd = ["tar", "-cf", "-", *excludes, "-C", str(REPO_ROOT), "."]
    if runner.get("shell", "bash") == "pwsh":
        # bsdtar on Windows reads stdin and -C to change dir.
        remote_extract = f'tar -xf - -C "{workdir}"'
        ssh_cmd = [
            "ssh", *_ssh_extra_args(runner), _ssh_target(runner),
            "pwsh", "-NoProfile", "-Command", remote_extract,
        ]
    else:
        remote_extract = f"tar -xf - -C {shlex.quote(workdir)}"
        ssh_cmd = [
            "ssh", *_ssh_extra_args(runner), _ssh_target(runner), remote_extract,
        ]

    src = subprocess.Popen(tar_cmd, stdout=subprocess.PIPE)
    dst = subprocess.Popen(ssh_cmd, stdin=src.stdout)
    assert src.stdout is not None
    src.stdout.close()  # let SIGPIPE propagate to src if dst dies
    dst.communicate()
    src.wait()
    if src.returncode or dst.returncode:
        fail("source sync failed (tar | ssh)")
        raise SystemExit(dst.returncode or src.returncode)


def _ps_quote(value: str) -> str:
    return "'" + value.replace("'", "''") + "'"


def build_bash_script(
    job_name: str, job: dict[str, Any], tag: str, workdir: str,
    extra_vars: dict[str, str],
) -> str:
    vars_ = ci_variables(job_name, tag, workdir)
    vars_.update({k: str(v) for k, v in (job.get("variables") or {}).items()})
    # Local secrets / overrides win over derived predefined ones.
    vars_.update(extra_vars)
    out = [
        "#!/usr/bin/env bash",
        "set -euo pipefail",
    ]
    for k, v in vars_.items():
        out.append(f"export {k}={shlex.quote(v)}")
    out.append(f"cd {shlex.quote(workdir)}")
    out.append("set -x")
    out.extend(script_lines(job))
    return "\n".join(out) + "\n"


def build_pwsh_script(
    job_name: str, job: dict[str, Any], tag: str, workdir: str,
    extra_vars: dict[str, str],
) -> str:
    vars_ = ci_variables(job_name, tag, workdir)
    vars_.update({k: str(v) for k, v in (job.get("variables") or {}).items()})
    vars_.update(extra_vars)
    out = [
        "$ErrorActionPreference = 'Stop'",
        "$ProgressPreference = 'SilentlyContinue'",
    ]
    # Set both env vars (so child processes see them) and PowerShell-local
    # vars (so YAML's `${CI_PROJECT_DIR}` resolves inside double-quoted
    # strings the way GitLab Runner does).
    for k, v in vars_.items():
        out.append(f"$env:{k} = {_ps_quote(v)}")
        out.append(f"${k} = {_ps_quote(v)}")
    out.append(f"Set-Location {_ps_quote(workdir)}")
    for line in script_lines(job):
        # Echo the command before running, like `set -x`.
        out.append(f"Write-Host ('+ ' + {_ps_quote(line)}) -ForegroundColor DarkGray")
        out.append(line)
    return "\r\n".join(out) + "\r\n"


def write_remote_script(runner: dict[str, Any], content: str, remote_path: str) -> None:
    target = _ssh_target(runner)
    if runner.get("shell", "bash") == "pwsh":
        # On Windows, force binary stdin → file via cmd to avoid CRLF mangling.
        cmd = ["ssh", *_ssh_extra_args(runner), target, "pwsh", "-NoProfile", "-Command",
               f"$bytes = [Console]::OpenStandardInput(); "
               f"$fs = [IO.File]::Create('{remote_path}'); "
               f"$bytes.CopyTo($fs); $fs.Close()"]
    else:
        cmd = ["ssh", *_ssh_extra_args(runner), target,
               f"cat > {shlex.quote(remote_path)}"]
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    p.communicate(content.encode("utf-8"))
    if p.returncode != 0:
        fail("failed to upload remote script")
        raise SystemExit(p.returncode)


def fetch_artifacts(runner: dict[str, Any], paths: list[str]) -> None:
    if not paths:
        return
    target = _ssh_target(runner)
    workdir = runner["workdir"]
    info(f"fetching artifacts: {', '.join(paths)}")
    for p in paths:
        # Strip trailing glob component → fetch parent dir.
        stripped = p.rstrip("/")
        if "*" in stripped or "?" in stripped:
            parent = stripped.rsplit("/", 1)[0] if "/" in stripped else "."
        else:
            parent = stripped
        local_parent = REPO_ROOT / parent
        local_parent.parent.mkdir(parents=True, exist_ok=True)
        remote = f"{target}:{workdir}/{parent}/"
        # rsync handles missing-source gracefully with --ignore-missing-args.
        cmd = ["rsync", "-az", "--ignore-missing-args", *_ssh_extra_args_for_rsync(runner),
               remote, str(local_parent) + "/"]
        res = subprocess.run(cmd)
        if res.returncode != 0:
            warn(f"could not fetch {parent} (rsync exit {res.returncode}); continuing")


def _ssh_extra_args_for_rsync(runner: dict[str, Any]) -> list[str]:
    extra = _ssh_extra_args(runner)
    if not extra:
        return []
    return ["-e", "ssh " + " ".join(shlex.quote(a) for a in extra)]


def run_ssh(
    job_name: str, job: dict[str, Any], tag: str, runner: dict[str, Any],
    extra_vars: dict[str, str],
) -> int:
    sync_to_remote(runner)
    workdir = runner["workdir"]
    shell = runner.get("shell", "bash")
    target = _ssh_target(runner)

    if shell == "pwsh":
        script = build_pwsh_script(job_name, job, tag, workdir, extra_vars)
        remote_script = f"{workdir.rstrip('/')}/.ci-local-job.ps1"
        write_remote_script(runner, script, remote_script)
        info(f"running {job_name} on {runner['host']} (pwsh)")
        rc = subprocess.call([
            "ssh", *_ssh_extra_args(runner), target,
            "pwsh", "-NoProfile", "-File", remote_script,
        ])
    elif shell == "bash":
        script = build_bash_script(job_name, job, tag, workdir, extra_vars)
        remote_script = f"{workdir.rstrip('/')}/.ci-local-job.sh"
        write_remote_script(runner, script, remote_script)
        _ssh(runner, "chmod", "+x", remote_script)
        info(f"running {job_name} on {runner['host']} (bash)")
        rc = subprocess.call([
            "ssh", *_ssh_extra_args(runner), target,
            "bash", remote_script,
        ])
    else:
        fail(f"unknown shell '{shell}' for runner {runner['host']}")
        return 2

    if rc == 0:
        artifacts = (job.get("artifacts") or {}).get("paths") or []
        fetch_artifacts(runner, list(artifacts))
    return rc


# ---------------------------------------------------------------------------
# main


def _detect_git_tag() -> str:
    """Return the tag at HEAD if any, else the most recent reachable tag, else ''.

    Mirrors what GitLab Runner sees: $CI_COMMIT_TAG is only set when the
    pipeline was triggered by a tag push (HEAD == tag). When developing
    locally, fall back to `git describe --tags --abbrev=0` so tag-gated
    jobs (`if: $CI_COMMIT_TAG =~ /^v/`) still match.
    """
    try:
        # Exact tag at HEAD — gives "" if HEAD is not a tag.
        out = subprocess.check_output(
            ["git", "-C", str(REPO_ROOT), "tag", "--points-at", "HEAD"],
            text=True, stderr=subprocess.DEVNULL,
        ).strip()
        if out:
            # Multiple tags possible; pick the first sorted by version.
            return sorted(out.splitlines())[-1]
        out = subprocess.check_output(
            ["git", "-C", str(REPO_ROOT), "describe", "--tags", "--abbrev=0"],
            text=True, stderr=subprocess.DEVNULL,
        ).strip()
        return out
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


def _parse_var_overrides(items: list[str]) -> dict[str, str]:
    out: dict[str, str] = {}
    for item in items:
        if "=" not in item:
            fail(f"--var expects KEY=VALUE, got: {item}")
            raise SystemExit(2)
        k, v = item.split("=", 1)
        out[k.strip()] = v
    return out


def main(argv: list[str] | None = None) -> int:
    global HOSTS_FILE
    ap = argparse.ArgumentParser(
        description="Run pyaxml-ng GitLab CI jobs locally with auto-SSH for Windows/macOS",
    )
    ap.add_argument("--tag", default=None,
                    help="Override CI_COMMIT_TAG. Default: read from `git tag --points-at HEAD` "
                         "or `git describe --tags --abbrev=0`.")
    ap.add_argument("--no-tag", action="store_true",
                    help="Force CI_COMMIT_TAG empty (skip tag-gated jobs).")
    ap.add_argument("--var", action="append", default=[], metavar="KEY=VALUE",
                    help="Inject variable for both Docker and SSH jobs (repeatable). "
                         "Wins over .gitlab-ci-local-variables.yml.")
    ap.add_argument("--manual", action="store_true",
                    help="Run jobs gated `when: manual` (e.g. release:pypi, release:crate).")
    ap.add_argument("--all", action="store_true", help="Run every applicable job")
    ap.add_argument("--skip-release", action="store_true",
                    help="Exclude release:* jobs (they need real tokens)")
    ap.add_argument("--list", action="store_true",
                    help="List jobs and their routing, then exit")
    ap.add_argument("--hosts-file", default=None,
                    help=f"Override path to hosts config (default: {HOSTS_FILE.name})")
    ap.add_argument("jobs", nargs="*", help="Job names to run")
    args = ap.parse_args(argv)

    if args.hosts_file:
        HOSTS_FILE = Path(args.hosts_file).resolve()

    if args.no_tag:
        tag = ""
    elif args.tag is not None:
        tag = args.tag
    else:
        tag = _detect_git_tag()
        if tag:
            info(f"using CI_COMMIT_TAG={tag} (from git)")
        else:
            info("no git tag found; tag-only jobs (rules: $CI_COMMIT_TAG =~ /^v/) will be skipped")

    cli_overrides = _parse_var_overrides(args.var)

    # Variable layering, low → high precedence:
    #   1. derived predefined CI_* (from git remote)        → ci_variables()
    #   2. ~/.gitlab-ci-local/variables.yml                 → load_local_variables()
    #   3. .gitlab-ci-local-variables.yml (in repo)         → load_local_variables()
    #   4. CI_COMMIT_TAG (--tag / git-detected)             → here
    #   5. --var KEY=VALUE on the CLI                       → cli_overrides
    file_vars = load_local_variables()
    extra_vars: dict[str, str] = {}
    extra_vars.update(file_vars)
    if tag:
        extra_vars["CI_COMMIT_TAG"] = tag
    extra_vars.update(cli_overrides)

    # Variables passed to gitlab-ci-local --preview so `rules:` evaluation
    # against $CI_COMMIT_TAG matches what we'll execute.
    preview_vars = dict(extra_vars)
    pipeline = load_pipeline(preview_vars)
    runners = load_hosts()
    jobs = list_jobs(pipeline)

    if args.all:
        selected = list(jobs.keys())
    else:
        selected = list(args.jobs)
    if args.skip_release:
        selected = [j for j in selected if not j.startswith("release:")]

    if not selected:
        fail("no jobs selected (pass job names, or --all)")
        return 2

    if args.list:
        width = max(len(n) for n in selected)
        for name in selected:
            job = jobs.get(name) or {}
            tags = job.get("tags") or []
            r = match_runner(tags, runners)
            where = f"ssh→{r['host']} ({r.get('shell', 'bash')})" if r else "docker (gitlab-ci-local)"
            print(f"  {name:<{width}}  tags={tags}  →  {where}")
        return 0

    if not runners:
        warn(f"no {HOSTS_FILE.name} found; jobs needing Windows/macOS will fail")

    rc_total = 0
    for name in selected:
        if name not in jobs:
            fail(f"job '{name}' not found in pipeline")
            return 2
        job = jobs[name]
        runner = match_runner(job.get("tags") or [], runners)
        if runner:
            if args.manual:
                warn(f"--manual has no effect on SSH-routed job {name}")
            rc = run_ssh(name, job, tag, runner, extra_vars)
        else:
            # gitlab-ci-local already auto-loads .gitlab-ci-local-variables.yml,
            # so we only need to forward --tag and --var overrides explicitly.
            docker_overrides = dict(cli_overrides)
            if tag:
                docker_overrides["CI_COMMIT_TAG"] = tag
            rc = run_local(name, None, docker_overrides, args.manual)
        if rc != 0:
            fail(f"job {name} failed (exit {rc})")
            rc_total = rc
            break
        info(f"{name} OK")
    return rc_total


if __name__ == "__main__":
    sys.exit(main())
