const jwt = require('jsonwebtoken');
const jwksClient = require('jwks-rsa');
const User = require('../models/user');
const axios = require('axios');

// Cache for JWKS clients to avoid re-fetching the configuration on every request.
const jwksClients = {};

// Function to get the JWKS URI from the OIDC discovery endpoint.
async function getJwksUri(issuer) {
  try {
    const discoveryUrl = `${issuer}/.well-known/openid-configuration`;
    console.log(`Fetching OIDC discovery document from: ${discoveryUrl}`);
    const response = await axios.get(discoveryUrl);
    console.log(`Successfully fetched OIDC discovery document. JWKS URI: ${response.data.jwks_uri}`);
    return response.data.jwks_uri;
  } catch (error) {
    console.error('Failed to fetch OIDC discovery document:', error.message);
    throw new Error('Could not fetch JWKS URI.');
  }
}

function fetchKey(client, header, callback) {
    console.log(`Fetching signing key for kid: ${header.kid}`);
    client.getSigningKey(header.kid, (err, key) => {
        if (err) {
            console.error('Error getting signing key:', err);
            return callback(err);
        }
        const signingKey = key.publicKey || key.rsaPublicKey;
        console.log('Successfully fetched signing key.');
        callback(null, signingKey);
    });
}


async function ssoHandler(req, res, next) {
    console.log('ssoHandler started.');
    const { token } = req.query;

    if (token) {
        console.log('Token found in query. Verifying...');
        const untrustedDecoded = jwt.decode(token, { complete: true });

        if (!untrustedDecoded) {
            return res.status(400).send('Invalid token.');
        }

        const issuer = untrustedDecoded.payload.iss;
        let client = jwksClients[issuer];

        if (!client) {
            try {
                console.log(`Creating new JWKS client for issuer: ${issuer}`);
                const jwksUri = await getJwksUri(issuer);
                client = jwksClient({
                    jwksUri: jwksUri,
                    cache: true,
                    rateLimit: true,
                });
                jwksClients[issuer] = client;
            } catch (err) {
                console.error('Could not create JWKS client.', err);
                return res.status(500).send('An error occurred during SSO processing due to a configuration error.');
            }
        } else {
            console.log(`Using existing JWKS client for issuer: ${issuer}`);
        }

        jwt.verify(token, (header, callback) => fetchKey(client, header, callback), { algorithms: ['RS256'] }, async (err, decoded) => {
            if (err) {
                console.error('JWT verification failed:', err);
                return res.status(500).send('An error occurred during SSO processing due to a verification failure.');
            }

            console.log('JWT verified successfully. Decoded token:', decoded);

            try {
                if (!decoded || !decoded.sub) {
                    console.error('Invalid token: "sub" claim is missing.');
                    return res.status(400).send('Invalid token: "sub" claim is missing.');
                }

                console.log(`Looking for user with sso_sub: ${decoded.sub}`);
                let user = await User.findBySsoSub(decoded.sub);

                if (!user) {
                    console.log('User not found. Creating a new user.');
                    user = await User.createUser({ sso_sub: decoded.sub });
                    console.log('New user created:', user);
                } else {
                    console.log('User found:', user);
                }

                req.session.userId = user.id;
                console.log(`User ID ${user.id} stored in session.`);

                const redirectUrl = req.path;
                console.log(`Redirecting to ${redirectUrl}`);
                return res.redirect(redirectUrl);

            } catch (error) {
                console.error('Error during user lookup or creation:', error);
                return res.status(500).send('An error occurred during user processing.');
            }
        });
    } else {
        console.log('No token in query. Proceeding to next middleware.');
        if (req.session.userId) {
            console.log(`Session found for user ID: ${req.session.userId}`);
            res.locals.user = await User.findById(req.session.userId);
        } else {
            console.log('No session found.');
            res.locals.user = null;
        }
        return next();
    }
}

module.exports = ssoHandler;