"""
PostgreSQL + FIPS AEAD provider for user OAuth credentials.

This provider stores encrypted OAuth tokens in PostgreSQL using
FIPS 140-3 compliant AES-256-GCM encryption.

Example:
    from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
    from sqlalchemy.orm import sessionmaker
    from empowernow_common.vault.user_oauth import (
        PostgresUserOAuthProvider,
        create_third_party_credential_model,
    )

    # Setup
    engine = create_async_engine("postgresql+asyncpg://...")
    async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    Base = declarative_base()
    ThirdPartyCredential = create_third_party_credential_model(Base)

    provider = PostgresUserOAuthProvider(
        session_factory=async_session,
        model_class=ThirdPartyCredential,
    )

    # Use
    cred = await provider.get_credential("arn:user:123", "salesforce")
"""

import logging
import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timezone, timedelta
from typing import Any, AsyncContextManager, Callable, Dict, List, Optional, Type, Union

from sqlalchemy import select, update, delete
from sqlalchemy.ext.asyncio import AsyncSession

from empowernow_common.crypto.fips_aead import encrypt, decrypt
from empowernow_common.vault.user_oauth.base import (
    BaseUserOAuthProvider,
    UserOAuthProviderCapabilities,
)
from empowernow_common.vault.user_oauth.types import (
    UserOAuthCredential,
    CredentialStatus,
)
from empowernow_common.vault.user_oauth.exceptions import (
    CredentialDecryptionError,
    CredentialStorageError,
)

logger = logging.getLogger(__name__)


class PostgresUserOAuthProvider(BaseUserOAuthProvider):
    """
    PostgreSQL storage provider with FIPS 140-3 encryption.

    This provider:
    - Encrypts tokens using AES-256-GCM before storage
    - Supports multi-tenant isolation via tenant_id
    - Tracks refresh failures and can auto-mark needs_reauth
    - Updates last_used_at on credential retrieval

    Attributes:
        PROVIDER_TYPE: "postgres"
        CAPABILITIES: Full feature support including batch fetch
    """

    PROVIDER_TYPE = "postgres"
    CAPABILITIES = UserOAuthProviderCapabilities(
        supports_refresh_tracking=True,
        supports_metadata=True,
        supports_multi_tenant=True,
        supports_list_credentials=True,
        supports_atomic_upsert=True,
        supports_batch_fetch=True,
        persistence_type="persistent",
    )

    def __init__(
        self,
        session_factory: Union[
            Callable[[], AsyncSession],
            Callable[[], AsyncContextManager[AsyncSession]],
        ],
        model_class: Type,
        refresh_failure_threshold: int = 3,
    ):
        """
        Initialize the PostgreSQL provider.

        Args:
            session_factory: Callable that returns an AsyncSession or
                async context manager yielding AsyncSession.
                Supports both patterns:
                - sessionmaker(engine, class_=AsyncSession)
                - async def get_session(): async with session() as s: yield s
            model_class: SQLAlchemy model class from create_third_party_credential_model
            refresh_failure_threshold: Number of refresh failures before
                marking credential as needs_reauth (default: 3)
        """
        self._session_factory = session_factory
        self._model = model_class
        self._refresh_threshold = refresh_failure_threshold

    @asynccontextmanager
    async def _get_session(self) -> AsyncSession:
        """Get a database session, handling both factory patterns."""
        session_or_cm = self._session_factory()

        # Check if it's an async context manager
        if hasattr(session_or_cm, "__aenter__"):
            async with session_or_cm as session:
                yield session
        else:
            # It's a session directly
            try:
                yield session_or_cm
            finally:
                await session_or_cm.close()

    def _build_query(
        self,
        user_arn: str,
        system_key: str,
        tenant_id: Optional[str],
    ):
        """Build base query for credential lookup."""
        query = select(self._model).where(
            self._model.user_arn == user_arn,
            self._model.system_key == system_key,
            self._model.type == "oauth2_user_token",
        )
        if tenant_id:
            query = query.where(self._model.tenant_id == tenant_id)
        else:
            query = query.where(self._model.tenant_id.is_(None))
        return query

    async def get_credential(
        self,
        user_arn: str,
        system_key: str,
        tenant_id: Optional[str] = None,
    ) -> Optional[UserOAuthCredential]:
        """
        Get a user's OAuth credential.

        Retrieves and decrypts the credential, updates last_used_at,
        and returns None if not found or not active.

        Args:
            user_arn: ARN identifying the user
            system_key: External system identifier
            tenant_id: Optional tenant ID

        Returns:
            UserOAuthCredential if found and active, None otherwise

        Raises:
            CredentialDecryptionError: If decryption fails
        """
        async with self._get_session() as session:
            query = self._build_query(user_arn, system_key, tenant_id)
            query = query.where(self._model.status == "active")

            result = await session.execute(query)
            cred = result.scalar_one_or_none()

            if not cred:
                logger.debug(
                    "Credential not found",
                    extra={
                        "user_arn": user_arn,
                        "system_key": system_key,
                        "tenant_id": tenant_id,
                    },
                )
                return None

            # Decrypt token data
            try:
                token_data = decrypt(cred.enc_blob)
            except Exception as e:
                logger.error(
                    "Credential decryption failed",
                    extra={
                        "credential_id": str(cred.id),
                        "error": str(e),
                    },
                )
                raise CredentialDecryptionError(
                    f"Failed to decrypt credential: {e}"
                ) from e

            # Parse expires_at
            expires_at = None
            if token_data.get("expires_at"):
                raw = token_data["expires_at"]
                if isinstance(raw, str):
                    # Handle ISO format with or without Z suffix
                    expires_at = datetime.fromisoformat(
                        raw.replace("Z", "+00:00")
                    )
                elif isinstance(raw, (int, float)):
                    # Unix timestamp
                    expires_at = datetime.fromtimestamp(raw, tz=timezone.utc)

            # Parse scopes
            scope_str = token_data.get("scope", "")
            scopes = scope_str.split() if scope_str else []

            # Update last_used_at (fire and forget, don't fail if it errors)
            try:
                await session.execute(
                    update(self._model)
                    .where(self._model.id == cred.id)
                    .values(last_used_at=datetime.now(timezone.utc))
                )
                await session.commit()
            except Exception as e:
                logger.warning(
                    "Failed to update last_used_at",
                    extra={"credential_id": str(cred.id), "error": str(e)},
                )

            return UserOAuthCredential(
                access_token=token_data["access_token"],
                refresh_token=token_data.get("refresh_token"),
                token_type=token_data.get("token_type", "Bearer"),
                expires_at=expires_at,
                scopes=scopes,
                system_key=system_key,
                user_arn=user_arn,
                metadata=cred.oauth_metadata or {},
                system_user_id=cred.system_user_id,
                system_user_email=cred.system_user_email,
                status=CredentialStatus(cred.status or "active"),
                credential_id=str(cred.id),
                tenant_id=cred.tenant_id,
                last_refreshed_at=cred.last_refreshed_at,
                last_used_at=cred.last_used_at,
                created_at=cred.created_at,
                refresh_failure_count=cred.refresh_failure_count or 0,
            )

    async def store_credential(
        self,
        user_arn: str,
        system_key: str,
        tokens: Dict[str, Any],
        metadata: Optional[Dict[str, Any]] = None,
        tenant_id: Optional[str] = None,
    ) -> str:
        """
        Store or update a user's OAuth credential.

        Encrypts the tokens and performs an atomic upsert.

        Args:
            user_arn: ARN identifying the user
            system_key: External system identifier
            tokens: OAuth token response with access_token, etc.
            metadata: Optional additional metadata
            tenant_id: Optional tenant ID

        Returns:
            Credential ID (UUID string)

        Raises:
            CredentialStorageError: If storage fails
        """
        try:
            # Calculate expiration
            expires_in = tokens.get("expires_in", 3600)
            expires_at = datetime.now(timezone.utc) + timedelta(
                seconds=int(expires_in)
            )

            # Build encrypted payload
            token_payload = {
                "access_token": tokens["access_token"],
                "refresh_token": tokens.get("refresh_token"),
                "token_type": tokens.get("token_type", "Bearer"),
                "expires_at": expires_at.isoformat(),
                "scope": tokens.get("scope", ""),
            }
            enc_blob = encrypt(token_payload)

            # Parse scopes
            scopes = (
                tokens.get("scope", "").split()
                if tokens.get("scope")
                else []
            )

            # Build unique name
            if tenant_id:
                name = f"{tenant_id}:{user_arn}:{system_key}"
            else:
                name = f"{user_arn}:{system_key}"

            # Extract user info from metadata
            system_user_id = None
            system_user_email = None
            if metadata:
                system_user_id = (
                    metadata.get("user_id")
                    or metadata.get("sub")
                    or metadata.get("id")
                )
                system_user_email = metadata.get("email")

            now = datetime.now(timezone.utc)

            async with self._get_session() as session:
                # Check for existing credential
                query = self._build_query(user_arn, system_key, tenant_id)
                result = await session.execute(query)
                existing = result.scalar_one_or_none()

                if existing:
                    # Update existing
                    existing.enc_blob = enc_blob
                    existing.oauth_metadata = metadata
                    existing.credential_scopes = scopes
                    existing.status = "active"
                    existing.last_refreshed_at = now
                    existing.refresh_failure_count = 0
                    existing.updated_at = now
                    if system_user_id:
                        existing.system_user_id = system_user_id
                    if system_user_email:
                        existing.system_user_email = system_user_email
                    credential_id = str(existing.id)
                    logger.info(
                        "Updated credential",
                        extra={
                            "credential_id": credential_id,
                            "system_key": system_key,
                        },
                    )
                else:
                    # Create new
                    new_id = uuid.uuid4()
                    new_cred = self._model(
                        id=new_id,
                        name=name,
                        type="oauth2_user_token",
                        enc_blob=enc_blob,
                        owner_arn=user_arn,
                        user_arn=user_arn,
                        tenant_id=tenant_id,
                        system_key=system_key,
                        oauth_metadata=metadata,
                        credential_scopes=scopes,
                        system_user_id=system_user_id,
                        system_user_email=system_user_email,
                        status="active",
                        last_refreshed_at=now,
                        refresh_failure_count=0,
                        created_at=now,
                        updated_at=now,
                    )
                    session.add(new_cred)
                    credential_id = str(new_id)
                    logger.info(
                        "Created credential",
                        extra={
                            "credential_id": credential_id,
                            "system_key": system_key,
                        },
                    )

                await session.commit()

            return credential_id

        except Exception as e:
            logger.error(
                "Failed to store credential",
                extra={
                    "user_arn": user_arn,
                    "system_key": system_key,
                    "error": str(e),
                },
            )
            raise CredentialStorageError(
                f"Failed to store credential: {e}"
            ) from e

    async def delete_credential(
        self,
        user_arn: str,
        system_key: str,
        tenant_id: Optional[str] = None,
    ) -> bool:
        """
        Delete a user's credential.

        Args:
            user_arn: ARN identifying the user
            system_key: External system identifier
            tenant_id: Optional tenant ID

        Returns:
            True if credential was deleted, False if not found
        """
        async with self._get_session() as session:
            query = delete(self._model).where(
                self._model.user_arn == user_arn,
                self._model.system_key == system_key,
                self._model.type == "oauth2_user_token",
            )
            if tenant_id:
                query = query.where(self._model.tenant_id == tenant_id)
            else:
                query = query.where(self._model.tenant_id.is_(None))

            result = await session.execute(query)
            await session.commit()

            deleted = result.rowcount > 0
            if deleted:
                logger.info(
                    "Deleted credential",
                    extra={
                        "user_arn": user_arn,
                        "system_key": system_key,
                    },
                )
            return deleted

    async def list_credentials(
        self,
        user_arn: str,
        tenant_id: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        List all credentials for a user (metadata only).

        Returns credential metadata without sensitive tokens.

        Args:
            user_arn: ARN identifying the user
            tenant_id: Optional tenant ID

        Returns:
            List of credential metadata dictionaries
        """
        async with self._get_session() as session:
            query = select(self._model).where(
                self._model.user_arn == user_arn,
                self._model.type == "oauth2_user_token",
            )
            if tenant_id:
                query = query.where(self._model.tenant_id == tenant_id)

            result = await session.execute(query)
            creds = result.scalars().all()

            return [
                {
                    "credential_id": str(c.id),
                    "system_key": c.system_key,
                    "status": c.status,
                    "system_user_id": c.system_user_id,
                    "system_user_email": c.system_user_email,
                    "scopes": c.credential_scopes or [],
                    "last_used_at": (
                        c.last_used_at.isoformat() if c.last_used_at else None
                    ),
                    "last_refreshed_at": (
                        c.last_refreshed_at.isoformat()
                        if c.last_refreshed_at
                        else None
                    ),
                    "created_at": (
                        c.created_at.isoformat() if c.created_at else None
                    ),
                    "refresh_failure_count": c.refresh_failure_count or 0,
                }
                for c in creds
            ]

    async def mark_needs_reauth(
        self,
        user_arn: str,
        system_key: str,
        tenant_id: Optional[str] = None,
    ) -> None:
        """
        Mark a credential as needing re-authorization.

        Args:
            user_arn: ARN identifying the user
            system_key: External system identifier
            tenant_id: Optional tenant ID
        """
        async with self._get_session() as session:
            query = (
                update(self._model)
                .where(
                    self._model.user_arn == user_arn,
                    self._model.system_key == system_key,
                    self._model.type == "oauth2_user_token",
                )
                .values(
                    status="needs_reauth",
                    updated_at=datetime.now(timezone.utc),
                )
            )

            if tenant_id:
                query = query.where(self._model.tenant_id == tenant_id)
            else:
                query = query.where(self._model.tenant_id.is_(None))

            await session.execute(query)
            await session.commit()

            logger.info(
                "Marked credential as needs_reauth",
                extra={
                    "user_arn": user_arn,
                    "system_key": system_key,
                },
            )

    async def increment_refresh_failure(
        self,
        user_arn: str,
        system_key: str,
        tenant_id: Optional[str] = None,
    ) -> int:
        """
        Increment the refresh failure count.

        If count exceeds threshold, marks credential as needs_reauth.

        Args:
            user_arn: ARN identifying the user
            system_key: External system identifier
            tenant_id: Optional tenant ID

        Returns:
            New failure count after increment
        """
        async with self._get_session() as session:
            query = self._build_query(user_arn, system_key, tenant_id)
            result = await session.execute(query)
            cred = result.scalar_one_or_none()

            if not cred:
                logger.warning(
                    "Cannot increment failure count: credential not found",
                    extra={
                        "user_arn": user_arn,
                        "system_key": system_key,
                    },
                )
                return 0

            new_count = (cred.refresh_failure_count or 0) + 1
            new_status = cred.status

            # Check threshold
            if new_count >= self._refresh_threshold:
                new_status = "needs_reauth"
                logger.warning(
                    "Refresh failure threshold exceeded, marking needs_reauth",
                    extra={
                        "user_arn": user_arn,
                        "system_key": system_key,
                        "failure_count": new_count,
                        "threshold": self._refresh_threshold,
                    },
                )

            await session.execute(
                update(self._model)
                .where(self._model.id == cred.id)
                .values(
                    refresh_failure_count=new_count,
                    status=new_status,
                    updated_at=datetime.now(timezone.utc),
                )
            )
            await session.commit()

            return new_count

    async def reset_refresh_failures(
        self,
        user_arn: str,
        system_key: str,
        tenant_id: Optional[str] = None,
    ) -> None:
        """
        Reset the refresh failure count after successful refresh.

        Args:
            user_arn: ARN identifying the user
            system_key: External system identifier
            tenant_id: Optional tenant ID
        """
        async with self._get_session() as session:
            query = (
                update(self._model)
                .where(
                    self._model.user_arn == user_arn,
                    self._model.system_key == system_key,
                    self._model.type == "oauth2_user_token",
                )
                .values(
                    refresh_failure_count=0,
                    updated_at=datetime.now(timezone.utc),
                )
            )

            if tenant_id:
                query = query.where(self._model.tenant_id == tenant_id)
            else:
                query = query.where(self._model.tenant_id.is_(None))

            await session.execute(query)
            await session.commit()

    async def get_credentials_batch(
        self,
        user_arn: str,
        system_keys: List[str],
        tenant_id: Optional[str] = None,
    ) -> Dict[str, Optional[UserOAuthCredential]]:
        """
        Fetch multiple credentials in a single database query.

        More efficient than calling get_credential multiple times.

        Args:
            user_arn: ARN identifying the user
            system_keys: List of system identifiers to fetch
            tenant_id: Optional tenant ID

        Returns:
            Dict mapping system_key to credential (or None if not found)
        """
        if not system_keys:
            return {}

        result_dict: Dict[str, Optional[UserOAuthCredential]] = {
            key: None for key in system_keys
        }

        async with self._get_session() as session:
            query = select(self._model).where(
                self._model.user_arn == user_arn,
                self._model.system_key.in_(system_keys),
                self._model.type == "oauth2_user_token",
                self._model.status == "active",
            )

            if tenant_id:
                query = query.where(self._model.tenant_id == tenant_id)
            else:
                query = query.where(self._model.tenant_id.is_(None))

            result = await session.execute(query)
            creds = result.scalars().all()

            for cred in creds:
                try:
                    token_data = decrypt(cred.enc_blob)

                    # Parse expires_at
                    expires_at = None
                    if token_data.get("expires_at"):
                        raw = token_data["expires_at"]
                        if isinstance(raw, str):
                            expires_at = datetime.fromisoformat(
                                raw.replace("Z", "+00:00")
                            )

                    # Parse scopes
                    scope_str = token_data.get("scope", "")
                    scopes = scope_str.split() if scope_str else []

                    result_dict[cred.system_key] = UserOAuthCredential(
                        access_token=token_data["access_token"],
                        refresh_token=token_data.get("refresh_token"),
                        token_type=token_data.get("token_type", "Bearer"),
                        expires_at=expires_at,
                        scopes=scopes,
                        system_key=cred.system_key,
                        user_arn=user_arn,
                        metadata=cred.oauth_metadata or {},
                        system_user_id=cred.system_user_id,
                        system_user_email=cred.system_user_email,
                        status=CredentialStatus(cred.status or "active"),
                        credential_id=str(cred.id),
                        tenant_id=cred.tenant_id,
                        last_refreshed_at=cred.last_refreshed_at,
                        last_used_at=cred.last_used_at,
                        created_at=cred.created_at,
                        refresh_failure_count=cred.refresh_failure_count or 0,
                    )
                except Exception as e:
                    logger.error(
                        "Failed to decrypt credential in batch",
                        extra={
                            "credential_id": str(cred.id),
                            "system_key": cred.system_key,
                            "error": str(e),
                        },
                    )
                    # Leave as None in result

        return result_dict
