// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using System.Diagnostics;

namespace TypeAgent.AIClient;

/// <summary>
/// Extension methods for working with Http 
/// </summary>
public static class HttpEx
{
    internal static async Task<TResponse> GetJsonResponseAsync<TRequest, TResponse>(
        this HttpClient client,
        string endpoint,
        TRequest request,
        string? apiToken = null,
        RetrySettings? settings = null,
        CancellationToken cancellationToken = default
    )
    {
        settings ??= new RetrySettings();

        int retryCount = 0;
        while (true)
        {
            HttpRequestMessage? httpRequest = null;
            int pauseMs = -1;
            try
            {
                var requestMessage = Json.ToJsonMessage(request);

                httpRequest = new HttpRequestMessage(HttpMethod.Post, endpoint)
                {
                    Content = requestMessage
                };
                if (!string.IsNullOrEmpty(apiToken))
                {
                    httpRequest.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiToken);
                }

                try
                {
                    using HttpResponseMessage response = await client.SendAsync(httpRequest, cancellationToken).ConfigureAwait(false);
                    if (response.StatusCode == HttpStatusCode.OK)
                    {
                        using Stream stream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
                        return Json.Parse<TResponse>(stream);
                    }
                    if (!response.StatusCode.IsTransientError() || retryCount >= settings.MaxRetries)
                    {
                        if (Debugger.IsAttached)
                        {
                            string content = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);

                            throw new HttpRequestException($"GetJsonResponse failed with '{content}'", null, response.StatusCode);
                        }
                        else
                        {
                            // Let HttpClient throw an exception
                            response.EnsureSuccessStatusCode();
                        }
                    }
                    if (response.StatusCode == (HttpStatusCode)429) // Too Many Requests
                    {
                        pauseMs = GetRetryAfterMs(response, pauseMs);
                    }
                }
                catch (HttpRequestException ex)
                {
                    if (retryCount >= settings.MaxRetries)
                    {
                        Debug.WriteLine(ex);
                        throw;
                    }
                }
            }
            finally
            {
                httpRequest?.Dispose();
                httpRequest = null;
            }

            if (pauseMs < 0)
            {
                pauseMs = settings.RetryPauseMs > 0
                            ? settings.RetryPauseMs * (1 << (retryCount)) // Exponential backoff
                            : 0;

            }
            retryCount++;

            if (pauseMs > 0)
            {
                pauseMs = settings.AdjustRetryPauseMs(pauseMs);
                await Task.Delay(pauseMs, cancellationToken).ConfigureAwait(false);
            }

        }
        throw new TypeAgentException("GetJsonResponse");
    }

    internal static bool IsTransientError(this HttpStatusCode status)
    {
        switch (status)
        {
            default:
                return false;

            case (HttpStatusCode)429: // Too many requests
            case HttpStatusCode.InternalServerError:
            case HttpStatusCode.BadGateway:
            case HttpStatusCode.ServiceUnavailable:
            case HttpStatusCode.GatewayTimeout:
            case HttpStatusCode.RequestTimeout:
                break;
        }

        return true;
    }

    /// <summary>
    /// When servers return a 429, they can include a Retry-After header that says how long the caller
    /// should wait before retrying.
    /// </summary>
    /// <param name="response">The HttpResponseMessage to inspect.</param>
    /// <param name="defaultValue">Default pause in milliseconds if header is missing or invalid.</param>
    /// <returns>Milliseconds to pause before retrying.</returns>
    internal static int GetRetryAfterMs(HttpResponseMessage response, int defaultValue)
    {
        int pauseMs = defaultValue;
        try
        {
            if (response.Headers.TryGetValues("Retry-After", out var values))
            {
                var pauseHeader = values.FirstOrDefault()?.Trim();
                if (!string.IsNullOrEmpty(pauseHeader))
                {
                    if (int.TryParse(pauseHeader, out int seconds))
                    {
                        pauseMs = seconds * 1000;
                    }
                    else if (DateTimeOffset.TryParse(pauseHeader, out var retryDate))
                    {
                        pauseMs = (int)(retryDate - DateTimeOffset.UtcNow).TotalMilliseconds;
                    }
                    if (pauseMs <= 0)
                    {
                        pauseMs = defaultValue;
                    }
                }
            }
        }
        catch (Exception ex)
        {
            Debug.WriteLine($"Failed to parse Retry-After header: {ex}");
            pauseMs = defaultValue;
        }
        return pauseMs;
    }
}
