using System;
using System.IO;
using System.Net;
using System.Net.WebSockets;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Builder;
using Microsoft.Extensions.DependencyInjection;

var builder = WebApplication.CreateBuilder(args);
builder.Services.AddHttpClient("GoBackend", client =>
{
    client.BaseAddress = new Uri("http://localhost:8080");
});

var app = builder.Build();

// Enable WebSocket support
app.UseWebSockets();

// Logging method
void LogToFile(string message)
{
    try
    {
        string logDir = Path.Combine(AppContext.BaseDirectory, "logs");
        Directory.CreateDirectory(logDir);
        string logFilePath = Path.Combine(logDir, "proxy_log.txt");
        File.AppendAllText(logFilePath, $"{DateTime.UtcNow}: {message}{Environment.NewLine}");
    }
    catch (Exception ex)
    {
        // Handle potential errors writing to log file
        Console.WriteLine($"Logging error: {ex.Message}");
    }
}

// Middleware to handle WebSocket requests
app.Use(async (context, next) =>
{
    if (context.WebSockets.IsWebSocketRequest && context.Request.Path.StartsWithSegments("/ws"))
    {
        LogToFile($"WebSocket request received for path: {context.Request.Path}");

        try
        {
            var backendUri = new UriBuilder("ws", "localhost", 8080)
            {
                Path = context.Request.Path,
                Query = context.Request.QueryString.ToString()
            }.Uri;

            using var backendSocket = new ClientWebSocket();
            await backendSocket.ConnectAsync(backendUri, context.RequestAborted);

            using var frontendSocket = await context.WebSockets.AcceptWebSocketAsync();
            var cts = new CancellationTokenSource();

            // WebSocket forwarding tasks
            var forwardToBackend = ForwardWebSocketAsync(frontendSocket, backendSocket, cts.Token);
            var forwardToFrontend = ForwardWebSocketAsync(backendSocket, frontendSocket, cts.Token);

            await Task.WhenAny(forwardToBackend, forwardToFrontend);
            cts.Cancel();
        }
        catch (Exception ex)
        {
            LogToFile($"WebSocket proxy error: {ex.Message}");
            context.Response.StatusCode = (int)HttpStatusCode.BadGateway;
            await context.Response.WriteAsync($"WebSocket proxy error: {ex.Message}");
        }
    }
    else
    {
        await next();
    }
});

// Middleware to handle HTTP requests
app.Use(async (context, next) =>
{
    if (context.WebSockets.IsWebSocketRequest)
    {
        await next();
        return;
    }

    var client = context.RequestServices.GetRequiredService<IHttpClientFactory>().CreateClient("GoBackend");

    try
    {
        var request = new HttpRequestMessage(new HttpMethod(context.Request.Method),
            context.Request.Path + context.Request.QueryString);

        foreach (var header in context.Request.Headers)
        {
            if (!request.Headers.TryAddWithoutValidation(header.Key, header.Value.ToArray()))
            {
                request.Content ??= new StreamContent(context.Request.Body);
                request.Content.Headers.TryAddWithoutValidation(header.Key, header.Value.ToArray());
            }
        }

        if (context.Request.ContentLength > 0 && request.Content == null)
        {
            request.Content = new StreamContent(context.Request.Body);
        }

        var response = await client.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, context.RequestAborted);
        context.Response.StatusCode = (int)response.StatusCode;

        foreach (var header in response.Headers)
        {
            context.Response.Headers[header.Key] = header.Value.ToArray();
        }

        foreach (var header in response.Content.Headers)
        {
            context.Response.Headers[header.Key] = header.Value.ToArray();
        }

        context.Response.Headers.Remove("transfer-encoding");
        await response.Content.CopyToAsync(context.Response.Body);
    }
    catch (HttpRequestException ex)
    {
        LogToFile($"HTTP proxy error: {ex.Message}");
        context.Response.StatusCode = (int)HttpStatusCode.BadGateway;
        await context.Response.WriteAsync($"Backend request failed: {ex.Message}");
    }
});

async Task ForwardWebSocketAsync(WebSocket source, WebSocket destination, CancellationToken cancellationToken)
{
    var buffer = new byte[4 * 1024];
    try
    {
        while (source.State == WebSocketState.Open &&
               destination.State == WebSocketState.Open &&
               !cancellationToken.IsCancellationRequested)
        {
            var result = await source.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                await destination.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "Closing", cancellationToken);
                break;
            }
            await destination.SendAsync(new ArraySegment<byte>(buffer, 0, result.Count), result.MessageType, result.EndOfMessage, cancellationToken);
        }
    }
    catch (WebSocketException ex)
    {
        LogToFile($"WebSocket forwarding error: {ex.Message}");
        await destination.CloseOutputAsync(WebSocketCloseStatus.InternalServerError, "Error", cancellationToken);
    }
}

app.Run();