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


var builder = WebApplication.CreateBuilder(args);

// Register HttpClient so we can proxy HTTP traffic
builder.Services.AddHttpClient();

var app = builder.Build();
var port8080 = builder.Configuration.GetValue<int>("ProxyConfig:Port8080", 8080);
var port4000 = builder.Configuration.GetValue<int>("ProxyConfig:Port4000", 4000);

// Enable WebSocket support
app.UseWebSockets();

// Forwarded headers (important if behind IIS or another proxy)
app.UseForwardedHeaders(new ForwardedHeadersOptions
{
    ForwardedHeaders = ForwardedHeaders.XForwardedFor | ForwardedHeaders.XForwardedProto,
    ForwardLimit = 2
});

// Simple file logger
// 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:u}: {message}{Environment.NewLine}");
//     }
//     catch (Exception ex)
//     {
//         Console.WriteLine($"Logging error: {ex.Message}");
//     }
// }

app.Use(async (context, next) =>
{
    if (context.Request.Headers.ContainsKey("Origin"))
    {
        var origin = context.Request.Headers["Origin"].ToString();
        context.Response.Headers["Access-Control-Allow-Origin"] = origin;
        context.Response.Headers["Vary"] = "Origin";
        context.Response.Headers["Access-Control-Allow-Credentials"] = "true";
        context.Response.Headers["Access-Control-Allow-Methods"] = "GET,POST,PUT,PATCH,DELETE,OPTIONS";
        context.Response.Headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization, X-Requested-With";
    }

    if (string.Equals(context.Request.Method, "OPTIONS", StringComparison.OrdinalIgnoreCase))
    {
        context.Response.StatusCode = StatusCodes.Status204NoContent;
        return;
    }

    await next();
});

// Single terminal middleware (Run instead of Use → no ambiguity)
app.Run(async (HttpContext context) =>
{
        var rawPath = context.Request.Path.Value ?? "";

    string backendHttpBase;
    string backendWsBase;

    if (rawPath.StartsWith("/socket.io", StringComparison.OrdinalIgnoreCase) ||
        rawPath.StartsWith("/lst/socket.io", StringComparison.OrdinalIgnoreCase))
    {
        backendHttpBase = $"http://localhost:{port8080}";
        backendWsBase   = $"ws://localhost:{port8080}";
    }
    else if (rawPath.StartsWith("/lst/api/controller", StringComparison.OrdinalIgnoreCase) || 
    rawPath.StartsWith("/api/controller", StringComparison.OrdinalIgnoreCase))
    {
        backendHttpBase = $"http://localhost:{port8080}";
        backendWsBase   = $"ws://localhost:{port8080}";

        // Now strip only once
        // var newPath = rawPath.Substring("/lst".Length);
        // context.Request.Path = new PathString(newPath);
    }
    else
    {
        backendHttpBase = $"http://localhost:{port4000}";
        backendWsBase   = $"ws://localhost:{port4000}";
    }

    // Handle WebSocket requests
    // if (context.WebSockets.IsWebSocketRequest)
    // {
    //     try
    //     {
    //         var backendUri = new UriBuilder(backendWsBase)
    //         {
    //             Path = context.Request.Path,
    //             Query = context.Request.QueryString.ToString()
    //         }.Uri;

    //         using var backendSocket = new ClientWebSocket();

    //         // Forward incoming headers
    //         foreach (var header in context.Request.Headers)
    //         {
    //             try { backendSocket.Options.SetRequestHeader(header.Key, header.Value); }
    //             catch { /* ignore headers WS client doesn't like */ }
    //         }

    //         await backendSocket.ConnectAsync(backendUri, context.RequestAborted);

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

    //         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}");
    //     }
    //     return;
    // }

if (context.WebSockets.IsWebSocketRequest)
{
    try
    {
        // Accept the incoming client WebSocket
        using var frontendSocket = await context.WebSockets.AcceptWebSocketAsync();

        // Build the backend WebSocket URI
        var backendUri = new UriBuilder(backendWsBase)
        {
            Path = context.Request.Path,
            Query = context.Request.QueryString.ToString()
        }.Uri;

        // Connect to the Node.js backend
        using var backendSocket = new ClientWebSocket();
        await backendSocket.ConnectAsync(backendUri, context.RequestAborted);

        // Forward frames both ways until one side closes
        var cts = new CancellationTokenSource();
        var forwardToBackend  = ForwardWebSocketAsync(frontendSocket, backendSocket, cts.Token);
        var forwardToFrontend = ForwardWebSocketAsync(backendSocket, frontendSocket, cts.Token);

        await Task.WhenAny(forwardToBackend, forwardToFrontend);

        // Cancel the other direction once one side closes
        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}");
    }

    return;
}
    // Otherwise: normal HTTP request
    var client = context.RequestServices.GetRequiredService<IHttpClientFactory>().CreateClient();
    var targetUri = backendHttpBase + context.Request.Path + context.Request.QueryString;

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

        // Copy headers
        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());
            }
        }

//Inject client IP into X-Forwarded-For
        var remoteIp = context.Connection.RemoteIpAddress?.ToString();
    if (!string.IsNullOrEmpty(remoteIp))
    {
        request.Headers.Remove("X-Forwarded-For");
        request.Headers.TryAddWithoutValidation("X-Forwarded-For", remoteIp);
    }

        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;
       // Copy backend headers
        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");

        // ✅ NOW inject/override CORS
        if (context.Request.Headers.ContainsKey("Origin"))
        {
            var origin = context.Request.Headers["Origin"].ToString();
            context.Response.Headers["Access-Control-Allow-Origin"] = origin;
            context.Response.Headers["Vary"] = "Origin";
            context.Response.Headers["Access-Control-Allow-Credentials"] = "true";
            context.Response.Headers["Access-Control-Allow-Methods"] =
                "GET,POST,PUT,PATCH,DELETE,OPTIONS";
            context.Response.Headers["Access-Control-Allow-Headers"] =
                "Content-Type, Authorization, X-Requested-With";
        }

        //await response.Content.CopyToAsync(context.Response.Body);

        // changes to manage the SSE so we get it all at once
        var stream = await response.Content.ReadAsStreamAsync(context.RequestAborted);
        var buffer = new byte[8192];
        int bytesRead;

        while ((bytesRead = await stream.ReadAsync(buffer.AsMemory(0, buffer.Length), context.RequestAborted)) > 0)
        {
            await context.Response.Body.WriteAsync(buffer.AsMemory(0, bytesRead), context.RequestAborted);

            // 🔑 Force flush so chunks go straight to client
            await context.Response.Body.FlushAsync(context.RequestAborted);
        }
    }
    catch (HttpRequestException ex)
    {
        LogToFile($"HTTP proxy error: {ex.Message}");
        context.Response.StatusCode = (int)HttpStatusCode.BadGateway;
        await context.Response.WriteAsync($"Backend request failed: {ex.Message}");
    }
});

// Helper to forward WS frames in both directions
static async Task ForwardWebSocketAsync(WebSocket source, WebSocket destination, CancellationToken cancellationToken)
{
    var buffer = new byte[8192];
    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)
    {
        Console.WriteLine($"WebSocket forwarding error: {ex.Message}");
        try
        {
            await destination.CloseOutputAsync(WebSocketCloseStatus.InternalServerError, "Error", cancellationToken);
        }
        catch { }
    }
}

app.Run();