#include "services/OidcRpService.h"

#include "common/Crypto.h"
#include "common/OidcCrypto.h"
#include "common/OidcHttp.h"
#include "common/OidcJwt.h"
#include "common/OidcRpConfig.h"
#include "common/OidcUserClaims.h"
#include "common/Validation.h"

#include <drogon/HttpClient.h>
#include <drogon/drogon.h>

#include <chrono>
#include <map>
#include <sstream>

namespace visione::services
{
namespace
{
constexpr std::string_view kCieProfile = "cie";
constexpr std::string_view kProviderAlias = "cie_demo";
constexpr std::string_view kOidcScope = "openid cie_demo";

Json::Value failure(const std::string &code, const std::string &message = {})
{
    Json::Value body;
    body["ok"] = false;
    body["error"] = code;
    if (!message.empty())
    {
        body["message"] = message;
    }
    return body;
}

Json::Value success()
{
    Json::Value body;
    body["ok"] = true;
    return body;
}

long long nowEpoch()
{
    return std::chrono::system_clock::to_time_t(
        std::chrono::system_clock::now());
}

int oidcTimeoutSeconds()
{
    return drogon::app().getCustomConfig()["oidc_rp"]
        .get("http_timeout_seconds", 8)
        .asInt();
}

std::string redirectUrlWithState(const std::string &baseUrl,
                                 const std::string &state)
{
    if (baseUrl.empty())
    {
        return {};
    }
    return baseUrl + (baseUrl.find('?') == std::string::npos ? "?" : "&") +
           "state=" + drogon::utils::urlEncode(state);
}

std::string issueVisioneToken(long long userId)
{
    const auto token = crypto::randomToken();
    const auto hours = drogon::app()
                           .getCustomConfig()["auth"]
                           .get("token_ttl_hours", 24)
                           .asInt();
    drogon::app().getDbClient()->execSqlSync(
        "INSERT INTO api_tokens (user_id, token_hash, expires_at) "
        "VALUES (?, ?, DATE_ADD(NOW(), INTERVAL ? HOUR))",
        userId,
        crypto::hashSecret(token),
        hours);
    return token;
}

std::vector<oidc::Jwk> parseJwks(const Json::Value &jwks)
{
    std::vector<oidc::Jwk> out;
    const auto &keys = jwks.isObject() ? jwks["keys"] : jwks;
    if (!keys.isArray())
    {
        return out;
    }
    for (const auto &item : keys)
    {
        out.push_back(oidc::parseJwk(item));
    }
    return out;
}

std::string sessionPublicStatus(const std::string &internalStatus)
{
    if (internalStatus == "AUTHENTICATED")
    {
        return "authenticated";
    }
    if (internalStatus == "FAILED")
    {
        return "failed";
    }
    return "pending";
}

Json::Value fetchUserinfoClaims(const oidc::RpConfig &config,
                                const Json::Value &providerMetadata,
                                const std::vector<oidc::Jwk> &providerJwks,
                                const std::string &accessToken)
{
    const auto userinfoEndpoint =
        providerMetadata.get("userinfo_endpoint", "").asString();
    if (userinfoEndpoint.empty())
    {
        throw std::runtime_error("userinfo_endpoint_missing");
    }

    const auto response = oidc::httpGet(
        userinfoEndpoint,
        {{"Authorization", "Bearer " + accessToken}},
        oidcTimeoutSeconds());
    if (response.statusCode != 200)
    {
        throw std::runtime_error("userinfo_request_failed");
    }

    try
    {
        return oidc::parseJsonResponse(response);
    }
    catch (const std::exception &)
    {
    }

    const auto jwe = oidc::parseJwe(response.body);
    const auto rpKid = jwe.header.get("kid", "").asString();
    const auto rpJwk = oidc::findJwkByKid(config.coreJwks, rpKid);
    if (!rpJwk)
    {
        throw std::runtime_error("userinfo_decryption_key_not_found");
    }

    auto jws = oidc::decryptJwe(response.body, *rpJwk);
    const auto jwsParts = oidc::parseJwt(jws);
    const auto providerJwk = oidc::findJwkByKid(
        providerJwks, jwsParts.header.get("kid", "").asString());
    if (!providerJwk || !oidc::verifyJwtSignature(jws, *providerJwk))
    {
        throw std::runtime_error("invalid_userinfo_signature");
    }

    return oidc::parseJsonObjectJwtPayload(jws);
}

std::string defaultFrontendRedirect(const oidc::RpConfig &config)
{
    return config.postLoginRedirectUri;
}

Json::Value tokenPayloadForResponse(const std::string &token)
{
    Json::Value body;
    body["token"] = token;
    body["token_type"] = "Bearer";
    body["expires_in_hours"] = drogon::app()
                                   .getCustomConfig()["auth"]
                                   .get("token_ttl_hours", 24)
                                   .asInt();
    return body;
}

long long findOrCreateCitizen(const oidc::NormalizedUserClaims &claims)
{
    auto db = drogon::app().getDbClient();

    const auto ext = db->execSqlSync(
        "SELECT user_id FROM user_external_identities "
        "WHERE provider = ? AND issuer = ? AND subject = ? LIMIT 1",
        std::string(kProviderAlias),
        claims.issuer,
        claims.subject);
    if (!ext.empty())
    {
        return ext[0]["user_id"].as<long long>();
    }

    const auto userByFiscal = db->execSqlSync(
        "SELECT id FROM users WHERE fiscal_code = ? LIMIT 1", claims.fiscalCode);
    long long userId = 0;
    if (!userByFiscal.empty())
    {
        userId = userByFiscal[0]["id"].as<long long>();
        db->execSqlSync(
            "UPDATE users SET first_name = COALESCE(NULLIF(?, ''), first_name), "
            "last_name = COALESCE(NULLIF(?, ''), last_name), "
            "email = COALESCE(NULLIF(?, ''), email) WHERE id = ?",
            claims.firstName,
            claims.lastName,
            claims.email,
            userId);
    }
    else
    {
        const auto email = !claims.email.empty()
                               ? claims.email
                               : (claims.fiscalCode + "@cie.local");
        db->execSqlSync(
            "INSERT INTO users "
            "(first_name, last_name, email, fiscal_code, birth_date, birth_place, phone) "
            "VALUES (?, ?, ?, ?, NULL, NULL, NULL)",
            claims.firstName.empty() ? claims.fiscalCode : claims.firstName,
            claims.lastName.empty() ? claims.fiscalCode : claims.lastName,
            email,
            claims.fiscalCode);
        const auto inserted = db->execSqlSync("SELECT LAST_INSERT_ID() AS id");
        userId = inserted[0]["id"].as<long long>();
        db->execSqlSync(
            "INSERT INTO user_roles (user_id, role_id) "
            "SELECT ?, id FROM roles WHERE name = 'citizen'",
            userId);
    }

    Json::StreamWriterBuilder writer;
    writer["indentation"] = "";
    db->execSqlSync(
        "INSERT INTO user_external_identities "
        "(user_id, provider, issuer, subject, fiscal_code, raw_claims) "
        "VALUES (?, ?, ?, ?, ?, CAST(? AS JSON)) "
        "ON DUPLICATE KEY UPDATE user_id = VALUES(user_id), "
        "fiscal_code = VALUES(fiscal_code), raw_claims = VALUES(raw_claims)",
        userId,
        std::string(kProviderAlias),
        claims.issuer,
        claims.subject,
        claims.fiscalCode,
        Json::writeString(writer, claims.rawClaims));
    return userId;
}

Json::Value providerResolveDocument(const oidc::RpProviderConfig &provider)
{
    const auto url = provider.trustAnchor + "/resolve?sub=" +
                     drogon::utils::urlEncode(provider.issuer) +
                     "&trust_anchor=" + drogon::utils::urlEncode(provider.trustAnchor) +
                     "&format=json";
    return oidc::parseJsonResponse(oidc::httpGet(url, oidcTimeoutSeconds()));
}
}  // namespace

Json::Value OidcRpService::startLogin(const Json::Value &body)
{
    const auto config = oidc::loadRpConfig();
    if (!config.enabled)
    {
        return failure("cie_login_disabled");
    }

    const auto providerName = body.get("provider", std::string(kProviderAlias)).asString();
    const auto provider = oidc::findProviderConfig(config, providerName);
    if (!provider)
    {
        return failure("unknown_provider");
    }

    const auto resolveDoc = providerResolveDocument(*provider);
    const auto metadata = resolveDoc["metadata"]["openid_provider"];
    if (!metadata.isObject())
    {
        return failure("provider_metadata_missing");
    }

    const auto authorizationEndpoint =
        metadata.get("authorization_endpoint", "").asString();
    const auto tokenEndpoint = metadata.get("token_endpoint", "").asString();
    if (authorizationEndpoint.empty() || tokenEndpoint.empty())
    {
        return failure("provider_endpoints_missing");
    }

    const auto state = oidc::randomUrlSafeString(24);
    const auto nonce = oidc::randomUrlSafeString(24);
    const auto pkceVerifier = oidc::randomUrlSafeString(32);
    const auto codeChallenge = oidc::sha256Base64Url(pkceVerifier);

    Json::Value claims;
    claims["id_token"]["given_name"]["essential"] = true;
    claims["id_token"]["family_name"]["essential"] = true;
    claims["userinfo"]["given_name"] = Json::nullValue;
    claims["userinfo"]["family_name"] = Json::nullValue;
    claims["userinfo"]["email"] = Json::nullValue;
    claims["userinfo"]["https://attributes.eid.gov.it/fiscal_number"] =
        Json::nullValue;

    Json::Value requestHeader;
    requestHeader["alg"] = "RS256";
    requestHeader["typ"] = "JWT";
    requestHeader["kid"] = config.coreJwks[1].kid;

    Json::Value requestPayload;
    requestPayload["iss"] = config.clientId;
    requestPayload["client_id"] = config.clientId;
    requestPayload["response_type"] = "code";
    requestPayload["scope"] = std::string(kOidcScope);
    requestPayload["redirect_uri"] = config.redirectUri;
    requestPayload["nonce"] = nonce;
    requestPayload["state"] = state;
    requestPayload["code_challenge"] = codeChallenge;
    requestPayload["code_challenge_method"] = "S256";
    requestPayload["acr_values"].append("https://www.spid.gov.it/SpidL2");
    requestPayload["claims"] = claims;
    requestPayload["prompt"] = "consent login";
    requestPayload["iat"] = static_cast<Json::Int64>(nowEpoch());
    requestPayload["exp"] = static_cast<Json::Int64>(nowEpoch() + 60);
    requestPayload["jti"] = drogon::utils::getUuid();
    requestPayload["aud"].append(provider->issuer);
    requestPayload["aud"].append(authorizationEndpoint);

    const auto requestJwt =
        oidc::signJwt(requestHeader, requestPayload, config.coreJwks[1]);

    Json::StreamWriterBuilder writer;
    writer["indentation"] = "";
    const auto frontendRedirect =
        body.get("frontend_redirect_uri", defaultFrontendRedirect(config)).asString();
    drogon::app().getDbClient()->execSqlSync(
        "INSERT INTO oidc_authentications "
        "(provider, provider_issuer, trust_anchor, state, nonce, pkce_verifier, "
        "code_challenge, frontend_redirect_uri, authorization_endpoint, "
        "token_endpoint, userinfo_endpoint, federation_resolve_endpoint, "
        "request_object_jwt, provider_metadata, trust_chain_json, status) "
        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, CAST(? AS JSON), "
        "CAST(? AS JSON), 'AUTHORIZATION_REDIRECT_READY')",
        providerName,
        provider->issuer,
        provider->trustAnchor,
        state,
        nonce,
        pkceVerifier,
        codeChallenge,
        frontendRedirect,
        authorizationEndpoint,
        tokenEndpoint,
        metadata.get("userinfo_endpoint", "").asString(),
        provider->trustAnchor + "/resolve",
        requestJwt,
        Json::writeString(writer, metadata),
        Json::writeString(writer, resolveDoc));

    std::ostringstream authorizationUrl;
    authorizationUrl << authorizationEndpoint
                     << "?client_id=" << drogon::utils::urlEncode(config.clientId)
                     << "&scope=" << drogon::utils::urlEncode(std::string(kOidcScope))
                     << "&response_type=code"
                     << "&code_challenge=" << drogon::utils::urlEncode(codeChallenge)
                     << "&code_challenge_method=S256"
                     << "&request=" << drogon::utils::urlEncode(requestJwt);

    auto response = success();
    response["status"] = "authorization_redirect_ready";
    response["provider"] = providerName;
    response["state"] = state;
    response["authorization_url"] = authorizationUrl.str();
    return response;
}

Json::Value OidcRpService::callback(const std::string &state,
                                    const std::string &code,
                                    const std::string &error,
                                    const std::string &errorDescription,
                                    const std::optional<std::string> &issuer)
{
    const auto config = oidc::loadRpConfig();
    auto db = drogon::app().getDbClient();

    const auto sessions = db->execSqlSync(
        "SELECT id, provider, provider_issuer, nonce, pkce_verifier, token_endpoint, "
        "frontend_redirect_uri, provider_metadata FROM oidc_authentications "
        "WHERE state = ? LIMIT 1",
        state);
    if (sessions.empty())
    {
        return failure("session_not_found");
    }

    const auto authId = sessions[0]["id"].as<long long>();
    const auto frontendRedirect =
        sessions[0]["frontend_redirect_uri"].as<std::string>();

    auto finishRedirect = [&](const std::string &statusCode,
                              const std::string &message) {
        db->execSqlSync(
            "UPDATE oidc_authentications SET status = 'FAILED', error_code = ?, "
            "error_message = ?, completed_at = NOW() WHERE id = ?",
            statusCode,
            message,
            authId);
        auto response = failure(statusCode, message);
        response["redirect_url"] = redirectUrlWithState(frontendRedirect, state);
        return response;
    };

    if (!error.empty())
    {
        return finishRedirect(error, errorDescription);
    }
    if (code.empty())
    {
        return finishRedirect("missing_code", "Provider callback did not include code");
    }
    if (issuer && *issuer != sessions[0]["provider_issuer"].as<std::string>())
    {
        return finishRedirect("issuer_mismatch", "Provider issuer mismatch");
    }
    db->execSqlSync(
        "UPDATE oidc_authentications SET status = 'CALLBACK_RECEIVED' WHERE id = ?",
        authId);

    Json::Value providerMetadata;
    {
        Json::CharReaderBuilder builder;
        std::string errors;
        std::unique_ptr<Json::CharReader> reader(builder.newCharReader());
        const auto providerMetadataText =
            sessions[0]["provider_metadata"].as<std::string>();
        reader->parse(providerMetadataText.data(),
                      providerMetadataText.data() + providerMetadataText.size(),
                      &providerMetadata,
                      &errors);
    }

    Json::Value assertionHeader;
    assertionHeader["alg"] = "RS256";
    assertionHeader["typ"] = "JWT";
    assertionHeader["kid"] = config.coreJwks[1].kid;
    Json::Value assertionPayload;
    assertionPayload["iss"] = config.clientId;
    assertionPayload["sub"] = config.clientId;
    assertionPayload["aud"].append(sessions[0]["token_endpoint"].as<std::string>());
    assertionPayload["iat"] = static_cast<Json::Int64>(nowEpoch());
    assertionPayload["exp"] = static_cast<Json::Int64>(nowEpoch() + 60);
    assertionPayload["jti"] = drogon::utils::getUuid();

    const auto clientAssertion =
        oidc::signJwt(assertionHeader, assertionPayload, config.coreJwks[1]);

    Json::Value tokenResponse;
    try
    {
        tokenResponse = oidc::parseJsonResponse(oidc::httpPostForm(
            sessions[0]["token_endpoint"].as<std::string>(),
            {
                {"grant_type", "authorization_code"},
                {"redirect_uri", config.redirectUri},
                {"client_id", config.clientId},
                {"state", state},
                {"code", code},
                {"code_verifier", sessions[0]["pkce_verifier"].as<std::string>()},
                {"client_assertion_type",
                 "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"},
                {"client_assertion", clientAssertion},
            },
            oidcTimeoutSeconds()));
    }
    catch (const std::exception &e)
    {
        return finishRedirect("token_exchange_failed", e.what());
    }

    if (!tokenResponse.isObject() || !tokenResponse.isMember("id_token") ||
        !tokenResponse.isMember("access_token"))
    {
        return finishRedirect("invalid_token_response",
                              "Provider token response missing tokens");
    }
    db->execSqlSync(
        "UPDATE oidc_authentications SET status = 'TOKENS_RECEIVED' WHERE id = ?",
        authId);

    const auto providerJwks = parseJwks(providerMetadata["jwks"]);
    const auto idToken = tokenResponse["id_token"].asString();
    const auto accessToken = tokenResponse["access_token"].asString();
    const auto parsedId = oidc::parseJwt(idToken);
    const auto parsedAccess = oidc::parseJwt(accessToken);
    const auto idJwk = oidc::findJwkByKid(providerJwks,
                                          parsedId.header.get("kid", "").asString());
    const auto accessJwk =
        oidc::findJwkByKid(providerJwks,
                           parsedAccess.header.get("kid", "").asString());
    if (!idJwk || !accessJwk ||
        !oidc::verifyJwtSignature(idToken, *idJwk) ||
        !oidc::verifyJwtSignature(accessToken, *accessJwk))
    {
        return finishRedirect("invalid_provider_signature",
                              "Provider token signature verification failed");
    }

    const auto idPayload = parsedId.payload;
    if (idPayload.get("iss", "").asString() !=
        sessions[0]["provider_issuer"].as<std::string>())
    {
        return finishRedirect("id_token_issuer_mismatch", "ID token issuer mismatch");
    }
    const auto aud = idPayload["aud"];
    bool audienceOk = false;
    if (aud.isArray())
    {
        for (const auto &item : aud)
        {
            if (item.asString() == config.clientId)
            {
                audienceOk = true;
                break;
            }
        }
    }
    else
    {
        audienceOk = aud.asString() == config.clientId;
    }
    if (!audienceOk || idPayload.get("nonce", "").asString() !=
                           sessions[0]["nonce"].as<std::string>())
    {
        return finishRedirect("id_token_validation_failed",
                              "ID token audience or nonce mismatch");
    }
    if (idPayload.get("at_hash", "").asString() !=
        oidc::leftHashBase64Url(accessToken))
    {
        return finishRedirect("at_hash_mismatch", "ID token at_hash mismatch");
    }

    Json::Value userinfoClaims;
    try
    {
        userinfoClaims =
            fetchUserinfoClaims(config, providerMetadata, providerJwks, accessToken);
    }
    catch (const std::exception &e)
    {
        return finishRedirect("invalid_userinfo_response", e.what());
    }
    db->execSqlSync(
        "UPDATE oidc_authentications SET status = 'USERINFO_RECEIVED' WHERE id = ?",
        authId);

    auto normalized = oidc::mergeUserClaims(
        sessions[0]["provider_issuer"].as<std::string>(),
        idPayload,
        userinfoClaims);
    if (!normalized)
    {
        return finishRedirect("missing_identity_claims",
                              "Required identity claims missing from merged OIDC claims");
    }

    const auto userId = findOrCreateCitizen(*normalized);
    const auto sessionToken = issueVisioneToken(userId);

    Json::StreamWriterBuilder writer;
    writer["indentation"] = "";
    db->execSqlSync(
        "INSERT INTO oidc_authentication_tokens "
        "(authentication_id, authorization_code, access_token, id_token, refresh_token, "
        "token_type, scope, expires_in_seconds) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
        authId,
        code,
        accessToken,
        idToken,
        tokenResponse.get("refresh_token", "").asString(),
        tokenResponse.get("token_type", "Bearer").asString(),
        tokenResponse.get("scope", "openid").asString(),
        tokenResponse.get("expires_in", 0).asInt());
    db->execSqlSync(
        "UPDATE oidc_authentications SET status = 'AUTHENTICATED', "
        "authenticated_user_id = ?, session_token = ?, "
        "session_token_expires_at = DATE_ADD(NOW(), INTERVAL ? HOUR), "
        "completed_at = NOW() WHERE id = ?",
        userId,
        sessionToken,
        drogon::app().getCustomConfig()["auth"].get("token_ttl_hours", 24).asInt(),
        authId);

    auto response = success();
    response["redirect_url"] = redirectUrlWithState(frontendRedirect, state);
    return response;
}

Json::Value OidcRpService::sessionStatus(const std::string &state)
{
    auto db = drogon::app().getDbClient();
    const auto result = db->execSqlSync(
        "SELECT state, status, provider, frontend_redirect_uri, "
        "authenticated_user_id, session_token, error_code, error_message, created_at, "
        "completed_at FROM oidc_authentications WHERE state = ? LIMIT 1",
        state);
    if (result.empty())
    {
        return failure("session_not_found");
    }

    auto response = success();
    response["state"] = result[0]["state"].as<std::string>();
    response["status"] =
        sessionPublicStatus(result[0]["status"].as<std::string>());
    response["phase"] = result[0]["status"].as<std::string>();
    response["provider"] = result[0]["provider"].as<std::string>();
    response["frontend_redirect_uri"] =
        result[0]["frontend_redirect_uri"].as<std::string>();
    if (!result[0]["authenticated_user_id"].isNull())
    {
        response["authenticated_user_id"] =
            Json::Int64(result[0]["authenticated_user_id"].as<long long>());
    }
    response["error_code"] = result[0]["error_code"].as<std::string>();
    response["error_message"] = result[0]["error_message"].as<std::string>();
    response["created_at"] = result[0]["created_at"].as<std::string>();
    response["completed_at"] = result[0]["completed_at"].isNull()
                                   ? ""
                                   : result[0]["completed_at"].as<std::string>();
    if (!result[0]["session_token"].isNull() &&
        !result[0]["session_token"].as<std::string>().empty())
    {
        response["auth"] = tokenPayloadForResponse(
            result[0]["session_token"].as<std::string>());
    }
    return response;
}

Json::Value OidcRpService::entityConfigurationDocument() const
{
    const auto config = oidc::loadRpConfig();
    return oidc::entityConfigurationDocument(config);
}

std::string OidcRpService::entityConfigurationJwt() const
{
    const auto config = oidc::loadRpConfig();
    Json::Value header;
    header["alg"] = "RS256";
    header["typ"] = "entity-statement+jwt";
    header["kid"] = config.federationJwks.front().kid;
    return oidc::signJwt(
        header, oidc::entityConfigurationDocument(config), config.federationJwks.front());
}

Json::Value OidcRpService::resolveDocument(const std::optional<std::string> &sub,
                                           const std::optional<std::string> &anchor) const
{
    const auto config = oidc::loadRpConfig();
    if (!sub || !anchor || *sub != config.clientId)
    {
        return failure("not_found");
    }
    bool anchorAllowed = false;
    for (const auto &value : config.trustAnchors)
    {
        if (value == *anchor)
        {
            anchorAllowed = true;
            break;
        }
    }
    if (!anchorAllowed)
    {
        return failure("not_found");
    }

    Json::Value body;
    body["iss"] = config.clientId;
    body["sub"] = config.clientId;
    body["iat"] = static_cast<Json::Int64>(nowEpoch());
    body["exp"] = static_cast<Json::Int64>(nowEpoch() + 2880 * 60);
    if (!config.trustMark.isNull())
    {
        body["trust_marks"].append(config.trustMark);
    }
    body["metadata"] = oidc::rpMetadataDocument(config);
    body["trust_chain"] = Json::arrayValue;
    return body;
}

std::string OidcRpService::resolveJwt(const std::optional<std::string> &sub,
                                      const std::optional<std::string> &anchor) const
{
    const auto document = resolveDocument(sub, anchor);
    if (document.get("ok", false).asBool() == false && document.isMember("error"))
    {
        return {};
    }
    const auto config = oidc::loadRpConfig();
    Json::Value header;
    header["alg"] = "RS256";
    header["typ"] = "resolve-response+jwt";
    header["kid"] = config.federationJwks.front().kid;
    return oidc::signJwt(header, document, config.federationJwks.front());
}

Json::Value OidcRpService::jwks() const
{
    return oidc::rpPublicJwksDocument(oidc::loadRpConfig());
}
}  // namespace visione::services