"""GitHub App installation token provider (multi-org).

Generates short-lived installation tokens by authenticating as a GitHub App
via RS256-signed JWTs.  Supports multiple installations (orgs/users) by
auto-discovering them via ``GET /app/installations``.

Environment variables
---------------------
``GITHUB_APP_ID``
    Numeric App ID (from the GitHub App settings page).
``GITHUB_APP_PRIVATE_KEY_PATH``
    Path to the PEM private key file downloaded from the App settings.
``GITHUB_APP_INSTALLATION_ID``
    (Optional, backward-compat) Static installation ID.  When set and no
    ``owner`` is provided to ``get_github_token()``, this installation is used.

Nyx identity (separate GitHub App for AI code-review tasks)
-----------------------------------------------------------
``NYX_GITHUB_APP_ID``
    Numeric App ID for the Nyx GitHub App (e.g. "nyx-ai" or "nyx-bot").
``NYX_GITHUB_APP_PRIVATE_KEY_PATH``
    Path to the PEM private key for Nyx's GitHub App.
``NYX_GITHUB_APP_INSTALLATION_ID``
    (Optional) Static installation ID for Nyx's app.

When ``NYX_GITHUB_APP_*`` vars are set, ``get_nyx_github_token()`` returns a
short-lived installation token that lets ``gh`` CLI commands run under Nyx's
GitHub identity rather than the peon bot's identity.
"""

from __future__ import annotations

import calendar
import json
import logging
import os
import time
import urllib.error
import urllib.request
from urllib.parse import urlparse

import jwt  # PyJWT

log = logging.getLogger("peon-loop")

# Safety buffer: refresh the token 10 minutes before it actually expires
# so we never hand a near-expired token to a subprocess.
_EXPIRY_BUFFER_SECONDS = 600

# How long to cache the installations list (seconds).
_INSTALLATIONS_TTL = 3600  # 1 hour


def _parse_github_owner(url: str) -> str | None:
    """Extract the owner (org or user) from a GitHub URL.

    >>> _parse_github_owner("https://github.com/pyrabbit/some-repo")
    'pyrabbit'
    >>> _parse_github_owner("git@github.com:pyrabbit/repo.git")
    'pyrabbit'
    >>> _parse_github_owner("not-a-github-url")
    """
    if not url:
        return None

    # HTTPS: https://github.com/owner/repo
    if "github.com/" in url:
        parsed = urlparse(url)
        parts = parsed.path.strip("/").split("/")
        if parts and parts[0]:
            return parts[0]

    # SSH: git@github.com:owner/repo.git
    if "github.com:" in url:
        _, _, rest = url.partition("github.com:")
        parts = rest.strip("/").split("/")
        if parts and parts[0]:
            return parts[0]

    return None


class GitHubAppTokenProvider:
    """Generates and caches GitHub App installation tokens for multiple orgs."""

    def __init__(
        self,
        app_id: str,
        private_key_path: str,
    ) -> None:
        self.app_id = app_id
        self._private_key = self._read_key(private_key_path)

        # Per-installation token cache: installation_id -> (token, expires_at)
        self._tokens: dict[str, tuple[str, float]] = {}

        # Installations cache: owner_login -> installation_id
        self._installations: dict[str, str] = {}
        self._installations_fetched_at: float = 0.0

        # Backward-compat: static installation ID from env var
        self._static_installation_id = os.environ.get(
            "GITHUB_APP_INSTALLATION_ID", ""
        ).strip()

    @staticmethod
    def _read_key(path: str) -> str:
        with open(path) as f:
            return f.read()

    def _generate_jwt(self) -> str:
        """Create a short-lived JWT for authenticating as the GitHub App."""
        now = int(time.time())
        payload = {
            "iat": now - 60,  # slight backdate for clock skew
            "exp": now + 600,  # 10-minute max per GitHub docs
            "iss": self.app_id,
        }
        return jwt.encode(payload, self._private_key, algorithm="RS256")

    def _fetch_installations(self, app_jwt: str) -> dict[str, str]:
        """Call GET /app/installations and return {owner_login: installation_id}."""
        url = "https://api.github.com/app/installations"
        req = urllib.request.Request(
            url,
            method="GET",
            headers={
                "Authorization": f"Bearer {app_jwt}",
                "Accept": "application/vnd.github+json",
                "X-GitHub-Api-Version": "2022-11-28",
            },
        )
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read().decode())

        result: dict[str, str] = {}
        for inst in data:
            account = inst.get("account") or {}
            login = account.get("login", "").lower()
            inst_id = str(inst.get("id", ""))
            if login and inst_id:
                result[login] = inst_id
        return result

    def get_installations(self) -> dict[str, str]:
        """Return cached installations map, refreshing if stale.

        Returns ``{owner_login: installation_id}``.
        """
        now = time.time()
        if self._installations and now < (
            self._installations_fetched_at + _INSTALLATIONS_TTL
        ):
            return dict(self._installations)

        try:
            app_jwt = self._generate_jwt()
            self._installations = self._fetch_installations(app_jwt)
            self._installations_fetched_at = now
            log.info(
                "Discovered %d GitHub App installation(s): %s",
                len(self._installations),
                list(self._installations.keys()),
            )
        except Exception:
            log.exception("Failed to fetch GitHub App installations")
            # Return stale data if available
            if self._installations:
                return dict(self._installations)
            raise

        return dict(self._installations)

    def invalidate_installations(self) -> None:
        """Force re-fetch of the installations list on next access."""
        self._installations_fetched_at = 0.0

    def _request_installation_token(
        self, app_jwt: str, installation_id: str
    ) -> tuple[str, float]:
        """Exchange the App JWT for an installation access token.

        Returns (token, expires_at_epoch).
        """
        url = (
            f"https://api.github.com/app/installations/"
            f"{installation_id}/access_tokens"
        )
        req = urllib.request.Request(
            url,
            method="POST",
            headers={
                "Authorization": f"Bearer {app_jwt}",
                "Accept": "application/vnd.github+json",
                "X-GitHub-Api-Version": "2022-11-28",
            },
        )
        with urllib.request.urlopen(req, timeout=15) as resp:
            data = json.loads(resp.read().decode())

        token: str = data["token"]
        expires_str: str = data["expires_at"].rstrip("Z")
        t = time.strptime(expires_str, "%Y-%m-%dT%H:%M:%S")
        expires_epoch = float(calendar.timegm(t))
        return token, expires_epoch

    def _resolve_installation_id(self, owner: str | None) -> str:
        """Resolve an owner to an installation ID.

        Priority:
        1. ``owner`` provided: look up in auto-discovered installations
        2. ``owner=None`` + ``GITHUB_APP_INSTALLATION_ID`` env var: use that
        3. ``owner=None`` + single installation: use the only one
        """
        if owner:
            installations = self.get_installations()
            owner_lower = owner.lower()
            inst_id = installations.get(owner_lower)
            if inst_id:
                return inst_id
            raise ValueError(
                f"No GitHub App installation found for owner '{owner}'. "
                f"Known installations: {list(installations.keys())}"
            )

        # No owner specified — try static env var first
        if self._static_installation_id:
            return self._static_installation_id

        # Fall back to auto-discovery
        installations = self.get_installations()
        if len(installations) == 1:
            return next(iter(installations.values()))
        if len(installations) == 0:
            raise ValueError(
                "No GitHub App installations found. "
                "Install the App on at least one organization or user account."
            )
        raise ValueError(
            f"Multiple installations found ({list(installations.keys())}). "
            "Pass an owner to specify which installation to use."
        )

    def get_token(self, owner: str | None = None) -> str:
        """Return a valid installation token, generating a new one if needed.

        Args:
            owner: GitHub org or user login. When provided, resolves to the
                   correct installation. When ``None``, uses the static
                   ``GITHUB_APP_INSTALLATION_ID`` or auto-discovery.
        """
        installation_id = self._resolve_installation_id(owner)

        now = time.time()
        cached = self._tokens.get(installation_id)
        if cached:
            token, expires_at = cached
            if now < (expires_at - _EXPIRY_BUFFER_SECONDS):
                return token

        owner_label = owner or installation_id
        log.info("Generating new GitHub App installation token for %s...", owner_label)
        app_jwt = self._generate_jwt()
        token, expires_at = self._request_installation_token(app_jwt, installation_id)
        self._tokens[installation_id] = (token, expires_at)
        log.info(
            "GitHub App token acquired for %s (expires in %.0f minutes)",
            owner_label,
            (expires_at - now) / 60,
        )
        return token

    def get_token_status(self, installation_id: str) -> str:
        """Return the cache status for an installation's token.

        Returns one of: "active", "expired", "none".
        """
        cached = self._tokens.get(installation_id)
        if not cached:
            return "none"
        _, expires_at = cached
        if time.time() < (expires_at - _EXPIRY_BUFFER_SECONDS):
            return "active"
        return "expired"

    def invalidate(self, owner: str | None = None) -> None:
        """Force regeneration on the next ``get_token()`` call.

        Args:
            owner: If provided, only invalidate this owner's token.
                   If ``None``, invalidate all cached tokens.
        """
        if owner:
            owner_lower = owner.lower()
            inst_id = self._installations.get(owner_lower)
            if inst_id:
                self._tokens.pop(inst_id, None)
            # Also try the static installation ID
            if self._static_installation_id:
                self._tokens.pop(self._static_installation_id, None)
        else:
            self._tokens.clear()


# ---------------------------------------------------------------------------
# Module-level singleton
# ---------------------------------------------------------------------------

_provider: GitHubAppTokenProvider | None = None
_provider_init_attempted = False


def _ensure_provider() -> GitHubAppTokenProvider | None:
    """Lazily initialise the provider singleton from env vars."""
    global _provider, _provider_init_attempted
    if _provider_init_attempted:
        return _provider

    _provider_init_attempted = True

    app_id = os.environ.get("GITHUB_APP_ID", "").strip()
    key_path = os.environ.get("GITHUB_APP_PRIVATE_KEY_PATH", "").strip()

    if not (app_id and key_path):
        return None

    if not os.path.isfile(key_path):
        log.warning(
            "GITHUB_APP_PRIVATE_KEY_PATH=%s does not exist",
            key_path,
        )
        return None

    try:
        _provider = GitHubAppTokenProvider(app_id, key_path)
        static_id = os.environ.get("GITHUB_APP_INSTALLATION_ID", "").strip()
        log.info(
            "GitHub App auth configured (app_id=%s%s)",
            app_id,
            f", static_installation_id={static_id}" if static_id else "",
        )
    except Exception:
        log.exception("Failed to initialise GitHub App token provider")
        _provider = None

    return _provider


def get_github_token(owner: str | None = None) -> str:
    """Return a GitHub App installation token.

    Args:
        owner: GitHub org or user login for multi-installation resolution.

    Returns:
        Installation token string, or empty string if the App is not configured.
    """
    provider = _ensure_provider()
    if provider is not None:
        try:
            return provider.get_token(owner=owner)
        except Exception:
            log.exception("Failed to get GitHub App installation token")

    return ""


def is_app_auth() -> bool:
    """Return True if GitHub App authentication is active."""
    return _ensure_provider() is not None


def get_provider() -> GitHubAppTokenProvider | None:
    """Return the provider singleton (or None if not configured)."""
    return _ensure_provider()


# ---------------------------------------------------------------------------
# Nyx identity — separate GitHub App for AI code-review tasks
# ---------------------------------------------------------------------------

_nyx_provider: GitHubAppTokenProvider | None = None
_nyx_provider_init_attempted = False


def _ensure_nyx_provider() -> GitHubAppTokenProvider | None:
    """Lazily initialise the Nyx provider singleton from NYX_GITHUB_APP_* env vars."""
    global _nyx_provider, _nyx_provider_init_attempted
    if _nyx_provider_init_attempted:
        return _nyx_provider

    _nyx_provider_init_attempted = True

    app_id = os.environ.get("NYX_GITHUB_APP_ID", "").strip()
    key_path = os.environ.get("NYX_GITHUB_APP_PRIVATE_KEY_PATH", "").strip()

    if not (app_id and key_path):
        return None

    if not os.path.isfile(key_path):
        log.warning(
            "NYX_GITHUB_APP_PRIVATE_KEY_PATH=%s does not exist",
            key_path,
        )
        return None

    try:
        _nyx_provider = GitHubAppTokenProvider(app_id, key_path)
        # Honour a static installation ID if provided
        static_id = os.environ.get("NYX_GITHUB_APP_INSTALLATION_ID", "").strip()
        if static_id:
            _nyx_provider._static_installation_id = static_id
        log.info(
            "Nyx GitHub App auth configured (app_id=%s%s)",
            app_id,
            f", static_installation_id={static_id}" if static_id else "",
        )
    except Exception:
        log.exception("Failed to initialise Nyx GitHub App token provider")
        _nyx_provider = None

    return _nyx_provider


def get_nyx_github_token(owner: str | None = None) -> str:
    """Return a GitHub App installation token for Nyx's identity.

    Used by the agent loop when spawning code-review tasks so that
    ``gh pr review`` commands appear under Nyx's GitHub App identity rather
    than the peon bot's.

    Args:
        owner: GitHub org or user login for multi-installation resolution.

    Returns:
        Installation token string, or empty string if Nyx's App is not configured.
    """
    provider = _ensure_nyx_provider()
    if provider is not None:
        try:
            return provider.get_token(owner=owner)
        except Exception:
            log.exception("Failed to get Nyx GitHub App installation token")

    return ""


def is_nyx_app_auth() -> bool:
    """Return True if Nyx GitHub App authentication is configured and active."""
    return _ensure_nyx_provider() is not None


def get_nyx_provider() -> GitHubAppTokenProvider | None:
    """Return the Nyx provider singleton (or None if not configured)."""
    return _ensure_nyx_provider()
