# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
#

"""
Simplified Okta OAuth 2.0 PKCE Authentication for CloudZero Operations CLI.

This module provides a streamlined implementation of the OAuth 2.0 PKCE flow
for authenticating with Okta.
"""

import base64
import hashlib
import platform
import secrets
import threading
import webbrowser
from datetime import UTC, datetime
from urllib.parse import urlencode, urlparse

import httpx
from flask import Flask, request
from rich.console import Console
from werkzeug.serving import make_server

from src.settings import Settings

from .token_storage import TokenStorage, TokenStorageError, get_token_storage
from .utils import (
    add_token_metadata,
    add_userinfo_metadata,
    extract_access_token,
    extract_refresh_token,
    is_token_valid,
    is_userinfo_valid,
)

# Flask app for OAuth callback
app = Flask(__name__)
server_token = {}
shutdown_event = threading.Event()

# Rich console for formatted output
console = Console()


@app.route("/callback")
def callback():
    """Handle OAuth callback from Okta on port 8042."""
    code = request.args.get("code")
    error = request.args.get("error")
    error_description = request.args.get("error_description", "No description provided")

    if error:
        server_token["error"] = f"{error}: {error_description}"
        shutdown_event.set()
        return (
            f"""
        <html>
        <head><title>Authentication Failed</title></head>
        <body style="font-family: Arial, sans-serif; text-align: center; padding: 50px;">
            <h2>❌ Authentication Failed</h2>
            <p>Error: {error}</p>
            <p>Description: {error_description}</p>
            <p>Please close this window and check the terminal for more information.</p>
        </body>
        </html>
        """,
            400,
        )

    if not code:
        server_token["error"] = "No authorization code provided"
        shutdown_event.set()
        return "No authorization code provided", 400

    server_token["code"] = code
    shutdown_event.set()
    return """
    <!DOCTYPE html>
    <html>
    <head>
        <title>Authentication Successful</title>
        <style>
            body {
                margin: 0;
                padding: 0;
                background: linear-gradient(to bottom,
                    #87CEEB 0%,
                    #B0C4DE 20%,
                    #FFB6C1 40%,
                    #FFA07A 60%,
                    #FF8C69 80%,
                    #FF7F50 100%);
                font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Helvetica Neue', Arial, sans-serif;
                color: #ffffff;
                overflow: hidden;
                display: flex;
                flex-direction: column;
                justify-content: center;
                align-items: center;
                height: 100vh;
            }
            #canvas {
                position: fixed;
                top: 0;
                left: 0;
                width: 100%;
                height: 100%;
                z-index: 1;
            }
            .content {
                position: relative;
                z-index: 2;
                text-align: center;
                width: 100%;
                max-width: 600px;
                padding: 0 40px;
            }
            .message {
                font-size: 1.5em;
                font-weight: 300;
                letter-spacing: 0.5px;
                margin-bottom: 30px;
                color: #ffffff;
                text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.3);
            }
            .divider {
                width: 100%;
                height: 1px;
                background: linear-gradient(90deg,
                    transparent 0%,
                    rgba(255, 255, 255, 0.3) 20%,
                    rgba(255, 255, 255, 0.8) 50%,
                    rgba(255, 255, 255, 0.3) 80%,
                    transparent 100%);
                margin: 0 auto;
            }
            .countdown {
                margin-top: 30px;
                font-size: 0.95em;
                font-weight: 300;
                color: rgba(255, 255, 255, 0.9);
                text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.3);
            }
            .pause-control {
                position: fixed;
                bottom: 20px;
                right: 20px;
                z-index: 3;
                display: flex;
                align-items: center;
                gap: 8px;
                background: rgba(255, 255, 255, 0.15);
                padding: 10px 15px;
                border-radius: 8px;
                backdrop-filter: blur(10px);
                border: 1px solid rgba(255, 255, 255, 0.3);
                cursor: pointer;
                user-select: none;
            }
            .pause-control:hover {
                background: rgba(255, 255, 255, 0.25);
            }
            .pause-control input[type="checkbox"] {
                width: 18px;
                height: 18px;
                cursor: pointer;
                accent-color: rgba(255, 255, 255, 0.8);
            }
            .pause-control label {
                font-size: 0.9em;
                font-weight: 400;
                color: rgba(255, 255, 255, 0.95);
                text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.3);
                cursor: pointer;
            }
        </style>
    </head>
    <body>
        <canvas id="canvas"></canvas>
        <div class="content">
            <div class="message">CZOPS Authentication Successful</div>
            <div class="divider"></div>
            <div class="countdown">Window closes in <span id="timer">20</span> seconds</div>
        </div>
        <div class="pause-control" onclick="document.getElementById('pauseCheckbox').click()">
            <input type="checkbox" id="pauseCheckbox" onclick="event.stopPropagation()">
            <label for="pauseCheckbox">Pause timer</label>
        </div>
        <script>
            // Murmuration of Mouse Pointer Sparrows
            const canvas = document.getElementById('canvas');
            const ctx = canvas.getContext('2d');

            canvas.width = window.innerWidth;
            canvas.height = window.innerHeight;

            const particles = [];
            const particleCount = 900;
            const time = { value: 0 };

            // Mouse tracking
            const mouse = {
                x: canvas.width / 2,
                y: canvas.height / 2,
                isMoving: false,
                lastMoveTime: 0
            };

            // Track mouse movement
            canvas.addEventListener('mousemove', (e) => {
                mouse.x = e.clientX;
                mouse.y = e.clientY;
                mouse.isMoving = true;
                mouse.lastMoveTime = Date.now();
            });

            // Check if mouse has stopped moving
            setInterval(() => {
                if (Date.now() - mouse.lastMoveTime > 300) {
                    mouse.isMoving = false;
                }
            }, 100);

            // Add birds when mouse button is held down
            let mousePressed = false;
            let spawnInterval = null;

            canvas.addEventListener('mousedown', (e) => {
                mousePressed = true;
                mouse.x = e.clientX;
                mouse.y = e.clientY;

                // Spawn birds continuously while mouse is held down
                spawnInterval = setInterval(() => {
                    if (mousePressed) {
                        // Add 5 birds at a time around the mouse position
                        for (let i = 0; i < 5; i++) {
                            const newBird = new Sparrow(particles.length);
                            // Spawn near mouse with slight random offset
                            const angle = Math.random() * Math.PI * 2;
                            const radius = Math.random() * 30;
                            newBird.x = mouse.x + Math.cos(angle) * radius;
                            newBird.y = mouse.y + Math.sin(angle) * radius;
                            // Give them initial velocity away from spawn point
                            newBird.vx = Math.cos(angle) * 2;
                            newBird.vy = Math.sin(angle) * 2;
                            particles.push(newBird);
                        }
                    }
                }, 100); // Spawn every 100ms
            });

            canvas.addEventListener('mouseup', () => {
                mousePressed = false;
                if (spawnInterval) {
                    clearInterval(spawnInterval);
                    spawnInterval = null;
                }
            });

            canvas.addEventListener('mouseleave', () => {
                mousePressed = false;
                if (spawnInterval) {
                    clearInterval(spawnInterval);
                    spawnInterval = null;
                }
            });

            class Sparrow {
                constructor(index) {
                    this.x = Math.random() * canvas.width;
                    this.y = Math.random() * canvas.height;
                    this.vx = (Math.random() - 0.5) * 2;
                    this.vy = (Math.random() - 0.5) * 2;
                    this.maxSpeed = 3;
                    this.maxForce = 0.05;
                }

                update() {
                    const perceptionRadius = 100;
                    const separationRadius = 30;
                    let alignmentX = 0, alignmentY = 0;
                    let cohesionX = 0, cohesionY = 0;
                    let separationX = 0, separationY = 0;
                    let alignCount = 0, cohesionCount = 0, separationCount = 0;

                    particles.forEach(other => {
                        if (other === this) return;

                        const dx = other.x - this.x;
                        const dy = other.y - this.y;
                        const dist = Math.sqrt(dx * dx + dy * dy);

                        if (dist > 0 && dist < perceptionRadius) {
                            // Alignment
                            alignmentX += other.vx;
                            alignmentY += other.vy;
                            alignCount++;

                            // Cohesion
                            cohesionX += other.x;
                            cohesionY += other.y;
                            cohesionCount++;
                        }

                        if (dist > 0 && dist < separationRadius) {
                            // Separation
                            separationX += (this.x - other.x) / dist;
                            separationY += (this.y - other.y) / dist;
                            separationCount++;
                        }
                    });

                    // Apply alignment
                    if (alignCount > 0) {
                        alignmentX /= alignCount;
                        alignmentY /= alignCount;
                        const mag = Math.sqrt(alignmentX * alignmentX + alignmentY * alignmentY);
                        if (mag > 0) {
                            alignmentX = (alignmentX / mag) * this.maxSpeed - this.vx;
                            alignmentY = (alignmentY / mag) * this.maxSpeed - this.vy;
                            const alignMag = Math.sqrt(alignmentX * alignmentX + alignmentY * alignmentY);
                            if (alignMag > this.maxForce) {
                                alignmentX = (alignmentX / alignMag) * this.maxForce;
                                alignmentY = (alignmentY / alignMag) * this.maxForce;
                            }
                            this.vx += alignmentX;
                            this.vy += alignmentY;
                        }
                    }

                    // Apply cohesion
                    if (cohesionCount > 0) {
                        cohesionX /= cohesionCount;
                        cohesionY /= cohesionCount;
                        cohesionX -= this.x;
                        cohesionY -= this.y;
                        const mag = Math.sqrt(cohesionX * cohesionX + cohesionY * cohesionY);
                        if (mag > 0) {
                            cohesionX = (cohesionX / mag) * this.maxSpeed - this.vx;
                            cohesionY = (cohesionY / mag) * this.maxSpeed - this.vy;
                            const cohMag = Math.sqrt(cohesionX * cohesionX + cohesionY * cohesionY);
                            if (cohMag > this.maxForce) {
                                cohesionX = (cohesionX / cohMag) * this.maxForce;
                                cohesionY = (cohesionY / cohMag) * this.maxForce;
                            }
                            this.vx += cohesionX;
                            this.vy += cohesionY;
                        }
                    }

                    // Apply separation
                    if (separationCount > 0) {
                        separationX /= separationCount;
                        separationY /= separationCount;
                        const mag = Math.sqrt(separationX * separationX + separationY * separationY);
                        if (mag > 0) {
                            separationX = (separationX / mag) * this.maxSpeed - this.vx;
                            separationY = (separationY / mag) * this.maxSpeed - this.vy;
                            const sepMag = Math.sqrt(separationX * separationX + separationY * separationY);
                            if (sepMag > this.maxForce) {
                                separationX = (separationX / sepMag) * this.maxForce;
                                separationY = (separationY / sepMag) * this.maxForce;
                            }
                            this.vx += separationX * 1.5;
                            this.vy += separationY * 1.5;
                        }
                    }

                    // Follow mouse if it's moving
                    if (mouse.isMoving) {
                        const dxToMouse = mouse.x - this.x;
                        const dyToMouse = mouse.y - this.y;
                        const distToMouse = Math.sqrt(dxToMouse * dxToMouse + dyToMouse * dyToMouse);

                        if (distToMouse > 0) {
                            // Normalize and apply force toward mouse
                            const mouseForceX = (dxToMouse / distToMouse) * this.maxForce * 2;
                            const mouseForceY = (dyToMouse / distToMouse) * this.maxForce * 2;
                            this.vx += mouseForceX;
                            this.vy += mouseForceY;
                        }
                    }

                    // Limit speed
                    const speed = Math.sqrt(this.vx * this.vx + this.vy * this.vy);
                    if (speed > this.maxSpeed) {
                        this.vx = (this.vx / speed) * this.maxSpeed;
                        this.vy = (this.vy / speed) * this.maxSpeed;
                    }

                    // Update position
                    this.x += this.vx;
                    this.y += this.vy;

                    // Wrap around edges
                    if (this.x < 0) this.x = canvas.width;
                    if (this.x > canvas.width) this.x = 0;
                    if (this.y < 0) this.y = canvas.height;
                    if (this.y > canvas.height) this.y = 0;
                }

                draw() {
                    // Calculate angle based on velocity
                    const angle = Math.atan2(this.vy, this.vx);

                    ctx.save();
                    ctx.translate(this.x, this.y);
                    ctx.rotate(angle);

                    // Draw black arrow/bird shape
                    ctx.beginPath();
                    ctx.moveTo(6, 0);           // Point forward
                    ctx.lineTo(-3, -3);         // Top wing
                    ctx.lineTo(-2, 0);          // Body
                    ctx.lineTo(-3, 3);          // Bottom wing
                    ctx.closePath();

                    ctx.fillStyle = 'rgba(0, 0, 0, 0.8)';
                    ctx.fill();

                    ctx.restore();
                }
            }

            // Initialize sparrows
            for (let i = 0; i < particleCount; i++) {
                particles.push(new Sparrow(i));
            }

            // Animation loop
            function animate() {
                time.value++;

                // Clear canvas (no trail effect - just clear background)
                ctx.clearRect(0, 0, canvas.width, canvas.height);

                particles.forEach(particle => {
                    particle.update();
                    particle.draw();
                });

                requestAnimationFrame(animate);
            }

            animate();

            // Countdown timer with pause functionality
            let seconds = 20;
            const timerElement = document.getElementById('timer');
            const pauseCheckbox = document.getElementById('pauseCheckbox');
            let countdownInterval = null;
            let closeTimeout = null;

            function startCountdown() {
                countdownInterval = setInterval(() => {
                    if (!pauseCheckbox.checked) {
                        seconds--;
                        timerElement.textContent = seconds;
                        if (seconds <= 0) {
                            clearInterval(countdownInterval);
                            window.close();
                        }
                    }
                }, 1000);
            }

            // Start countdown
            startCountdown();

            // Update timer text when checkbox changes
            pauseCheckbox.addEventListener('change', (e) => {
                if (e.target.checked) {
                    timerElement.textContent = seconds + ' (paused)';
                } else {
                    timerElement.textContent = seconds;
                }
            });

            // Handle window resize
            window.addEventListener('resize', () => {
                canvas.width = window.innerWidth;
                canvas.height = window.innerHeight;
            });
        </script>
    </body>
    </html>
    """


class OktaPKCEAuthenticator:
    """Simplified Okta PKCE authenticator with secure file-based storage."""

    def __init__(self, token_storage: TokenStorage | None = None):
        """Initialize the authenticator with settings and token storage.

        Args:
            token_storage: Optional custom token storage instance for testing
        """
        self.settings = Settings()
        self.client_id = self.settings.api_v4.auth["CLIENT_ID"]
        self.okta_domain = self.settings.api_v4.auth["OKTA_DOMAIN"]
        self.redirect_uri = self.settings.api_v4.auth["LOGIN_REDIRECT_URI"]

        # Okta endpoints - using default API authorization server
        self.authorization_endpoint = f"https://{self.okta_domain}/oauth2/default/v1/authorize"
        self.token_endpoint = f"https://{self.okta_domain}/oauth2/default/v1/token"
        self.userinfo_endpoint = f"https://{self.okta_domain}/oauth2/default/v1/userinfo"

        self.scopes = ["openid", "profile", "email", "offline_access", "groups"]

        # Initialize token storage (allows dependency injection for testing)
        self.token_storage = token_storage or get_token_storage()

    def authenticate(self, force=False, debug=False):
        """
        Perform PKCE authentication flow.

        Args:
            force: Force re-authentication even if valid token exists
            debug: Show debug information during authentication

        Returns:
            dict: Token response containing access_token, refresh_token, etc.
        """
        # Check for existing valid token
        if not force:
            existing_token = self.get_stored_token()
            if existing_token and self.is_token_valid(existing_token):
                return existing_token

        # Generate PKCE code verifier and challenge
        code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("utf-8").rstrip("=")
        code_challenge = (
            base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode("utf-8")).digest()).decode("utf-8").rstrip("=")
        )

        # Generate state for CSRF protection
        state = secrets.token_urlsafe(32)

        # Build authorization URL
        auth_params = {
            "client_id": self.client_id,
            "response_type": "code",
            "scope": " ".join(self.scopes),
            "redirect_uri": self.redirect_uri,
            "state": state,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
            # Request group claims in both userinfo and id_token
            "claims": '{"userinfo":{"groups":{"essential":true}},"id_token":{"groups":{"essential":true}}}',
        }
        uri = f"{self.authorization_endpoint}?{urlencode(auth_params)}"

        # Store the code verifier and state for later use
        server_token["verifier"] = code_verifier
        server_token["state"] = state

        if debug:
            print("\n🔐 Opening browser for authentication...")
            print(f"If the browser doesn't open automatically, visit:\n{uri}\n")

        # Show debug information for authorization parameters
        if debug:
            print("Authorization Request Parameters:")
            for key, value in auth_params.items():
                if key == "claims":
                    print(f"  {key}: {value}")
                else:
                    print(f"  {key}: {value[:50]}..." if len(str(value)) > 50 else f"  {key}: {value}")

        # Show debug information only when requested
        if debug:
            print("Debug Information:")
            print(f"  Client ID: {self.client_id}")
            print(f"  Redirect URI: {self.redirect_uri}")
            print(f"  Okta Domain: {self.okta_domain}")
            print(f"  Authorization Endpoint: {self.authorization_endpoint}")

        # Check for configuration mismatch (always show warnings)
        if self.redirect_uri and "8080" in self.redirect_uri:
            print("\n⚠️  WARNING: Your settings are using port 8080, but Okta app expects port 8042!")
            print("   Please run: czops configure --auth")
            print("   And update LOGIN_REDIRECT_URI to: http://127.0.0.1:8042/callback")

        print()

        # Try to open browser with enhanced error handling
        print("🔐 Opening browser for authentication...")
        browser_opened = webbrowser.open(uri)

        # Always display the URL for manual access (helps with Windows/WSL issues)
        console.print("\n[bold cyan]Authentication URL:[/bold cyan]")
        console.print(f"[link={uri}]{uri}[/link]")

        if browser_opened:
            console.print("\n[green]✓[/green] Browser opened automatically")
        else:
            console.print("\n[yellow]⚠ Browser did not open automatically[/yellow]")
            console.print("[yellow]Please copy the URL above and paste it into your browser[/yellow]")

            if platform.system() == "Windows":
                console.print("\n[dim]Windows tip: Right-click to copy the URL, or try running as Administrator[/dim]")

        console.print("\n[dim]Waiting for you to complete authentication in your browser...[/dim]")

        # Start local server to capture callback
        authorization_code = self._run_callback_server()

        if not authorization_code:
            raise AuthenticationError("Failed to get authorization code")

        # Exchange code for tokens
        token = self._exchange_code_for_token(authorization_code, server_token["verifier"])

        if debug:
            print("Received Token:")
            for k, v in token.items():
                if k in ["access_token", "refresh_token", "id_token"]:
                    print(f"  {k}: {'*' * 8} (hidden for security)")
                else:
                    print(f"  {k}: {v}")
            print()

        # Store the token
        self.store_token(token)

        # Fetch and cache user info
        self._fetch_and_cache_userinfo(extract_access_token(token))

        return token

    def _run_callback_server(self, timeout=300):
        """
        Run a local Flask server to capture the OAuth callback.

        Args:
            timeout: Maximum time to wait for callback in seconds

        Returns:
            str: Authorization code or None if failed
        """
        # Clear previous state but preserve verifier and state
        code_verifier = server_token.get("verifier")
        state = server_token.get("state")
        server_token.clear()
        if code_verifier:
            server_token["verifier"] = code_verifier
        if state:
            server_token["state"] = state
        shutdown_event.clear()

        # Extract host and port from redirect URI
        parsed_uri = urlparse(self.redirect_uri)
        host = parsed_uri.hostname or "127.0.0.1"
        port = parsed_uri.port or 8042  # Default to 8042 if not specified

        # Create server in a separate thread with logging disabled
        import logging

        log = logging.getLogger("werkzeug")
        log.setLevel(logging.ERROR)  # Suppress INFO logs from Flask/Werkzeug

        server = make_server(host, port, app, threaded=True)
        server_thread = threading.Thread(target=server.serve_forever)
        server_thread.daemon = True
        server_thread.start()

        # Wait for callback or timeout
        if not shutdown_event.wait(timeout):
            server.shutdown()
            raise AuthenticationError("Authentication timed out")

        server.shutdown()

        if "error" in server_token:
            raise AuthenticationError(f"Authentication failed: {server_token['error']}")

        return server_token.get("code")

    def _exchange_code_for_token(self, code, code_verifier):
        """
        Exchange authorization code for access and refresh tokens.

        Args:
            code: Authorization code from callback
            code_verifier: PKCE code verifier

        Returns:
            dict: Token response
        """
        data = {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": self.redirect_uri,
            "client_id": self.client_id,
            "code_verifier": code_verifier,
        }

        response = httpx.post(
            self.token_endpoint, data=data, headers={"Content-Type": "application/x-www-form-urlencoded"}, timeout=30
        )

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

        token = response.json()

        # Add metadata using utility function
        enhanced_token = add_token_metadata(token)

        return enhanced_token

    def refresh_token(self, refresh_token):
        """
        Refresh an access token using a refresh token.

        Args:
            refresh_token: Refresh token to use

        Returns:
            dict: New token response
        """
        data = {
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
            "client_id": self.client_id,
        }

        response = httpx.post(
            self.token_endpoint, data=data, headers={"Content-Type": "application/x-www-form-urlencoded"}, timeout=30
        )

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

        token = response.json()

        # Add metadata using utility function
        enhanced_token = add_token_metadata(token)

        # Store the refreshed token
        self.store_token(enhanced_token)

        # Refresh user info with new token
        self._fetch_and_cache_userinfo(extract_access_token(enhanced_token))

        return enhanced_token

    def get_user_info(self, access_token):
        """
        Get user information from Okta.

        Args:
            access_token: Valid access token

        Returns:
            dict: User information
        """
        response = httpx.get(self.userinfo_endpoint, headers={"Authorization": f"Bearer {access_token}"}, timeout=30)

        if response.status_code != 200:
            raise AuthenticationError(f"Failed to get user info: {response.text}")

        return response.json()

    def get_extended_user_info(self, access_token=None):
        """
        Get extended user information including groups and metadata from multiple sources.

        This method combines information from:
        1. Okta UserInfo endpoint (basic profile + groups if available)
        2. JWT token claims
        3. CloudZero backend user verification

        Args:
            access_token: Valid access token (optional, will get current if not provided)

        Returns:
            dict: Extended user information with groups, department, title, etc.
        """
        if not access_token:
            access_token = self.get_access_token()
            if not access_token:
                raise AuthenticationError("No valid access token available")

        # Get basic user info from Okta
        userinfo = self.get_user_info(access_token)

        # Store original Okta response for debugging
        userinfo["_debug_okta_raw_response"] = dict(userinfo)  # Make a copy for debugging

        # Try to get additional info from JWT tokens (both Access Token and ID Token)
        try:
            import base64
            import json

            # Decode Access Token payload to get additional claims
            parts = access_token.split(".")
            if len(parts) == 3:
                payload_data = parts[1] + "=" * (4 - len(parts[1]) % 4)  # Add padding
                payload = json.loads(base64.urlsafe_b64decode(payload_data))

                # Add Access Token JWT claims to userinfo if they exist
                access_token_claims = {
                    "uid": payload.get("uid"),
                    "groups_from_access_token": payload.get("groups", []),
                    "scp": payload.get("scp", []),
                    "auth_time": payload.get("auth_time"),
                    "iat": payload.get("iat"),
                    "exp": payload.get("exp"),
                }
                userinfo.update({k: v for k, v in access_token_claims.items() if v is not None})
        except Exception:
            # If JWT parsing fails, continue with basic userinfo
            pass

        # Also check ID Token for additional claims (especially groups)
        try:
            import base64
            import json

            # Get stored token to check for ID token
            stored_token = self.get_stored_token()
            if stored_token and "id_token" in stored_token:
                id_token = stored_token["id_token"]
                parts = id_token.split(".")
                if len(parts) == 3:
                    payload_data = parts[1] + "=" * (4 - len(parts[1]) % 4)  # Add padding
                    id_payload = json.loads(base64.urlsafe_b64decode(payload_data))

                    # Add ID Token JWT claims to userinfo if they exist
                    id_token_claims = {
                        "groups_from_id_token": id_payload.get("groups", []),
                        "id_token_uid": id_payload.get("uid"),
                        "id_token_sub": id_payload.get("sub"),
                        "id_token_aud": id_payload.get("aud"),
                        "id_token_iss": id_payload.get("iss"),
                        "id_token_exp": id_payload.get("exp"),
                        "id_token_iat": id_payload.get("iat"),
                    }
                    userinfo.update({k: v for k, v in id_token_claims.items() if v is not None})

                    # If ID token has groups, also add them to the general groups_from_jwt field
                    if id_payload.get("groups"):
                        userinfo["groups_from_jwt"] = id_payload.get("groups", [])
        except Exception:
            # If ID Token parsing fails, continue
            pass

        # Try to get additional metadata from CloudZero backend
        try:
            from src.api.authenticated import get_authenticated

            backend_info = get_authenticated("auth/verify", debug=False)

            # Merge backend user metadata
            backend_fields = {
                "user_groups": backend_info.get("user_groups", []),
                "user_department": backend_info.get("user_department"),
                "user_title": backend_info.get("user_title"),
                "user_full_name": backend_info.get("user_full_name"),
                "backend_user_id": backend_info.get("user_id"),
            }
            userinfo.update({k: v for k, v in backend_fields.items() if v is not None})
        except Exception:
            # If backend call fails, continue with Okta data only
            pass

        return userinfo

    def _fetch_and_cache_userinfo(self, access_token):
        """
        Fetch user info from Okta and cache it.

        Args:
            access_token: Valid access token
        """
        if not access_token:
            return

        try:
            userinfo = self.get_user_info(access_token)
            # Add metadata using utility function
            token = self.get_stored_token()
            token_expires_at = token.get("expires_at") if token else datetime.now(UTC).isoformat()
            enhanced_userinfo = add_userinfo_metadata(userinfo, token_expires_at)
            self.store_userinfo(enhanced_userinfo)
        except Exception:
            # If userinfo fetch fails, don't break the authentication flow
            pass

    def store_userinfo(self, userinfo):
        """Store user info securely using file storage."""
        try:
            self.token_storage.store_userinfo(userinfo)
        except TokenStorageError as e:
            raise AuthenticationError(f"Failed to store user info: {e}") from e

    def get_stored_userinfo(self):
        """Retrieve stored user info."""
        try:
            return self.token_storage.get_stored_userinfo()
        except TokenStorageError as e:
            raise AuthenticationError(f"Failed to retrieve user info: {e}") from e

    def is_userinfo_valid(self, userinfo):
        """Check if cached userinfo is still valid based on token expiry."""
        return is_userinfo_valid(userinfo)

    def get_cached_userinfo(self):
        """
        Get cached user info if valid, otherwise fetch fresh info.

        Returns:
            dict: User information or None if not available
        """
        # Try to get cached userinfo first
        cached_userinfo = self.get_stored_userinfo()
        if cached_userinfo and self.is_userinfo_valid(cached_userinfo):
            return cached_userinfo

        # If cached userinfo is invalid, try to fetch fresh info
        access_token = self.get_access_token()
        if access_token:
            try:
                self._fetch_and_cache_userinfo(access_token)
                return self.get_stored_userinfo()
            except Exception:
                pass

        return cached_userinfo  # Return cached even if expired, better than nothing

    def store_token(self, token):
        """Store token securely using file storage."""
        try:
            self.token_storage.store_token(token)
        except TokenStorageError as e:
            raise AuthenticationError(f"Failed to store token: {e}") from e

    def get_stored_token(self):
        """Retrieve stored token."""
        try:
            return self.token_storage.get_stored_token()
        except TokenStorageError as e:
            raise AuthenticationError(f"Failed to retrieve token: {e}") from e

    def clear_token(self):
        """Clear stored token and userinfo."""
        try:
            self.token_storage.clear_all()
        except TokenStorageError as e:
            raise AuthenticationError(f"Failed to clear tokens: {e}") from e

    def is_token_valid(self, token):
        """Check if token is still valid."""
        return is_token_valid(token)

    def is_authenticated(self):
        """Check if user is authenticated with a valid token."""
        token = self.get_stored_token()
        return token is not None and self.is_token_valid(token)

    def get_access_token(self):
        """
        Get a valid access token, refreshing if necessary.

        Returns:
            str: Valid access token or None
        """
        token = self.get_stored_token()
        if not token:
            return None

        # If token is valid, return it
        if self.is_token_valid(token):
            return extract_access_token(token)

        # Try to refresh if we have a refresh token
        refresh_token_str = extract_refresh_token(token)
        if refresh_token_str:
            try:
                refreshed = self.refresh_token(refresh_token_str)
                return extract_access_token(refreshed)
            except Exception:
                # Refresh failed, need to re-authenticate
                pass

        return None


class AuthenticationError(Exception):
    """Authentication-related errors."""

    pass
