from __future__ import annotations

import base64
import json
import os
import sys
import tempfile
from pathlib import Path

AUTH_NAMESPACE = "https://api.openai.com/auth"
PI_PROVIDER_KEY = "openai-codex"


class PiAuthSyncError(RuntimeError):
    pass


class PiAuthSync:
    """Mirrors the active codex account into pi's auth store.

    pi's `openai-codex` provider uses the same OAuth client as the codex CLI
    (client id app_EMoamEEZ73f0CkXaXp7hrann), so pi credentials are derived
    from the account's CODEX_HOME/auth.json. Rotation is non-destructive
    (a rotated-away refresh token stays valid), so pi and codex may each
    refresh their own copy; every switch re-derives from codex's copy.
    """

    def __init__(
        self,
        agent_dir: Path,
        accounts_dir: Path,
        account_home_name: str = "CODEX_HOME",
    ) -> None:
        self.agent_dir = agent_dir
        self.accounts_dir = accounts_dir
        self._account_home_name = account_home_name

    @property
    def active_auth_path(self) -> Path:
        return self.agent_dir / "auth.json"

    def prepare(self, slug: str) -> Path | None:
        """Derive and write the per-account pi auth file; no active-link mutation.

        Returns the per-account file to activate, or None (with a warning)
        when pi is not installed. Raises PiAuthSyncError when pi is installed
        but the account's codex credentials cannot be derived.
        """
        if not self.agent_dir.is_dir():
            print(
                f"pi agent dir '{self.agent_dir}' not found; "
                f"switching codex only for account '{slug}'",
                file=sys.stderr,
            )
            return None

        codex_auth = self.accounts_dir / slug / self._account_home_name / "auth.json"
        credential = self._derive_credential(codex_auth, slug)

        target = self.accounts_dir / slug / "PI_HOME" / "auth.json"
        merged = self._read_json_dict(target)
        merged[PI_PROVIDER_KEY] = credential
        self._atomic_write_private(target, json.dumps(merged, indent=2) + "\n")
        return target

    def activate(self, target: Path) -> None:
        self._atomic_symlink(target, self.active_auth_path)

    def clear_active(self) -> None:
        if self.active_auth_path.is_symlink():
            self.active_auth_path.unlink()

    def _derive_credential(self, codex_auth: Path, slug: str) -> dict[str, object]:
        try:
            payload = json.loads(codex_auth.read_text(encoding="utf-8"))
            tokens = payload["tokens"]
            access = tokens["access_token"]
            refresh = tokens["refresh_token"]
            if not (isinstance(access, str) and access.strip()):
                raise ValueError("empty access_token")
            if not (isinstance(refresh, str) and refresh.strip()):
                raise ValueError("empty refresh_token")
            claims = self._decode_jwt_payload(access.split(".")[1])
            expires_s = claims["exp"]
            if not isinstance(expires_s, (int, float)):
                raise ValueError("access_token exp claim is not numeric")
            auth_claims = claims.get(AUTH_NAMESPACE)
            if not isinstance(auth_claims, dict):
                auth_claims = {}
            account_id = tokens.get("account_id") or auth_claims.get("chatgpt_account_id")
            if not (isinstance(account_id, str) and account_id):
                raise ValueError("missing account_id")
        except Exception as exc:
            raise PiAuthSyncError(
                f"cannot derive pi credentials for account '{slug}' from {codex_auth}: {exc}"
            ) from exc

        return {
            "type": "oauth",
            "access": access,
            "refresh": refresh,
            "expires": int(expires_s * 1000),
            "accountId": account_id,
        }

    @staticmethod
    def _read_json_dict(path: Path) -> dict:
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            return {}
        return payload if isinstance(payload, dict) else {}

    @staticmethod
    def _decode_jwt_payload(segment: str) -> dict:
        padding = "=" * (-len(segment) % 4)
        data = base64.urlsafe_b64decode(segment + padding)
        return json.loads(data.decode("utf-8"))

    @staticmethod
    def _atomic_write_private(path: Path, content: str) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            dir=path.parent,
            delete=False,
        ) as handle:
            handle.write(content)
            os.fchmod(handle.fileno(), 0o600)
            temp_name = handle.name
        os.replace(temp_name, path)

    @staticmethod
    def _atomic_symlink(source: Path, target: Path) -> None:
        if not source.exists():
            raise FileNotFoundError(source)
        target.parent.mkdir(parents=True, exist_ok=True)
        temp_path = target.parent / f".{target.name}.tmp-{os.getpid()}"
        if temp_path.exists() or temp_path.is_symlink():
            temp_path.unlink()
        temp_path.symlink_to(source)
        os.replace(temp_path, target)