"""HTTP client with JWT auth for the Personal Assistant API."""

import time

import httpx

from personal_assistant_mcp.config import Credentials, save_credentials_google


class AuthClient:
    """HTTP client that handles JWT login, token refresh, and authenticated requests."""

    TOKEN_EXPIRY_BUFFER = 60  # refresh 60s before expiry

    def __init__(self, credentials: Credentials) -> None:
        self._credentials = credentials
        self._http = httpx.AsyncClient(base_url=credentials.api_url)
        self._access_token: str | None = credentials.access_token or None
        self._refresh_token: str | None = credentials.refresh_token or None
        self._token_obtained_at: float = 0
        self._token_expires_in: int = 1800  # 30min default

        # If we loaded tokens from config (google auth), mark as recently obtained
        if self._access_token:
            self._token_obtained_at = time.time()

    async def close(self) -> None:
        await self._http.aclose()

    async def login(self) -> None:
        """Authenticate with email/password and store tokens."""
        if self._credentials.auth_method == "google":
            # Google users can't re-login with password, try refresh instead
            if self._refresh_token:
                await self._refresh()
                return
            raise RuntimeError(
                "Google auth session expired. Run 'personal-assistant-mcp setup' again."
            )

        resp = await self._http.post(
            "/api/auth/login",
            json={
                "email": self._credentials.email,
                "password": self._credentials.password,
            },
        )
        resp.raise_for_status()
        self._store_tokens(resp.json())

    async def _refresh(self) -> None:
        """Refresh the access token using the refresh token."""
        resp = await self._http.post(
            "/api/auth/refresh",
            json={"refresh_token": self._refresh_token},
        )
        if resp.status_code == 401:
            if self._credentials.auth_method == "email":
                await self.login()
                return
            raise RuntimeError(
                "Google auth session expired. Run 'personal-assistant-mcp setup' again."
            )
        resp.raise_for_status()
        self._store_tokens(resp.json())

    def _store_tokens(self, data: dict) -> None:
        self._access_token = data["access_token"]
        self._refresh_token = data["refresh_token"]
        self._token_obtained_at = time.time()

        # Persist refreshed tokens for Google auth so they survive restarts
        if self._credentials.auth_method == "google":
            save_credentials_google(self._access_token, self._refresh_token)

    def _token_needs_refresh(self) -> bool:
        if not self._access_token:
            return True
        elapsed = time.time() - self._token_obtained_at
        return elapsed >= (self._token_expires_in - self.TOKEN_EXPIRY_BUFFER)

    async def _ensure_auth(self) -> None:
        """Ensure we have a valid access token."""
        if self._access_token is None:
            await self.login()
        elif self._token_needs_refresh():
            await self._refresh()

    def _auth_headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self._access_token}"}

    async def get(self, path: str, **kwargs) -> httpx.Response:
        await self._ensure_auth()
        resp = await self._http.get(path, headers=self._auth_headers(), **kwargs)
        if resp.status_code == 401:
            await self.login()
            resp = await self._http.get(path, headers=self._auth_headers(), **kwargs)
        resp.raise_for_status()
        return resp

    async def post(self, path: str, **kwargs) -> httpx.Response:
        await self._ensure_auth()
        resp = await self._http.post(path, headers=self._auth_headers(), **kwargs)
        if resp.status_code == 401:
            await self.login()
            resp = await self._http.post(path, headers=self._auth_headers(), **kwargs)
        resp.raise_for_status()
        return resp

    async def put(self, path: str, **kwargs) -> httpx.Response:
        await self._ensure_auth()
        resp = await self._http.put(path, headers=self._auth_headers(), **kwargs)
        if resp.status_code == 401:
            await self.login()
            resp = await self._http.put(path, headers=self._auth_headers(), **kwargs)
        resp.raise_for_status()
        return resp

    async def patch(self, path: str, **kwargs) -> httpx.Response:
        await self._ensure_auth()
        resp = await self._http.patch(path, headers=self._auth_headers(), **kwargs)
        if resp.status_code == 401:
            await self.login()
            resp = await self._http.patch(path, headers=self._auth_headers(), **kwargs)
        resp.raise_for_status()
        return resp
