"""PAT / OAuth token management — parameterized port of reddit-mcp's token_manager.py.

Includes server-side token verification via auth_gateway's GET /api/me
endpoint, with local file cache to avoid hitting the gateway on every startup.
"""

from __future__ import annotations

import hashlib
import json
import os
import time
from pathlib import Path
from typing import Mapping

import httpx

from .config_helpers import AUTH_URL, ensure_config_dir, ensure_shared_auth_dir
from .shared_auth import read_shared_auth, write_shared_auth

# Verification cache: valid for 1 hour
_VERIFY_CACHE_TTL_S = 3600


class TokenManager:
    """Read/write/refresh PAT and JWT tokens stored in auth.json."""

    def __init__(
        self,
        auth_url: str = AUTH_URL,
        *,
        auth_json_path: Path | str,
        env_token_name: str,
        env: Mapping[str, str] | None = None,
    ) -> None:
        """Initialize TokenManager.
        
        Args:
            auth_url: Auth gateway base URL (default: https://mkterswingman.com)
            auth_json_path: Path to shared auth.json (e.g., ~/.mkterswingman/auth.json)
            env_token_name: Environment variable name to check for PAT (e.g., REDDIT_MCP_TOKEN)
            env: Optional env dict (defaults to os.environ)
        """
        self.auth_url = auth_url
        self.auth_json_path = Path(auth_json_path)
        self.env_token_name = env_token_name
        self.env = env if env is not None else os.environ
        self._verify_cache_path = self.auth_json_path.parent / "verify-cache.json"

    async def get_valid_token(self) -> str | None:
        """Return a valid token, or None if not authenticated.

        Priority: env var > PAT in auth.json > JWT (auto-refresh if expired).
        """
        # 1. Env var PAT takes precedence
        env_pat = self.env.get(self.env_token_name)
        if env_pat:
            return env_pat

        # 2. Read persisted auth
        auth = read_shared_auth(self.auth_json_path)
        if auth is None:
            return None

        if auth.get("type") == "pat" and auth.get("pat"):
            return auth["pat"]

        if auth.get("type") == "jwt":
            access = auth.get("access_token")
            expires_at = auth.get("expires_at", 0)
            if access and time.time() * 1000 < expires_at:
                return access

            # Try refresh
            refresh = auth.get("refresh_token")
            if refresh:
                try:
                    refreshed = await self._refresh_tokens(refresh, auth.get("client_id"))
                    if refreshed:
                        await self.save_tokens(
                            refreshed["access_token"],
                            refreshed["refresh_token"],
                            refreshed["expires_in"],
                            auth.get("client_id"),
                        )
                        return refreshed["access_token"]
                except Exception:
                    pass  # refresh failed → treat as unauthenticated
            return None

        return None

    async def verify_token_with_server(self, token: str) -> dict:
        """Verify token against auth_gateway GET /api/me.

        Client sends its own Bearer token; server validates JWT signature +
        expiry + PAT revocation and returns {active: true/false}.

        Uses a local file cache (1h TTL) to avoid hitting the gateway every startup.
        """
        # Check cache first
        cached = self._read_verify_cache(token)
        if cached is not None:
            return cached

        # Call auth_gateway — standard Bearer auth, no shared secret needed
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.get(
                    f"{self.auth_url}/api/me",
                    headers={"Authorization": f"Bearer {token}"},
                )
                if resp.status_code == 401 or resp.status_code == 403:
                    result = {"active": False}
                elif resp.status_code == 200:
                    result = resp.json()
                else:
                    return {"active": False, "error": f"Server returned {resp.status_code}"}
        except Exception as e:
            return {"active": False, "error": f"Verification request failed: {e}"}

        # Cache result
        self._write_verify_cache(token, result)
        return result

    async def save_tokens(
        self,
        access_token: str,
        refresh_token: str,
        expires_in: int,
        client_id: str | None = None,
    ) -> None:
        """Persist JWT tokens to auth.json."""
        data = {
            "type": "jwt",
            "access_token": access_token,
            "refresh_token": refresh_token,
            # expires_at in ms (matching yt-mcp convention)
            "expires_at": int(time.time() * 1000) + expires_in * 1000,
            "client_id": client_id,
        }
        write_shared_auth(self.auth_json_path, data)

    async def save_pat(self, pat: str) -> None:
        """Persist a PAT to auth.json."""
        data = {"type": "pat", "pat": pat}
        write_shared_auth(self.auth_json_path, data)

    # -- private ----------------------------------------------------------------

    async def _refresh_tokens(
        self, refresh_token: str, client_id: str | None
    ) -> dict | None:
        url = f"{self.auth_url}/oauth/token"
        body: dict[str, str] = {
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
        }
        if client_id:
            body["client_id"] = client_id

        async with httpx.AsyncClient() as client:
            resp = await client.post(url, json=body)
            if resp.status_code != 200:
                return None
            data = resp.json()
            if not data.get("access_token"):
                return None
            return data

    # -- verify cache -----------------------------------------------------------

    def _read_verify_cache(self, token: str) -> dict | None:
        """Read cached verification result if still valid."""
        if not self._verify_cache_path.exists():
            return None
        try:
            cache = json.loads(self._verify_cache_path.read_text("utf-8"))
            # Cache keyed by SHA-256 hash of full token to avoid collisions
            key = hashlib.sha256(token.encode()).hexdigest()[:32]
            entry = cache.get(key)
            if not entry:
                return None
            if time.time() - entry.get("cached_at", 0) > _VERIFY_CACHE_TTL_S:
                return None
            return entry.get("result")
        except (json.JSONDecodeError, OSError):
            return None

    def _write_verify_cache(self, token: str, result: dict) -> None:
        """Write verification result to cache."""
        ensure_config_dir(self._verify_cache_path.parent)
        key = hashlib.sha256(token.encode()).hexdigest()[:32]
        cache: dict = {}
        if self._verify_cache_path.exists():
            try:
                cache = json.loads(self._verify_cache_path.read_text("utf-8"))
            except (json.JSONDecodeError, OSError):
                pass
        cache[key] = {"result": result, "cached_at": time.time()}
        self._verify_cache_path.write_text(json.dumps(cache, indent=2), "utf-8")
        try:
            self._verify_cache_path.chmod(0o600)
        except OSError:
            pass
