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;
    }

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