# common/codex_auth.py
"""
Codex OAuth Token Manager (Simplified for Canifa AI Platform).

Loads Codex access tokens from local file ~/.codex/auth.json (from `codex login`).
Auto-refreshes when expired.

Copied from contract-ai, simplified: removed MongoDB dependency.
"""
import base64
import json
import logging
import os
import threading
import time
from pathlib import Path
from typing import Any

import requests

logger = logging.getLogger(__name__)

# OAuth endpoints
_AUTH_TOKEN_URL = "https://auth.openai.com/oauth/token"
_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"  # Codex CLI client ID

# Token cache (in-memory)
_lock = threading.Lock()
_cached_token: str | None = None
_cached_account_id: str | None = None
_token_expiry: float = 0.0


# ─── JWT Helpers ──────────────────────────────────────────────────────

def _decode_jwt_payload(token: str) -> dict:
    """Decode JWT payload without verification. Returns claims dict."""
    try:
        parts = token.split(".")
        if len(parts) != 3:
            return {}
        payload_b64 = parts[1] + "=" * (4 - len(parts[1]) % 4)
        return json.loads(base64.urlsafe_b64decode(payload_b64))
    except Exception:
        return {}


def _decode_jwt_expiry(token: str) -> float:
    """Extract expiry timestamp from JWT without verification."""
    return float(_decode_jwt_payload(token).get("exp", 0))


def _extract_account_id(token: str) -> str | None:
    """Extract chatgpt_account_id from JWT claims."""
    payload = _decode_jwt_payload(token)
    auth_claims = payload.get("https://api.openai.com/auth", {})
    return auth_claims.get("chatgpt_account_id")


# ─── Local File Token Loading ────────────────────────────────────────

def _get_auth_json_path() -> Path:
    """Get path to Codex auth.json file."""
    codex_home = os.getenv("CODEX_HOME", "")
    if codex_home:
        return Path(codex_home) / "auth.json"
    return Path.home() / ".codex" / "auth.json"


def _read_auth_json() -> dict[str, Any] | None:
    """Read and parse auth.json file."""
    path = _get_auth_json_path()
    if not path.exists():
        logger.debug("Codex auth.json not found at %s", path)
        return None
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data
    except Exception as e:
        logger.error("❌ Failed to read Codex auth.json: %s", e)
        return None


# ─── Token Refresh ───────────────────────────────────────────────────

def _refresh_token(refresh_token: str) -> str | None:
    """Refresh the access token using the refresh token."""
    try:
        resp = requests.post(
            _AUTH_TOKEN_URL,
            json={
                "grant_type": "refresh_token",
                "client_id": _CLIENT_ID,
                "refresh_token": refresh_token,
            },
            headers={"Content-Type": "application/json"},
            timeout=15,
        )
        if resp.status_code == 200:
            data = resp.json()
            new_access = data.get("access_token")
            new_refresh = data.get("refresh_token")
            if new_access:
                logger.info("✅ Codex token refreshed successfully")
                _update_auth_json(new_access, new_refresh or refresh_token)
                return new_access
        logger.error("❌ Codex token refresh failed: %d %s", resp.status_code, resp.text[:200])
        return None
    except Exception as e:
        logger.error("❌ Codex token refresh error: %s", e)
        return None


def _update_auth_json(access_token: str, refresh_token: str):
    """Write updated tokens back to auth.json."""
    path = _get_auth_json_path()
    try:
        data = {}
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)

        if "tokens" not in data:
            data["tokens"] = {}
        data["tokens"]["access_token"] = access_token
        data["tokens"]["refresh_token"] = refresh_token
        data["last_refresh"] = time.strftime("%Y-%m-%dT%H:%M:%S.000000000Z", time.gmtime())

        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
        logger.debug("📝 Updated auth.json with refreshed token")
    except Exception as e:
        logger.warning("⚠️ Failed to update auth.json: %s", e)


# ─── Public API ──────────────────────────────────────────────────────

def get_codex_access_token() -> str | None:
    """Get a valid Codex access token from local auth.json.
    Auto-refreshes if expired."""
    global _cached_token, _cached_account_id, _token_expiry

    with _lock:
        now = time.time()

        # Return cached token if still valid (with 60s buffer)
        if _cached_token and _token_expiry > (now + 60):
            return _cached_token

        # Read auth.json
        auth_data = _read_auth_json()
        if not auth_data:
            return None

        tokens = auth_data.get("tokens", {})
        access_token = tokens.get("access_token")
        refresh_token = tokens.get("refresh_token")

        if not access_token:
            logger.warning("⚠️ No access_token in Codex auth.json")
            return None

        # Extract account_id from JWT claims
        _cached_account_id = _extract_account_id(access_token)

        # Check if token is expired
        expiry = _decode_jwt_expiry(access_token)

        if expiry > (now + 60):
            # Token still valid
            _cached_token = access_token
            _token_expiry = expiry
            logger.info("🔑 Codex token valid (expires in %dm)", int((expiry - now) / 60))
            return access_token

        # Token expired → refresh
        if refresh_token:
            logger.info("🔄 Codex token expired, refreshing...")
            new_token = _refresh_token(refresh_token)
            if new_token:
                _cached_token = new_token
                _cached_account_id = _extract_account_id(new_token)
                _token_expiry = _decode_jwt_expiry(new_token)
                return new_token

        # Fallback: return expired token anyway (might still work briefly)
        logger.warning("⚠️ Codex token may be expired, using anyway")
        _cached_token = access_token
        _token_expiry = 0
        return access_token


def get_codex_account_id() -> str | None:
    """Get the ChatGPT account ID from the cached token."""
    if _cached_account_id:
        return _cached_account_id
    get_codex_access_token()
    return _cached_account_id


def get_refresh_token() -> str:
    """Get the refresh token from auth.json."""
    auth_data = _read_auth_json()
    if auth_data:
        return auth_data.get("tokens", {}).get("refresh_token", "")
    return ""


def is_codex_available() -> bool:
    """Check if Codex auth is configured and accessible."""
    return get_codex_access_token() is not None
