// Copyright Epic Games, Inc. All Rights Reserved.

#include "zenhttp/auth/oidc.h"
#include <zenhttp/httpclient.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <fmt/format.h>
#include <json11.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

namespace details {

	using StringArray = std::vector<std::string>;

	StringArray ToStringArray(const json11::Json JsonArray)
	{
		StringArray Result;

		const auto& Items = JsonArray.array_items();

		for (const auto& Item : Items)
		{
			Result.push_back(Item.string_value());
		}

		return Result;
	}

}  // namespace details

using namespace std::literals;

// Return the "scheme://host[:port]" prefix of an absolute URL, or empty on
// malformed input.  RFC 8414 §3 requires every endpoint advertised by an
// OP to live under the issuer's origin, so exact-prefix comparison on this
// value is sufficient to pin an endpoint to the configured IdP.
static std::string_view
OriginOf(std::string_view Url)
{
	// scheme://...
	auto SchemeEnd = Url.find("://");
	if (SchemeEnd == std::string_view::npos)
	{
		return {};
	}
	// First path char (or query / fragment) after the authority.
	const size_t AuthorityStart = SchemeEnd + 3;
	size_t		 OriginEnd		= Url.size();
	for (size_t I = AuthorityStart; I < Url.size(); ++I)
	{
		const char C = Url[I];
		if (C == '/' || C == '?' || C == '#')
		{
			OriginEnd = I;
			break;
		}
	}
	// Require at least one character in the authority.
	if (OriginEnd == AuthorityStart)
	{
		return {};
	}
	return Url.substr(0, OriginEnd);
}

// True if the URL uses a scheme / host that we trust for discovery even
// without HTTPS — narrowly, only loopback for dev / test setups.
static bool
IsLoopbackHttp(std::string_view Url)
{
	constexpr std::string_view Prefixes[] = {
		"http://localhost"sv,
		"http://127.0.0.1"sv,
		"http://[::1]"sv,
	};
	for (std::string_view P : Prefixes)
	{
		if (Url.size() >= P.size() && Url.substr(0, P.size()) == P)
		{
			// Ensure the next char (if any) ends the authority cleanly.
			if (Url.size() == P.size())
			{
				return true;
			}
			const char Next = Url[P.size()];
			if (Next == ':' || Next == '/' || Next == '?' || Next == '#')
			{
				return true;
			}
		}
	}
	return false;
}

static std::string
FormUrlEncode(std::string_view Input)
{
	std::string Result;
	Result.reserve(Input.size());
	for (char C : Input)
	{
		if ((C >= 'A' && C <= 'Z') || (C >= 'a' && C <= 'z') || (C >= '0' && C <= '9') || C == '-' || C == '_' || C == '.' || C == '~')
		{
			Result.push_back(C);
		}
		else
		{
			Result.append(fmt::format("%{:02X}", static_cast<uint8_t>(C)));
		}
	}
	return Result;
}

OidcClient::OidcClient(const OidcClient::Options& Options)
{
	m_BaseUrl  = std::string(Options.BaseUrl);
	m_ClientId = std::string(Options.ClientId);
}

OidcClient::InitResult
OidcClient::Initialize()
{
	// The OIDC discovery document determines where we send refresh tokens, so
	// the transport to the discovery endpoint has to be trustworthy.  Require
	// HTTPS on the configured BaseUrl.  Loopback is permitted over plain HTTP
	// for developer setups that run a local IdP mock — no meaningful attack
	// surface on the loopback interface.
	if (m_BaseUrl.size() < 8 || m_BaseUrl.substr(0, 8) != "https://"sv)
	{
		if (!IsLoopbackHttp(m_BaseUrl))
		{
			return {.Reason = "BaseUrl must use https:// (or a http://localhost / 127.0.0.1 / [::1] loopback)"};
		}
	}

	HttpClient			 Http{m_BaseUrl};
	HttpClient::Response Response = Http.Get("/.well-known/openid-configuration"sv);

	if (!Response)
	{
		return {.Reason = Response.ErrorMessage("")};
	}

	if (Response.StatusCode != HttpResponseCode::OK)
	{
		return {.Reason = std::string{ToString(Response.StatusCode)}};
	}

	std::string	 JsonError;
	json11::Json Json = json11::Json::parse(std::string{Response.AsText()}, JsonError);

	if (JsonError.empty() == false)
	{
		return {.Reason = std::move(JsonError)};
	}

	// RFC 8414 §3: the discovery document's `issuer` value MUST identify the
	// OP and MUST be the origin used to fetch the document.  Without this
	// check, an attacker who can intercept discovery (or a misconfigured
	// intermediate) can swap the issuer identity without detection.  Accept
	// a trailing '/' divergence since OPs vary.
	const std::string Issuer = Json["issuer"].string_value();
	{
		std::string_view ExpectedBase = m_BaseUrl;
		while (!ExpectedBase.empty() && ExpectedBase.back() == '/')
		{
			ExpectedBase.remove_suffix(1);
		}
		std::string_view ActualIssuer = Issuer;
		while (!ActualIssuer.empty() && ActualIssuer.back() == '/')
		{
			ActualIssuer.remove_suffix(1);
		}
		if (ActualIssuer.empty() || ActualIssuer != ExpectedBase)
		{
			return {.Reason = fmt::format("discovery issuer mismatch (expected '{}')", ExpectedBase)};
		}
	}

	// Pin every endpoint we actually use to the same origin as BaseUrl.  This
	// is the last defense against a tampered discovery document redirecting
	// token submissions to an attacker-controlled host.  We check the
	// endpoints we may call later; endpoints this client never dispatches to
	// are left alone so a discovery document with unrelated auxiliary URLs
	// isn't rejected for no reason.
	const std::string_view BaseOrigin = OriginOf(m_BaseUrl);
	if (BaseOrigin.empty())
	{
		return {.Reason = "BaseUrl is malformed"};
	}

	const std::string TokenEndpoint	   = Json["token_endpoint"].string_value();
	const std::string UserInfoEndpoint = Json["userinfo_endpoint"].string_value();
	const std::string JwksUri		   = Json["jwks_uri"].string_value();

	auto CheckOrigin = [&](std::string_view Name, std::string_view Url) -> std::optional<std::string> {
		if (Url.empty())
		{
			return std::nullopt;
		}
		const std::string_view Origin = OriginOf(Url);
		if (Origin != BaseOrigin)
		{
			return fmt::format("discovery endpoint '{}' is off-origin (expected origin '{}')", Name, BaseOrigin);
		}
		return std::nullopt;
	};

	if (auto Err = CheckOrigin("token_endpoint"sv, TokenEndpoint); Err.has_value())
	{
		return {.Reason = std::move(*Err)};
	}
	if (auto Err = CheckOrigin("userinfo_endpoint"sv, UserInfoEndpoint); Err.has_value())
	{
		return {.Reason = std::move(*Err)};
	}
	if (auto Err = CheckOrigin("jwks_uri"sv, JwksUri); Err.has_value())
	{
		return {.Reason = std::move(*Err)};
	}

	// token_endpoint is required for the refresh flow we implement; fail early
	// rather than at RefreshToken time if the OP omitted it.
	if (TokenEndpoint.empty())
	{
		return {.Reason = "discovery document is missing token_endpoint"};
	}

	m_Config = {.Issuer							   = Issuer,
				.AuthorizationEndpoint			   = Json["authorization_endpoint"].string_value(),
				.TokenEndpoint					   = TokenEndpoint,
				.UserInfoEndpoint				   = UserInfoEndpoint,
				.RegistrationEndpoint			   = Json["registration_endpoint"].string_value(),
				.EndSessionEndpoint				   = Json["end_session_endpoint"].string_value(),
				.DeviceAuthorizationEndpoint	   = Json["device_authorization_endpoint"].string_value(),
				.JwksUri						   = JwksUri,
				.SupportedResponseTypes			   = details::ToStringArray(Json["response_types_supported"]),
				.SupportedResponseModes			   = details::ToStringArray(Json["response_modes_supported"]),
				.SupportedGrantTypes			   = details::ToStringArray(Json["grant_types_supported"]),
				.SupportedScopes				   = details::ToStringArray(Json["scopes_supported"]),
				.SupportedTokenEndpointAuthMethods = details::ToStringArray(Json["token_endpoint_auth_methods_supported"]),
				.SupportedClaims				   = details::ToStringArray(Json["claims_supported"])};

	return {.Ok = true};
}

OidcClient::RefreshTokenResult
OidcClient::RefreshToken(std::string_view RefreshToken)
{
	const std::string Body =
		fmt::format("grant_type=refresh_token&refresh_token={}&client_id={}", FormUrlEncode(RefreshToken), FormUrlEncode(m_ClientId));

	HttpClient Http{m_Config.TokenEndpoint};

	HttpClient::KeyValueMap Headers{{"Content-Type", "application/x-www-form-urlencoded"}};

	HttpClient::Response Response = Http.Post("", IoBufferBuilder::MakeFromMemory(MemoryView{Body.data(), Body.size()}), Headers);

	if (!Response)
	{
		return {.Reason = std::string{Response.ErrorMessage("")}};
	}

	if (Response.StatusCode != HttpResponseCode::OK)
	{
		// Do NOT include Response.AsText() in the reason string.  Some IdPs
		// echo the submitted refresh_token (or a prefix of it) in their error
		// body — plumbing that into the Reason string causes AuthMgrImpl's
		// ZEN_WARN in the refresh paths to write the token into the log.
		// Only the status code is safe to surface up to the log sites.
		return {.Reason = fmt::format("{} (provider returned {} bytes)", ToString(Response.StatusCode), Response.AsText().size())};
	}

	std::string	 JsonError;
	json11::Json Json = json11::Json::parse(std::string{Response.AsText()}, JsonError);

	if (JsonError.empty() == false)
	{
		return {.Reason = std::move(JsonError)};
	}

	// Note: id_token is intentionally not parsed.  It is a JWT whose contents
	// are meaningful only after signature / issuer / audience / expiry
	// verification against the provider's JWKS, and nothing downstream
	// currently consumes it.  Leaving it unparsed avoids planting an
	// unauthenticated identity claim in the OpenIdToken cache.
	return {.TokenType		  = Json["token_type"].string_value(),
			.AccessToken	  = Json["access_token"].string_value(),
			.RefreshToken	  = Json["refresh_token"].string_value(),
			.Scope			  = Json["scope"].string_value(),
			.ExpiresInSeconds = static_cast<int64_t>(Json["expires_in"].int_value()),
			.Ok				  = true};
}

}  // namespace zen