# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
SSO Authentication Module - Phase 2

Provides Single Sign-On integration with:
- Google Workspace (most nonprofits)
- Microsoft Entra ID (larger orgs)
- Generic OIDC (Okta, Auth0, etc.)

Authentication Flow:
1. User clicks "Sign in with Google/Microsoft"
2. Redirect to IdP authorization endpoint
3. User authenticates with IdP
4. IdP redirects back with authorization code
5. Exchange code for tokens
6. Validate ID token and extract user info
7. Create/update user in Familiar
8. Create session and redirect to app

Configuration:
    auth:
      method: sso  # or "magic_link"
      sso:
        provider: google  # google, microsoft, oidc
        client_id: ${GOOGLE_CLIENT_ID}
        client_secret: ${GOOGLE_CLIENT_SECRET}
        allowed_domains:
          - nonprofit.org
          - partner.org
        auto_provision: true  # Create users on first login
        default_role: staff   # Role for new users
"""

from __future__ import annotations

import base64
import hashlib
import json  # noqa: F401
import logging
import os
import secrets
from abc import ABC, abstractmethod  # noqa: F401
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from enum import Enum
from typing import Any, Dict, List, Optional  # noqa: F401
from urllib.parse import parse_qs, urlencode, urlparse  # noqa: F401

logger = logging.getLogger(__name__)

# Optional dependencies
try:
    import httpx

    HAS_HTTPX = True
except ImportError:
    HAS_HTTPX = False
    logger.debug("httpx not installed — SSO unavailable.")

try:
    import jwt
    from jwt import PyJWKClient

    HAS_JWT = True
except ImportError:
    HAS_JWT = False
    logger.debug("PyJWT not installed — SSO unavailable.")


class SSOProvider(str, Enum):
    """Supported SSO providers."""

    GOOGLE = "google"
    MICROSOFT = "microsoft"
    OIDC = "oidc"  # Generic OIDC


@dataclass
class SSOConfig:
    """SSO configuration."""

    provider: SSOProvider
    client_id: str
    client_secret: str
    redirect_uri: str  # e.g., "https://familiar.nonprofit.org/auth/callback"

    # Optional settings
    allowed_domains: List[str] = field(default_factory=list)  # Restrict to org domains
    auto_provision: bool = True  # Create users on first login
    default_role: str = "staff"  # Role for auto-provisioned users

    # OIDC endpoints (auto-discovered for Google/Microsoft)
    issuer: str = None
    authorization_endpoint: str = None
    token_endpoint: str = None
    userinfo_endpoint: str = None
    jwks_uri: str = None

    # Scopes
    scopes: List[str] = field(default_factory=lambda: ["openid", "email", "profile"])

    @classmethod
    def from_dict(cls, data: Dict) -> "SSOConfig":
        """Create from config dict."""
        provider = SSOProvider(data.get("provider", "google"))

        return cls(
            provider=provider,
            client_id=data.get(
                "client_id", os.environ.get(f"{provider.value.upper()}_CLIENT_ID", "")
            ),
            client_secret=data.get(
                "client_secret", os.environ.get(f"{provider.value.upper()}_CLIENT_SECRET", "")
            ),
            redirect_uri=data.get("redirect_uri", "http://localhost:5001/auth/callback"),
            allowed_domains=data.get("allowed_domains", []),
            auto_provision=data.get("auto_provision", True),
            default_role=data.get("default_role", "staff"),
            issuer=data.get("issuer"),
            authorization_endpoint=data.get("authorization_endpoint"),
            token_endpoint=data.get("token_endpoint"),
            userinfo_endpoint=data.get("userinfo_endpoint"),
            jwks_uri=data.get("jwks_uri"),
            scopes=data.get("scopes", ["openid", "email", "profile"]),
        )


@dataclass
class SSOUser:
    """User info from SSO provider."""

    sub: str  # Unique ID from IdP
    email: str
    name: str
    given_name: str = None
    family_name: str = None
    picture: str = None
    email_verified: bool = True
    provider: SSOProvider = None
    raw_claims: Dict = field(default_factory=dict)

    @property
    def domain(self) -> str:
        """Get email domain."""
        return self.email.split("@")[-1].lower()


@dataclass
class SSOTokens:
    """OAuth tokens from SSO provider."""

    access_token: str
    id_token: str
    refresh_token: str = None
    expires_in: int = 3600
    token_type: str = "Bearer"
    scope: str = None


# ============================================================
# WELL-KNOWN OIDC CONFIGURATIONS
# ============================================================

GOOGLE_OIDC = {
    "issuer": "https://accounts.google.com",
    "authorization_endpoint": "https://accounts.google.com/o/oauth2/v2/auth",
    "token_endpoint": "https://oauth2.googleapis.com/token",
    "userinfo_endpoint": "https://openidconnect.googleapis.com/v1/userinfo",
    "jwks_uri": "https://www.googleapis.com/oauth2/v3/certs",
}

MICROSOFT_OIDC = {
    "issuer": "https://login.microsoftonline.com/common/v2.0",
    "authorization_endpoint": "https://login.microsoftonline.com/common/oauth2/v2.0/authorize",
    "token_endpoint": "https://login.microsoftonline.com/common/oauth2/v2.0/token",
    "userinfo_endpoint": "https://graph.microsoft.com/oidc/userinfo",
    "jwks_uri": "https://login.microsoftonline.com/common/discovery/v2.0/keys",
}


# ============================================================
# SSO CLIENT
# ============================================================


class SSOClient:
    """
    SSO client for OAuth 2.0 / OIDC authentication.

    Usage:
        config = SSOConfig(
            provider=SSOProvider.GOOGLE,
            client_id="xxx",
            client_secret="xxx",
            redirect_uri="https://app.example.org/auth/callback"
        )

        client = SSOClient(config)

        # Step 1: Generate auth URL
        auth_url, state = client.get_authorization_url()
        # Redirect user to auth_url

        # Step 2: Handle callback
        tokens = await client.exchange_code(code)
        user = await client.get_user_info(tokens)

        # Step 3: Create/update Familiar user
        familiar_user = client.provision_user(user)
    """

    def __init__(self, config: SSOConfig):
        self.config = config
        self._setup_endpoints()

        # State tokens for CSRF protection.
        # NOTE: These are stored in-memory, so they do not survive process
        # restarts and are not shared across workers.  For multi-process
        # deployments, replace with a shared store (Redis, DB, etc.).
        self._pending_states: Dict[str, datetime] = {}

    def _setup_endpoints(self):
        """Set up OIDC endpoints based on provider."""
        if self.config.provider == SSOProvider.GOOGLE:
            endpoints = GOOGLE_OIDC
        elif self.config.provider == SSOProvider.MICROSOFT:
            endpoints = MICROSOFT_OIDC
        else:
            # Custom OIDC - use config values
            endpoints = {}

        # Apply defaults if not set in config
        self.issuer = self.config.issuer or endpoints.get("issuer")
        self.authorization_endpoint = self.config.authorization_endpoint or endpoints.get(
            "authorization_endpoint"
        )
        self.token_endpoint = self.config.token_endpoint or endpoints.get("token_endpoint")
        self.userinfo_endpoint = self.config.userinfo_endpoint or endpoints.get("userinfo_endpoint")
        self.jwks_uri = self.config.jwks_uri or endpoints.get("jwks_uri")

    def _cleanup_expired_states(self):
        """Remove expired SSO state tokens to prevent memory leaks."""
        now = datetime.now(timezone.utc)
        expired = [
            key
            for key, expiry in self._pending_states.items()
            if isinstance(expiry, datetime) and expiry < now
        ]
        for key in expired:
            del self._pending_states[key]
            # Also remove the companion PKCE verifier if present
            verifier_key = f"{key}_verifier"
            self._pending_states.pop(verifier_key, None)

    def get_authorization_url(self, extra_params: Dict = None) -> tuple[str, str]:
        """
        Generate the authorization URL to redirect user to.

        Returns:
            (auth_url, state) - URL to redirect to and state token to verify
        """
        self._cleanup_expired_states()

        # Generate state for CSRF protection
        state = secrets.token_urlsafe(32)
        self._pending_states[state] = datetime.now(timezone.utc) + timedelta(minutes=10)

        # Generate PKCE code verifier and challenge
        code_verifier = secrets.token_urlsafe(64)
        code_challenge = (
            base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest())
            .decode()
            .rstrip("=")
        )

        # Store code verifier with state
        self._pending_states[f"{state}_verifier"] = code_verifier

        params = {
            "client_id": self.config.client_id,
            "redirect_uri": self.config.redirect_uri,
            "response_type": "code",
            "scope": " ".join(self.config.scopes),
            "state": state,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
            "access_type": "offline",  # Get refresh token
            "prompt": "select_account",  # Always show account picker
        }

        # Add extra params
        if extra_params:
            params.update(extra_params)

        # Microsoft-specific
        if self.config.provider == SSOProvider.MICROSOFT:
            params["response_mode"] = "query"

        auth_url = f"{self.authorization_endpoint}?{urlencode(params)}"

        return auth_url, state

    def verify_state(self, state: str) -> bool:
        """Verify state token from callback."""
        self._cleanup_expired_states()

        if state not in self._pending_states:
            return False

        expires = self._pending_states[state]
        if datetime.now(timezone.utc) > expires:
            del self._pending_states[state]
            return False

        return True

    async def exchange_code(self, code: str, state: str) -> SSOTokens:
        """
        Exchange authorization code for tokens.

        Args:
            code: Authorization code from callback
            state: State token for verification

        Returns:
            SSOTokens with access_token, id_token, etc.
        """
        if not HAS_HTTPX:
            raise RuntimeError("httpx required for SSO.")

        # Get code verifier
        verifier_key = f"{state}_verifier"
        code_verifier = self._pending_states.get(verifier_key)

        # Clean up state
        if state in self._pending_states:
            del self._pending_states[state]
        if verifier_key in self._pending_states:
            del self._pending_states[verifier_key]

        data = {
            "client_id": self.config.client_id,
            "client_secret": self.config.client_secret,
            "code": code,
            "redirect_uri": self.config.redirect_uri,
            "grant_type": "authorization_code",
        }

        if code_verifier:
            data["code_verifier"] = code_verifier

        async with httpx.AsyncClient() as client:
            response = await client.post(
                self.token_endpoint, data=data, headers={"Accept": "application/json"}
            )

            if response.status_code != 200:
                logger.error(f"Token exchange failed: {response.text}")
                raise ValueError(f"Token exchange failed: {response.status_code}")

            token_data = response.json()

        return SSOTokens(
            access_token=token_data["access_token"],
            id_token=token_data.get("id_token", ""),
            refresh_token=token_data.get("refresh_token"),
            expires_in=token_data.get("expires_in", 3600),
            token_type=token_data.get("token_type", "Bearer"),
            scope=token_data.get("scope"),
        )

    async def refresh_access_token(self, refresh_token: str) -> SSOTokens:
        """Refresh an expired access token."""
        if not HAS_HTTPX:
            raise RuntimeError("httpx required for SSO")

        data = {
            "client_id": self.config.client_id,
            "client_secret": self.config.client_secret,
            "refresh_token": refresh_token,
            "grant_type": "refresh_token",
        }

        async with httpx.AsyncClient() as client:
            response = await client.post(
                self.token_endpoint, data=data, headers={"Accept": "application/json"}
            )

            if response.status_code != 200:
                raise ValueError(f"Token refresh failed: {response.status_code}")

            token_data = response.json()

        return SSOTokens(
            access_token=token_data["access_token"],
            id_token=token_data.get("id_token", ""),
            refresh_token=token_data.get("refresh_token", refresh_token),
            expires_in=token_data.get("expires_in", 3600),
        )

    async def get_user_info(self, tokens: SSOTokens) -> SSOUser:
        """
        Get user info from ID token or userinfo endpoint.

        Args:
            tokens: OAuth tokens from exchange_code()

        Returns:
            SSOUser with email, name, etc.
        """
        # First try to decode ID token (faster, no network)
        if tokens.id_token and HAS_JWT:
            try:
                return self._decode_id_token(tokens.id_token)
            except Exception as e:
                logger.warning(f"ID token decode failed, falling back to userinfo: {e}")

        # Fall back to userinfo endpoint
        return await self._fetch_userinfo(tokens.access_token)

    def _decode_id_token(self, id_token: str) -> SSOUser:
        """Decode and validate ID token."""
        if not HAS_JWT:
            raise RuntimeError("PyJWT required for SSO token validation.")

        # Get signing keys
        jwks_client = PyJWKClient(self.jwks_uri)
        signing_key = jwks_client.get_signing_key_from_jwt(id_token)

        # Decode and verify
        claims = jwt.decode(
            id_token,
            signing_key.key,
            algorithms=["RS256"],
            audience=self.config.client_id,
            issuer=self.issuer,
        )

        return SSOUser(
            sub=claims["sub"],
            email=claims.get("email", ""),
            name=claims.get("name", claims.get("email", "").split("@")[0]),
            given_name=claims.get("given_name"),
            family_name=claims.get("family_name"),
            picture=claims.get("picture"),
            email_verified=claims.get("email_verified", False),
            provider=self.config.provider,
            raw_claims=claims,
        )

    async def _fetch_userinfo(self, access_token: str) -> SSOUser:
        """Fetch user info from userinfo endpoint."""
        if not HAS_HTTPX:
            raise RuntimeError("httpx required for SSO")

        async with httpx.AsyncClient() as client:
            response = await client.get(
                self.userinfo_endpoint,
                headers={
                    "Authorization": f"Bearer {access_token}",
                    "Accept": "application/json",
                },
            )

            if response.status_code != 200:
                raise ValueError(f"Userinfo fetch failed: {response.status_code}")

            info = response.json()

        return SSOUser(
            sub=info["sub"],
            email=info.get("email", ""),
            name=info.get("name", info.get("email", "").split("@")[0]),
            given_name=info.get("given_name"),
            family_name=info.get("family_name"),
            picture=info.get("picture"),
            email_verified=info.get("email_verified", False),
            provider=self.config.provider,
            raw_claims=info,
        )

    def validate_domain(self, user: SSOUser) -> bool:
        """Check if user's email domain is allowed."""
        if not self.config.allowed_domains:
            return True  # No restrictions

        return user.domain in [d.lower() for d in self.config.allowed_domains]

    def provision_user(self, sso_user: SSOUser, user_manager=None):
        """
        Create or update Familiar user from SSO user.

        Args:
            sso_user: User info from IdP
            user_manager: UserManager instance (optional, will get default)

        Returns:
            Familiar User object
        """
        # Validate domain
        if not self.validate_domain(sso_user):
            raise ValueError(f"Email domain not allowed: {sso_user.domain}")

        # Get user manager
        if user_manager is None:
            from ..users import get_user_manager

            user_manager = get_user_manager()

        # Check if user exists
        from ..users import UserRole, UserStatus

        existing = user_manager.get_user_by_email(sso_user.email)

        if existing:
            # Update SSO info
            user_manager.update_user(
                existing.id,
                name=sso_user.name,
                sso_provider=self.config.provider.value,
                sso_subject=sso_user.sub,
            )
            user_manager.update_last_active(existing.id)
            return user_manager.get_user(existing.id)

        # Auto-provision new user
        if not self.config.auto_provision:
            raise ValueError(f"User not found and auto-provision disabled: {sso_user.email}")

        # Map role
        role = UserRole(self.config.default_role)

        # Create user
        user = user_manager.create_user(
            email=sso_user.email,
            name=sso_user.name,
            role=role,
            status=UserStatus.ACTIVE,
        )

        # Set SSO info
        user_manager.update_user(
            user.id,
            sso_provider=self.config.provider.value,
            sso_subject=sso_user.sub,
        )

        logger.info(f"Auto-provisioned SSO user: {sso_user.email} ({role.value})")

        return user


# ============================================================
# FLASK INTEGRATION
# ============================================================


def setup_sso_routes(app, sso_client: SSOClient, user_manager=None):
    """
    Add SSO routes to Flask app.

    Routes:
        GET /auth/login - Redirect to IdP
        GET /auth/callback - Handle IdP callback
        GET /auth/logout - Clear session

    Usage:
        from flask import Flask
        from familiar.core.auth.sso import SSOClient, SSOConfig, setup_sso_routes

        app = Flask(__name__)
        config = SSOConfig.from_dict(app.config["SSO"])
        client = SSOClient(config)
        setup_sso_routes(app, client)
    """
    import asyncio

    from flask import flash, redirect, request, session, url_for

    @app.route("/auth/login")
    def sso_login():
        """Redirect to SSO provider."""
        auth_url, state = sso_client.get_authorization_url()
        session["sso_state"] = state
        return redirect(auth_url)

    @app.route("/auth/callback")
    def sso_callback():
        """Handle SSO callback."""
        # Check for errors
        error = request.args.get("error")
        if error:
            flash(f"SSO error: {error}", "error")
            return redirect(url_for("login"))

        # Verify state
        code = request.args.get("code")
        state = request.args.get("state")

        if state != session.pop("sso_state", None) or not sso_client.verify_state(state):
            flash("Invalid or expired login session", "error")
            return redirect(url_for("login"))

        # Exchange code for tokens
        try:
            tokens = asyncio.run(sso_client.exchange_code(code, state))
            sso_user = asyncio.run(sso_client.get_user_info(tokens))
        except Exception as e:
            logger.error(f"SSO callback failed: {e}")
            flash(f"Login failed: {str(e)}", "error")
            return redirect(url_for("login"))

        # Validate domain
        if not sso_client.validate_domain(sso_user):
            flash(f"Email domain not allowed: {sso_user.domain}", "error")
            return redirect(url_for("login"))

        # Provision/update user
        try:
            user = sso_client.provision_user(sso_user, user_manager)
        except Exception as e:
            logger.error(f"User provisioning failed: {e}")
            flash(f"Account setup failed: {str(e)}", "error")
            return redirect(url_for("login"))

        # Session regeneration: clear old session data to prevent fixation
        session.clear()

        # Create session
        if user_manager:
            token = user_manager.create_session(user)
            session["session_token"] = token
        else:
            session["user_id"] = user.id

        session["user_email"] = user.email
        session["user_name"] = user.name

        flash(f"Welcome, {user.name}!", "success")
        return redirect(url_for("dashboard"))

    @app.route("/auth/logout")
    def sso_logout():
        """Clear session."""
        session.clear()
        flash("You have been logged out", "info")
        return redirect(url_for("login"))

    return app


# ============================================================
# HELPER FUNCTIONS
# ============================================================


def create_sso_client(config_dict: Dict = None) -> SSOClient:
    """
    Create SSO client from config dict or environment.

    Args:
        config_dict: Configuration dict, or None to use env vars

    Returns:
        Configured SSOClient
    """
    if config_dict is None:
        # Build from environment
        provider = os.environ.get("SSO_PROVIDER", "google")
        config_dict = {
            "provider": provider,
            "client_id": os.environ.get(f"{provider.upper()}_CLIENT_ID"),
            "client_secret": os.environ.get(f"{provider.upper()}_CLIENT_SECRET"),
            "redirect_uri": os.environ.get(
                "SSO_REDIRECT_URI", "http://localhost:5001/auth/callback"
            ),
            "allowed_domains": [
                d.strip() for d in os.environ.get("SSO_ALLOWED_DOMAINS", "").split(",") if d.strip()
            ],
        }

    config = SSOConfig.from_dict(config_dict)
    return SSOClient(config)


def get_sso_provider_name(provider: SSOProvider) -> str:
    """Get display name for SSO provider."""
    return {
        SSOProvider.GOOGLE: "Google",
        SSOProvider.MICROSOFT: "Microsoft",
        SSOProvider.OIDC: "SSO",
    }.get(provider, "SSO")
