#!/usr/bin/env python3
"""Minimal self-update helper for Python scripts."""

from __future__ import annotations

import hashlib
import json
import os
import stat
import sys
import tempfile
from pathlib import Path
from typing import Any, Callable
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urljoin
from urllib.request import Request, urlopen


class UpdateError(RuntimeError):
    """Raised when the update flow cannot be completed safely."""


ReleaseHandler = Callable[[Path, dict[str, Any]], None]


def ensure_updated(
    base_url: str,
    api_key: str,
    project_slug: str,
    current_version: str,
    apply_update: ReleaseHandler | None = None,
) -> bool:
    """
    Check for a newer release and restart the current script after update.

    If `apply_update` is omitted, the current script file is replaced in place.
    If `apply_update` is provided, it receives `(downloaded_artifact_path, release_dict)`
    and is responsible for applying the downloaded artifact before restart.
    """
    if not api_key or api_key == "upsk_replace_me":
        raise UpdateError("set a real API key before calling ensure_updated")

    script_path = _current_script_path()
    release = _load_latest_release(
        base_url=base_url,
        api_key=api_key,
        project_slug=project_slug,
    )

    latest_version = str(release.get("version") or "")
    if not latest_version or latest_version == current_version:
        return False

    download_url = urljoin(base_url.rstrip("/") + "/", str(release.get("download_url") or ""))
    if not download_url:
        raise UpdateError("latest release does not contain a download URL")

    temp_path = _download_release(
        url=download_url,
        api_key=api_key,
        checksum_sha256=str(release.get("checksum_sha256") or ""),
        size_bytes=int(release.get("size_bytes") or 0),
        target_name=script_path.name,
    )

    try:
        if apply_update is None:
            _replace_file(temp_path, script_path)
        else:
            apply_update(temp_path, release)
    finally:
        temp_path.unlink(missing_ok=True)

    _restart_current_process(script_path)
    return True


def _current_script_path() -> Path:
    main_module = sys.modules.get("__main__")
    main_file = getattr(main_module, "__file__", None)
    if not main_file:
        raise UpdateError("cannot resolve the current script path")
    return Path(main_file).resolve()


def _make_request(url: str, api_key: str, accept_json: bool = False) -> Request:
    headers = {
        "Authorization": f"Bearer {api_key}",
        "User-Agent": "autoupdate/1.0",
    }
    if accept_json:
        headers["Accept"] = "application/json"
    return Request(url, headers=headers)


def _load_latest_release(
    base_url: str,
    api_key: str,
    project_slug: str,
) -> dict[str, Any]:
    latest_url = urljoin(
        base_url.rstrip("/") + "/",
        f"api/v1/projects/{quote(project_slug, safe='')}/releases/latest",
    )
    request = _make_request(latest_url, api_key, accept_json=True)

    try:
        with urlopen(request, timeout=15) as response:
            charset = response.headers.get_content_charset() or "utf-8"
            payload = json.loads(response.read().decode(charset))
    except HTTPError as exc:
        if exc.code == 404:
            raise UpdateError(f"no releases found for project '{project_slug}'") from exc
        raise UpdateError(f"HTTP {exc.code} while checking updates: {exc.reason}") from exc
    except URLError as exc:
        raise UpdateError(f"network error while checking updates: {exc.reason}") from exc

    release = payload.get("release")
    if not isinstance(release, dict):
        raise UpdateError("update server returned an unexpected response")
    return release


def _download_release(
    url: str,
    api_key: str,
    checksum_sha256: str,
    size_bytes: int,
    target_name: str,
) -> Path:
    request = _make_request(url, api_key)
    digest = hashlib.sha256()
    written = 0
    tmp_path: Path | None = None

    try:
        with urlopen(request, timeout=60) as response:
            with tempfile.NamedTemporaryFile(
                prefix=target_name + ".",
                suffix=".download",
                delete=False,
            ) as tmp_file:
                tmp_path = Path(tmp_file.name)
                while True:
                    chunk = response.read(64 * 1024)
                    if not chunk:
                        break
                    tmp_file.write(chunk)
                    digest.update(chunk)
                    written += len(chunk)

                tmp_file.flush()
                os.fsync(tmp_file.fileno())
    except HTTPError as exc:
        if tmp_path is not None:
            tmp_path.unlink(missing_ok=True)
        raise UpdateError(f"HTTP {exc.code} while downloading update: {exc.reason}") from exc
    except URLError as exc:
        if tmp_path is not None:
            tmp_path.unlink(missing_ok=True)
        raise UpdateError(f"network error while downloading update: {exc.reason}") from exc
    except Exception:
        if tmp_path is not None:
            tmp_path.unlink(missing_ok=True)
        raise

    actual_sha256 = digest.hexdigest()
    if checksum_sha256 and actual_sha256 != checksum_sha256:
        tmp_path.unlink(missing_ok=True)
        raise UpdateError(
            f"checksum mismatch: expected {checksum_sha256}, got {actual_sha256}"
        )

    if size_bytes and written != size_bytes:
        tmp_path.unlink(missing_ok=True)
        raise UpdateError(f"size mismatch: expected {size_bytes}, got {written}")

    return tmp_path


def _replace_file(downloaded_path: Path, target_path: Path) -> None:
    target_path.parent.mkdir(parents=True, exist_ok=True)

    if target_path.exists():
        current_mode = stat.S_IMODE(target_path.stat().st_mode)
        os.chmod(downloaded_path, current_mode)

    os.replace(downloaded_path, target_path)


def _restart_current_process(script_path: Path) -> None:
    os.execv(sys.executable, [sys.executable, str(script_path), *sys.argv[1:]])