using System.Buffers.Text;
using System.Globalization;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;

namespace DodoSSH.Client.Auth;

/// <summary>
/// Obtains an identity-provider signature over a set of public keys.
/// </summary>
/// <remarks>
/// Narrower than the whole OIDC client on purpose. Enrollment needs exactly this one capability, and
/// depending on the full client would drag discovery, token exchange and refresh into every test of
/// it — which is how a test for key binding ends up needing a stubbed token endpoint.
/// </remarks>
public interface IKeyBindingAuthorizer
{
    /// <summary>
    /// Runs an authorization whose <c>nonce</c> is a key statement's hash.
    /// </summary>
    /// <param name="bindingNonce">From <c>KeyStatementCodec.ComputeNonce</c>.</param>
    /// <param name="cancellationToken">Cancels the wait for the browser.</param>
    /// <returns>The ID token to hand to the enrollment endpoint.</returns>
    Task<string> AuthorizeKeyBindingAsync(string bindingNonce, CancellationToken cancellationToken);
}

/// <summary>
/// Authorization Code with PKCE on a loopback redirect, for a native public client.
/// </summary>
/// <remarks>
/// <para>
/// Hand-rolled rather than delegating to an OIDC library, for one specific reason: the
/// identity-provider key binding needs the <c>nonce</c> set to an exact value — the hash of a key
/// statement — and libraries generate and validate their own nonce as an internal detail. Fighting
/// that is worse than owning a flow the RFCs specify completely.
/// </para>
/// <para>
/// State and PKCE are per-request and never reused. The response is rejected unless the state
/// matches, which is what stops an attacker feeding the listener a code of their own — the loopback
/// port is reachable by any local process.
/// </para>
/// </remarks>
public sealed class OidcClient(
    HttpClient http,
    IBrowserLauncher browser,
    TimeProvider clock,
    OidcClientOptions options) : IKeyBindingAuthorizer
{
    private readonly OidcDiscoveryClient discovery = new(http);

    /// <summary>Runs an interactive sign-in and returns the resulting tokens.</summary>
    public async Task<TokenSet> SignInAsync(CancellationToken cancellationToken)
    {
        var metadata = await GetMetadataAsync(cancellationToken).ConfigureAwait(false);

        var (code, redirectUri, pkce) = await AuthorizeAsync(
                metadata,
                options.Scopes,
                extraParameters: null,
                cancellationToken)
            .ConfigureAwait(false);

        return await ExchangeAsync(
                metadata,
                new Dictionary<string, string>(StringComparer.Ordinal)
                {
                    ["grant_type"] = "authorization_code",
                    ["code"] = code,
                    ["redirect_uri"] = redirectUri.AbsoluteUri,
                    ["client_id"] = options.ClientId,
                    ["code_verifier"] = pkce.CodeVerifier,
                },
                cancellationToken)
            .ConfigureAwait(false);
    }

    /// <summary>Exchanges a refresh token for a fresh access token.</summary>
    /// <remarks>
    /// A provider may rotate the refresh token, so the caller must persist whatever comes back
    /// rather than keeping the one it sent. Treating rotation as optional is how a client ends up
    /// permanently signed out after one refresh.
    /// </remarks>
    public async Task<TokenSet> RefreshAsync(string refreshToken, CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(refreshToken);

        var metadata = await GetMetadataAsync(cancellationToken).ConfigureAwait(false);

        var refreshed = await ExchangeAsync(
                metadata,
                new Dictionary<string, string>(StringComparer.Ordinal)
                {
                    ["grant_type"] = "refresh_token",
                    ["refresh_token"] = refreshToken,
                    ["client_id"] = options.ClientId,
                },
                cancellationToken)
            .ConfigureAwait(false);

        // Providers that do not rotate omit the field entirely; carry the old one forward so the
        // caller can persist one value unconditionally.
        return refreshed.RefreshToken is null
            ? refreshed with { RefreshToken = refreshToken }
            : refreshed;
    }

    /// <summary>
    /// Runs a second authorization whose sole purpose is to have the provider sign over a set of
    /// public keys.
    /// </summary>
    /// <remarks>
    /// <para>
    /// This is the primary public-key trust anchor. The <paramref name="bindingNonce"/> is the hash
    /// of the key statement (docs/crypto.md §7.1), so the ID token that comes back is the provider's
    /// signature over exactly those keys. The DodoSSH server cannot mint that signature, so it cannot
    /// fabricate a key for a user who never enrolled — the attack that would otherwise let an
    /// operator read every vault by publishing its own key as yours. See ADR 0001.
    /// </para>
    /// <para>
    /// <c>prompt=login</c> forces a fresh authentication rather than reusing an existing session, so
    /// the assertion attests to a user present at this moment rather than to a session opened at some
    /// unknown earlier time.
    /// </para>
    /// <para>
    /// Only <c>openid</c> is requested. A second refresh token here would be one more long-lived
    /// credential to store for no benefit.
    /// </para>
    /// </remarks>
    /// <param name="bindingNonce">
    /// From <c>KeyStatementCodec.ComputeNonce</c>. Must be the statement being enrolled.
    /// </param>
    /// <param name="cancellationToken">Cancels the wait for the browser.</param>
    /// <returns>The ID token to hand to <c>POST /api/v1/me/enrollment</c>.</returns>
    public async Task<string> AuthorizeKeyBindingAsync(
        string bindingNonce,
        CancellationToken cancellationToken)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(bindingNonce);

        var metadata = await GetMetadataAsync(cancellationToken).ConfigureAwait(false);

        var (code, redirectUri, pkce) = await AuthorizeAsync(
                metadata,
                ["openid"],
                new Dictionary<string, string>(StringComparer.Ordinal)
                {
                    ["nonce"] = bindingNonce,
                    ["prompt"] = "login",
                },
                cancellationToken)
            .ConfigureAwait(false);

        var tokens = await ExchangeAsync(
                metadata,
                new Dictionary<string, string>(StringComparer.Ordinal)
                {
                    ["grant_type"] = "authorization_code",
                    ["code"] = code,
                    ["redirect_uri"] = redirectUri.AbsoluteUri,
                    ["client_id"] = options.ClientId,
                    ["code_verifier"] = pkce.CodeVerifier,
                },
                cancellationToken)
            .ConfigureAwait(false);

        if (tokens.IdToken is null)
        {
            throw new OidcException(
                "The provider returned no ID token, so nothing binds the keys. Check that the "
                + "'openid' scope is permitted for this client.");
        }

        var returned = ReadNonceClaim(tokens.IdToken);

        if (!string.Equals(returned, bindingNonce, StringComparison.Ordinal))
        {
            // A token that does not carry our nonce binds some other statement. Enrolling it would
            // store evidence that verifies against keys we are not publishing.
            throw new OidcException(
                "The ID token's nonce is not the key statement's hash, so it does not bind these keys.");
        }

        return tokens.IdToken;
    }

    private Task<OidcProviderMetadata> GetMetadataAsync(CancellationToken cancellationToken) =>
        discovery.GetAsync(options.Authority, options.RequireHttpsMetadata, cancellationToken);

    /// <summary>Opens the browser and waits for a matching callback.</summary>
    private async Task<(string Code, Uri RedirectUri, PkcePair Pkce)> AuthorizeAsync(
        OidcProviderMetadata metadata,
        IReadOnlyList<string> scopes,
        IReadOnlyDictionary<string, string>? extraParameters,
        CancellationToken cancellationToken)
    {
        if (metadata.CodeChallengeMethodsSupported.Count > 0
            && !metadata.CodeChallengeMethodsSupported.Contains(PkcePair.Method, StringComparer.Ordinal))
        {
            // Continuing without PKCE is not an option for a public client, so this is fatal rather
            // than a downgrade.
            throw new OidcException(
                $"The provider does not advertise the {PkcePair.Method} code challenge method.");
        }

        var pkce = PkcePair.Create();
        var state = Base64Url.EncodeToString(RandomNumberGenerator.GetBytes(32));

        using var listener = new LoopbackCallbackListener(options.RedirectPath);
        var redirectUri = listener.RedirectUri;

        var authorizeUri = BuildAuthorizeUri(
            metadata, options.ClientId, scopes, pkce, state, redirectUri, extraParameters);

        await browser.OpenAsync(authorizeUri, cancellationToken).ConfigureAwait(false);

        // Constructed with the TimeProvider rather than CancelAfter, so a test can advance a fake
        // clock instead of waiting out a five-minute browser timeout.
        using var timeout = new CancellationTokenSource(options.BrowserTimeout, clock);
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(
            cancellationToken,
            timeout.Token);

        CallbackResult callback;
        try
        {
            callback = await listener
                .WaitForCallbackAsync(options.CompletionHtml, linked.Token)
                .ConfigureAwait(false);
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested)
        {
            throw new OidcException(
                $"No authorization response arrived within {options.BrowserTimeout}.");
        }

        return (ExtractCode(callback, state), redirectUri, pkce);
    }

    /// <summary>Checks a callback and returns its authorization code.</summary>
    private static string ExtractCode(CallbackResult callback, string expectedState)
    {
        if (callback.Error is { } error)
        {
            throw new OidcException(
                $"The provider refused authorization: {callback.ErrorDescription ?? error}.",
                error);
        }

        // Checked before the code is touched. Constant-time comparison is unnecessary — state is not
        // a secret being guessed byte by byte — but it must be exact and ordinal.
        if (!string.Equals(callback.State, expectedState, StringComparison.Ordinal))
        {
            throw new OidcException(
                "The authorization response carried the wrong state and was discarded. Any local "
                + "process can reach the loopback port, so an unmatched response is treated as hostile.");
        }

        if (string.IsNullOrEmpty(callback.Code))
        {
            throw new OidcException("The authorization response carried no code.");
        }

        return callback.Code;
    }

    /// <remarks>
    /// Extra parameters are applied last and may override the defaults. That is what lets the key
    /// binding flow supply its own <c>nonce</c> and <c>prompt</c>, which is the reason this is
    /// hand-rolled at all.
    /// </remarks>
    private static Uri BuildAuthorizeUri(
        OidcProviderMetadata metadata,
        string clientId,
        IReadOnlyList<string> scopes,
        PkcePair pkce,
        string state,
        Uri redirectUri,
        IReadOnlyDictionary<string, string>? extraParameters)
    {
        var parameters = new Dictionary<string, string>(StringComparer.Ordinal)
        {
            ["response_type"] = "code",
            ["client_id"] = clientId,
            ["redirect_uri"] = redirectUri.AbsoluteUri,
            ["scope"] = string.Join(' ', scopes),
            ["state"] = state,
            ["code_challenge"] = pkce.CodeChallenge,
            ["code_challenge_method"] = PkcePair.Method,
        };

        if (extraParameters is not null)
        {
            foreach (var (name, value) in extraParameters)
            {
                parameters[name] = value;
            }
        }

        // The endpoint may already carry a query — Keycloak's does not, but some providers pin a
        // tenant or an audience there, and clobbering it would break them.
        var uri = new StringBuilder(metadata.AuthorizationEndpoint.AbsoluteUri);
        uri.Append(metadata.AuthorizationEndpoint.Query.Length > 0 ? '&' : '?');

        var first = true;
        foreach (var (name, value) in parameters)
        {
            if (!first)
            {
                uri.Append('&');
            }

            first = false;
            uri.Append(Uri.EscapeDataString(name)).Append('=').Append(Uri.EscapeDataString(value));
        }

        return new Uri(uri.ToString(), UriKind.Absolute);
    }

    private async Task<TokenSet> ExchangeAsync(
        OidcProviderMetadata metadata,
        Dictionary<string, string> form,
        CancellationToken cancellationToken)
    {
        using var content = new FormUrlEncodedContent(form);
        using var response = await http
            .PostAsync(metadata.TokenEndpoint, content, cancellationToken)
            .ConfigureAwait(false);

        var body = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);

        using var document = ParseOrThrow(body, response.StatusCode);
        var root = document.RootElement;

        if (!response.IsSuccessStatusCode)
        {
            var error = ReadString(root, "error") ?? "unknown_error";

            throw new OidcException(
                $"The token endpoint returned {(int)response.StatusCode}: "
                + $"{ReadString(root, "error_description") ?? error}.",
                error);
        }

        var accessToken = ReadString(root, "access_token")
            ?? throw new OidcException("The token response contained no access token.");

        // Absent expires_in is legal and means unspecified. Treating that as "never expires" would
        // hand the caller a token it never refreshes; a short assumed lifetime degrades to an extra
        // refresh instead.
        var lifetime = root.TryGetProperty("expires_in", out var expires)
            && expires.ValueKind == JsonValueKind.Number
            && expires.TryGetInt64(out var seconds)
                ? TimeSpan.FromSeconds(seconds)
                : TimeSpan.FromMinutes(5);

        return new TokenSet(
            AccessToken: accessToken,
            RefreshToken: ReadString(root, "refresh_token"),
            IdToken: ReadString(root, "id_token"),
            ExpiresAtUtc: clock.GetUtcNow() + lifetime,
            Scope: ReadString(root, "scope"));
    }

    private static JsonDocument ParseOrThrow(string body, System.Net.HttpStatusCode status)
    {
        try
        {
            return JsonDocument.Parse(body);
        }
        catch (JsonException exception)
        {
            throw new OidcException(
                $"The token endpoint returned {(int)status} with a body that is not JSON: "
                + exception.Message);
        }
    }

    private static string? ReadString(JsonElement root, string name) =>
        root.TryGetProperty(name, out var property) && property.ValueKind == JsonValueKind.String
            ? property.GetString()
            : null;

    /// <summary>
    /// Reads the <c>nonce</c> claim out of an ID token without verifying its signature.
    /// </summary>
    /// <remarks>
    /// Sanctioned by OpenID Connect Core §3.1.3.7: when an ID token is received by direct
    /// communication with the token endpoint, TLS server authentication may stand in for checking the
    /// token signature. It is the channel that establishes who sent this, so the nonce check here is
    /// about detecting a provider bug or a mixed-up response, not about trusting an unsigned token.
    /// <para>
    /// This reasoning does <b>not</b> extend to another user's key binding, which arrives through the
    /// DodoSSH server and must be verified against the provider's JWKS fetched directly. That is the
    /// directory work in M3.
    /// </para>
    /// </remarks>
    private static string? ReadNonceClaim(string idToken)
    {
        var segments = idToken.Split('.');
        if (segments.Length < 2)
        {
            throw new OidcException("The ID token is not a JWT.");
        }

        byte[] payload;
        try
        {
            payload = Base64Url.DecodeFromChars(segments[1]);
        }
        catch (FormatException)
        {
            throw new OidcException("The ID token's payload is not valid base64url.");
        }

        try
        {
            using var document = JsonDocument.Parse(payload);
            return ReadString(document.RootElement, "nonce");
        }
        catch (JsonException)
        {
            throw new OidcException("The ID token's payload is not JSON.");
        }
    }
}