#!/usr/bin/env python3
"""Install and start the Windows Provider after a Python runtime is available."""

from __future__ import annotations

import base64
import hashlib
import importlib.util
import importlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time
import site
from importlib.machinery import SourceFileLoader
from urllib.request import Request, urlopen


MANIFEST_URL = "__MANIFEST_URL__"
PROVIDER_URL = "__PROVIDER_URL__"
PROGRESS_INTERVAL = __PROGRESS_INTERVAL_SECONDS__
PROVIDER_ROOT = Path.home() / ".unidesk" / "unideskcprovider"
CONFIG_PATH = Path.home() / ".unidesk" / "provider.yaml"
PROVIDER_PATH = PROVIDER_ROOT / "unideskcprovider.py"
RESTART_SCRIPT_PATH = PROVIDER_ROOT / "restart-provider.ps1"
RESTART_STATUS_PATH = PROVIDER_ROOT / "restart-status.json"


def fail(code: str, message: str) -> None:
    print(
        "ERROR code={} stage=provider-install message={} recovery=check-provider-yaml-and-network".format(
            code, message
        ),
        file=sys.stderr,
    )
    raise SystemExit(1)


def progress(label: str, completed: int, total: int, started: float, source: str) -> None:
    elapsed = max(0.001, time.monotonic() - started)
    speed = completed / elapsed
    eta = max(0, int((total - completed) / speed + 0.5)) if total and speed > 0 else -1
    size = "{:.1f}/{:.1f} MiB".format(completed / 1048576, total / 1048576) if total else "{} bytes".format(completed)
    percent = "{:.1f}%".format(completed * 100 / total) if total else "unknown"
    remaining = "{}s".format(eta) if eta >= 0 else "unknown"
    print(
        "[download] {} {} {} speed={:.1f} KiB/s elapsed={}s eta={} source={}".format(
            label, percent, size, speed / 1024, int(elapsed), remaining, source
        ),
        flush=True,
    )


def download(url: str, path: Path, label: str) -> None:
    request = Request(url, headers={"User-Agent": "unideskcprovider-installer/1"})
    started = time.monotonic()
    try:
        response = urlopen(request, timeout=__CONNECT_TIMEOUT_SECONDS__)
        total = int(response.headers.get("Content-Length") or 0)
        completed = 0
        last = 0.0
        with path.open("wb") as output:
            while True:
                chunk = response.read(1024 * 256)
                if not chunk:
                    break
                output.write(chunk)
                completed += len(chunk)
                now = time.monotonic()
                if now - last >= PROGRESS_INTERVAL:
                    progress(label, completed, total, started, url)
                    last = now
        progress(label, completed, total, started, url)
    except Exception as exc:
        raise RuntimeError("download failed: {}".format(exc)) from exc


def ensure_yaml(pip_index: str) -> None:
    try:
        import yaml  # noqa: F401

        return
    except ImportError:
        pass
    print("[5/6] Installing PyYAML (index={})".format(pip_index), flush=True)
    runtime_root = Path.home() / ".unidesk" / "runtime"
    executable = Path(sys.executable).resolve()
    private_runtime = runtime_root == executable or runtime_root in executable.parents
    target = executable.parent / "Lib" / "site-packages"
    command = [
        sys.executable,
        "-m",
        "pip",
        "install",
        "--disable-pip-version-check",
        "--index-url",
        pip_index,
    ]
    if private_runtime:
        command.extend(["--target", str(target)])
    else:
        command.append("--user")
    command.append("PyYAML")
    result = subprocess.call(command)
    if result:
        fail("pyyaml_install_failed", "pip returned {}".format(result))
    if private_runtime:
        site.addsitedir(str(target))
    else:
        site.addsitedir(site.getusersitepackages())
    importlib.invalidate_caches()
    try:
        import yaml  # noqa: F401
    except ImportError:
        fail("pyyaml_unavailable", "PyYAML is unavailable after pip completed")


def load_provider(path: Path):
    loader = SourceFileLoader("unideskcprovider_install_probe", str(path))
    spec = importlib.util.spec_from_loader(loader.name, loader)
    if spec is None:
        fail("provider_import_failed", "could not create import spec")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    loader.exec_module(module)
    return module


def existing_provider_pids() -> list[int]:
    source = (
        "Get-WmiObject Win32_Process | "
        "Where-Object { $_.CommandLine -match 'unideskcprovider[.]py' } | "
        "ForEach-Object { Write-Output $_.ProcessId }"
    )
    result = subprocess.check_output(
        [
            "powershell.exe",
            "-NoLogo",
            "-NoProfile",
            "-ExecutionPolicy",
            "Bypass",
            "-Command",
            source,
        ],
        creationflags=0x08000000,
        universal_newlines=True,
    )
    return sorted(
        {
            int(line.strip())
            for line in result.splitlines()
            if line.strip().isdigit() and int(line.strip()) != os.getpid()
        }
    )


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


def schedule_provider_handoff(pythonw: Path, old_pids: list[int]) -> None:
    pid_list = ",".join(str(pid) for pid in old_pids)
    script = """$ErrorActionPreference = 'Stop'
$statusPath = {status}
try {{
  Start-Sleep -Seconds 2
  @({pids}) | ForEach-Object {{
    Stop-Process -Id $_ -Force -ErrorAction SilentlyContinue
  }}
  Start-Sleep -Seconds 1
  $process = Start-Process -FilePath {pythonw} -ArgumentList @({provider}) -WorkingDirectory {root} -WindowStyle Hidden -PassThru
  Start-Sleep -Seconds 2
  if ($process.HasExited) {{ throw "Provider exited with code $($process.ExitCode)" }}
  [ordered]@{{ ok = $true; pid = $process.Id; startedAt = (Get-Date).ToString('o') }} |
    ConvertTo-Json -Compress | Set-Content -LiteralPath $statusPath -Encoding UTF8
}} catch {{
  [ordered]@{{ ok = $false; error = $_.Exception.Message; failedAt = (Get-Date).ToString('o') }} |
    ConvertTo-Json -Compress | Set-Content -LiteralPath $statusPath -Encoding UTF8
}}
Remove-Item -LiteralPath $MyInvocation.MyCommand.Path -Force -ErrorAction SilentlyContinue
""".format(
        status=powershell_string(str(RESTART_STATUS_PATH)),
        pids=pid_list,
        pythonw=powershell_string(str(pythonw)),
        provider=powershell_string(str(PROVIDER_PATH)),
        root=powershell_string(str(PROVIDER_ROOT)),
    )
    RESTART_SCRIPT_PATH.write_text(script, encoding="utf-8-sig")
    try:
        RESTART_STATUS_PATH.unlink()
    except OSError:
        pass
    command = (
        'powershell.exe -NoLogo -NoProfile -ExecutionPolicy Bypass '
        '-WindowStyle Hidden -File "{}"'.format(RESTART_SCRIPT_PATH)
    )
    wmi_source = (
        "$result = ([wmiclass]'Win32_Process').Create({}); "
        "if ($result.ReturnValue -ne 0) {{ exit $result.ReturnValue }}; "
        "Write-Output $result.ProcessId"
    ).format(powershell_string(command))
    encoded = base64.b64encode(wmi_source.encode("utf-16le")).decode("ascii")
    result = subprocess.check_output(
        [
            "powershell.exe",
            "-NoLogo",
            "-NoProfile",
            "-ExecutionPolicy",
            "Bypass",
            "-EncodedCommand",
            encoded,
        ],
        creationflags=0x08000000,
        universal_newlines=True,
    ).strip()
    if not result or not result.splitlines()[-1].strip().isdigit():
        fail("provider_handoff_failed", "WMI did not return a restart process id")
    print(
        "[6/6] Provider restart handoff scheduled oldPids={} handoffPid={}".format(
            ",".join(str(pid) for pid in old_pids),
            result.splitlines()[-1].strip(),
        ),
        flush=True,
    )


def main() -> None:
    if os.name != "nt":
        fail("windows_required", "native Windows is required")
    if not CONFIG_PATH.is_file():
        fail("config_missing", str(CONFIG_PATH))
    PROVIDER_ROOT.mkdir(parents=True, exist_ok=True)
    print("[5/6] Downloading and validating Provider release", flush=True)
    manifest_path = PROVIDER_ROOT / "update.json.download"
    provider_path = PROVIDER_ROOT / "unideskcprovider.py.download"
    try:
        download(MANIFEST_URL, manifest_path, "manifest")
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        expected_bytes = int(manifest["bytes"])
        expected_sha = str(manifest["sha256"]).lower()
        download(PROVIDER_URL, provider_path, "provider")
        content = provider_path.read_bytes()
        actual_sha = hashlib.sha256(content).hexdigest()
        if len(content) != expected_bytes or actual_sha != expected_sha:
            fail("artifact_integrity_failed", "Provider bytes or SHA-256 mismatch")
        if 'APP_NAME = "UniDesk Windows Provider"' not in content.decode("utf-8"):
            fail("artifact_identity_failed", "Provider identity marker missing")
        compile_result = subprocess.call([sys.executable, "-m", "py_compile", str(provider_path)])
        if compile_result:
            fail("provider_compile_failed", "py_compile returned {}".format(compile_result))
        ensure_yaml(os.environ.get("UNIDESK_PIP_INDEX", "https://pypi.org/simple"))
        module = load_provider(provider_path)
        config = module.ProviderConfig(CONFIG_PATH)
        old_pids = existing_provider_pids()
        print(
            "[5/6] Preparing Provider handoff oldPids={}".format(
                ",".join(str(pid) for pid in old_pids) if old_pids else "none"
            ),
            flush=True,
        )
        if PROVIDER_PATH.exists():
            backup = PROVIDER_PATH.with_suffix(".py.install.bak")
            try:
                shutil.copy2(str(PROVIDER_PATH), str(backup))
            except OSError:
                pass
        os.replace(str(provider_path), str(PROVIDER_PATH))
        pythonw = Path(sys.executable).with_name("pythonw.exe")
        if not pythonw.is_file():
            pythonw = Path(sys.executable)
        if old_pids:
            schedule_provider_handoff(pythonw, old_pids)
            process_id = "handoff"
        else:
            print("[6/6] Starting Provider runtime={}".format(sys.executable), flush=True)
            provider_environment = os.environ.copy()
            provider_environment.pop("UNIDESK_PIP_INDEX", None)
            process = subprocess.Popen(
                [str(pythonw), str(PROVIDER_PATH)],
                cwd=str(PROVIDER_ROOT),
                creationflags=0x08000000,
                env=provider_environment,
            )
            time.sleep(1)
            if process.poll() is not None:
                fail("provider_start_failed", "Provider exited with code {}".format(process.returncode))
            process_id = str(process.pid)
        print(
            "OK operation=unideskcprovider.install providerId={} version={} python={} pid={}".format(
                config.id, module.APP_VERSION, sys.executable, process_id
            ),
            flush=True,
        )
    except SystemExit:
        raise
    except Exception as exc:
        fail("provider_install_failed", str(exc))
    finally:
        for temporary in (manifest_path, provider_path):
            try:
                temporary.unlink()
            except OSError:
                pass


if __name__ == "__main__":
    main()
