using System.Net;
using System.Security.Cryptography;
using System.Text;

namespace NATS.Server.WebSocket;

/// <summary>
/// WebSocket HTTP upgrade handshake handler.
/// Ported from golang/nats-server/server/websocket.go lines 731-917.
/// </summary>
public static class WsUpgrade
{
    public static async Task<WsUpgradeResult> TryUpgradeAsync(
        Stream inputStream, Stream outputStream, WebSocketOptions options,
        CancellationToken ct = default)
    {
        try
        {
            using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            cts.CancelAfter(options.HandshakeTimeout);
            var (method, path, queryString, headers) = await ReadHttpRequestAsync(inputStream, cts.Token);

            if (!string.Equals(method, "GET", StringComparison.OrdinalIgnoreCase))
                return await FailAsync(outputStream, 405, "request method must be GET");

            if (!headers.ContainsKey("Host"))
                return await FailAsync(outputStream, 400, "'Host' missing in request");

            if (!HeaderContains(headers, "Upgrade", "websocket"))
                return await FailAsync(outputStream, 400, "invalid value for header 'Upgrade'");

            if (!HeaderContains(headers, "Connection", "Upgrade"))
                return await FailAsync(outputStream, 400, "invalid value for header 'Connection'");

            if (!headers.TryGetValue("Sec-WebSocket-Key", out var key) || string.IsNullOrEmpty(key))
                return await FailAsync(outputStream, 400, "key missing");

            if (!HeaderContains(headers, "Sec-WebSocket-Version", "13"))
                return await FailAsync(outputStream, 400, "invalid version");

            var kind = path switch
            {
                _ when path.EndsWith("/leafnode") => WsClientKind.Leaf,
                _ when path.EndsWith("/mqtt") => WsClientKind.Mqtt,
                _ => WsClientKind.Client,
            };

            // Origin checking
            if (options.SameOrigin || options.AllowedOrigins is { Count: > 0 })
            {
                var checker = new WsOriginChecker(options.SameOrigin, options.AllowedOrigins);
                headers.TryGetValue("Origin", out var origin);
                if (string.IsNullOrEmpty(origin))
                    headers.TryGetValue("Sec-WebSocket-Origin", out origin);
                var originErr = checker.CheckOrigin(origin, headers.GetValueOrDefault("Host", ""), isTls: false);
                if (originErr != null)
                    return await FailAsync(outputStream, 403, $"origin not allowed: {originErr}");
            }

            // Compression negotiation (RFC 7692)
            bool compress = options.Compression;
            WsDeflateParams? deflateParams = null;
            if (compress)
            {
                headers.TryGetValue("Sec-WebSocket-Extensions", out var ext);
                deflateParams = WsDeflateNegotiator.Negotiate(ext);
                compress = deflateParams != null;
            }

            // No-masking support (leaf nodes only -- browser clients must always mask)
            bool noMasking = kind == WsClientKind.Leaf &&
                             headers.TryGetValue(WsConstants.NoMaskingHeader, out var nmVal) &&
                             string.Equals(nmVal.Trim(), WsConstants.NoMaskingValue, StringComparison.OrdinalIgnoreCase);

            // Browser detection
            bool browser = false;
            bool noCompFrag = false;
            if (kind is WsClientKind.Client or WsClientKind.Mqtt &&
                headers.TryGetValue("User-Agent", out var ua) && ua.StartsWith("Mozilla/"))
            {
                browser = true;
                // Disable fragmentation of compressed frames for Safari browsers.
                // Safari has both "Version/" and "Safari/" in the user agent string,
                // while Chrome on macOS has "Safari/" but not "Version/".
                noCompFrag = compress && ua.Contains("Version/") && ua.Contains("Safari/");
            }

            // Cookie extraction
            string? cookieJwt = null, cookieUsername = null, cookiePassword = null, cookieToken = null;
            if ((kind is WsClientKind.Client or WsClientKind.Mqtt) &&
                headers.TryGetValue("Cookie", out var cookieHeader))
            {
                var cookies = ParseCookies(cookieHeader);
                if (options.JwtCookie != null) cookies.TryGetValue(options.JwtCookie, out cookieJwt);
                if (options.UsernameCookie != null) cookies.TryGetValue(options.UsernameCookie, out cookieUsername);
                if (options.PasswordCookie != null) cookies.TryGetValue(options.PasswordCookie, out cookiePassword);
                if (options.TokenCookie != null) cookies.TryGetValue(options.TokenCookie, out cookieToken);
            }

            // JWT extraction from multiple sources (E11):
            // Priority: Authorization header > cookie > query parameter
            // Reference: NATS WebSocket JWT auth — browser clients often pass JWT
            // via cookie or query param since they cannot set custom headers.
            string? jwt = null;
            if (headers.TryGetValue("Authorization", out var authHeader))
            {
                jwt = ExtractBearerToken(authHeader);
            }

            jwt ??= cookieJwt;

            if (jwt == null && queryString != null)
            {
                var queryParams = ParseQueryString(queryString);
                queryParams.TryGetValue("jwt", out jwt);
            }

            // X-Forwarded-For client IP extraction
            string? clientIp = null;
            if (headers.TryGetValue(WsConstants.XForwardedForHeader, out var xff))
            {
                var ip = xff.Split(',')[0].Trim();
                if (IPAddress.TryParse(ip, out _))
                    clientIp = ip;
            }

            // Build the 101 Switching Protocols response
            var response = new StringBuilder();
            response.Append("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: ");
            response.Append(ComputeAcceptKey(key));
            response.Append("\r\n");
            if (compress && deflateParams != null)
            {
                response.Append("Sec-WebSocket-Extensions: ");
                response.Append(deflateParams.Value.ToResponseHeaderValue());
                response.Append("\r\n");
            }

            if (noMasking)
                response.Append(WsConstants.NoMaskingFullResponse);
            if (options.Headers != null)
            {
                foreach (var (k, v) in options.Headers)
                {
                    response.Append(k);
                    response.Append(": ");
                    response.Append(v);
                    response.Append("\r\n");
                }
            }

            response.Append("\r\n");

            var responseBytes = Encoding.ASCII.GetBytes(response.ToString());
            await outputStream.WriteAsync(responseBytes);
            await outputStream.FlushAsync();

            return new WsUpgradeResult(
                Success: true, Compress: compress, Browser: browser, NoCompFrag: noCompFrag,
                MaskRead: !noMasking, MaskWrite: false,
                CookieJwt: cookieJwt, CookieUsername: cookieUsername,
                CookiePassword: cookiePassword, CookieToken: cookieToken,
                ClientIp: clientIp, Kind: kind,
                DeflateParams: deflateParams, Jwt: jwt);
        }
        catch (Exception)
        {
            return WsUpgradeResult.Failed;
        }
    }

    /// <summary>
    /// Computes the Sec-WebSocket-Accept value per RFC 6455 Section 4.2.2.
    /// </summary>
    public static string ComputeAcceptKey(string clientKey)
    {
        var combined = Encoding.ASCII.GetBytes(clientKey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
        var hash = SHA1.HashData(combined);
        return Convert.ToBase64String(hash);
    }

    /// <summary>
    /// Extracts a bearer token from an Authorization header value.
    /// Supports both "Bearer {token}" and bare "{token}" formats.
    /// </summary>
    internal static string? ExtractBearerToken(string? authHeader)
    {
        if (string.IsNullOrWhiteSpace(authHeader))
            return null;

        var trimmed = authHeader.Trim();
        if (trimmed.StartsWith("Bearer ", StringComparison.OrdinalIgnoreCase))
            return trimmed["Bearer ".Length..].Trim();

        // Some clients send the token directly without "Bearer" prefix
        return trimmed;
    }

    /// <summary>
    /// Parses a query string into key-value pairs.
    /// </summary>
    internal static Dictionary<string, string> ParseQueryString(string queryString)
    {
        var result = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
        if (queryString.StartsWith('?'))
            queryString = queryString[1..];

        foreach (var pair in queryString.Split('&'))
        {
            var eqIdx = pair.IndexOf('=');
            if (eqIdx > 0)
            {
                var name = Uri.UnescapeDataString(pair[..eqIdx]);
                var value = Uri.UnescapeDataString(pair[(eqIdx + 1)..]);
                result[name] = value;
            }
            else if (pair.Length > 0)
            {
                result[Uri.UnescapeDataString(pair)] = string.Empty;
            }
        }

        return result;
    }

    private static async Task<WsUpgradeResult> FailAsync(Stream output, int statusCode, string reason)
    {
        var statusText = statusCode switch
        {
            400 => "Bad Request",
            401 => "Unauthorized",
            403 => "Forbidden",
            405 => "Method Not Allowed",
            _ => "Internal Server Error",
        };
        var response = $"HTTP/1.1 {statusCode} {statusText}\r\nSec-WebSocket-Version: 13\r\nContent-Type: text/plain\r\nContent-Length: {reason.Length}\r\n\r\n{reason}";
        await output.WriteAsync(Encoding.ASCII.GetBytes(response));
        await output.FlushAsync();
        return statusCode == 401
            ? WsUpgradeResult.Unauthorized
            : WsUpgradeResult.Failed;
    }

    /// <summary>
    /// Sends a 401 Unauthorized response and returns a failed upgrade result.
    /// Used by the server when JWT authentication fails during WS upgrade.
    /// </summary>
    public static async Task<WsUpgradeResult> FailUnauthorizedAsync(Stream output, string reason)
    {
        return await FailAsync(output, 401, reason);
    }

    private static async Task<(string method, string path, string? queryString, Dictionary<string, string> headers)> ReadHttpRequestAsync(
        Stream stream, CancellationToken ct)
    {
        var headerBytes = new List<byte>(4096);
        var buf = new byte[512];
        while (true)
        {
            int n = await stream.ReadAsync(buf, ct);
            if (n == 0) throw new IOException("connection closed during handshake");
            for (int i = 0; i < n; i++)
            {
                headerBytes.Add(buf[i]);
                if (headerBytes.Count >= 4 &&
                    headerBytes[^4] == '\r' && headerBytes[^3] == '\n' &&
                    headerBytes[^2] == '\r' && headerBytes[^1] == '\n')
                    goto done;
                if (headerBytes.Count > 8192)
                    throw new InvalidOperationException("HTTP header too large");
            }
        }
        done:;

        var text = Encoding.ASCII.GetString(headerBytes.ToArray());
        var lines = text.Split("\r\n", StringSplitOptions.None);
        if (lines.Length < 1) throw new InvalidOperationException("invalid HTTP request");

        var parts = lines[0].Split(' ');
        if (parts.Length < 3) throw new InvalidOperationException("invalid HTTP request line");
        var method = parts[0];
        var requestUri = parts[1];

        // Split path and query string
        string path;
        string? queryString = null;
        var qIdx = requestUri.IndexOf('?');
        if (qIdx >= 0)
        {
            path = requestUri[..qIdx];
            queryString = requestUri[qIdx..]; // includes the '?'
        }
        else
        {
            path = requestUri;
        }

        var headers = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
        for (int i = 1; i < lines.Length; i++)
        {
            var line = lines[i];
            if (string.IsNullOrEmpty(line)) break;
            var colonIdx = line.IndexOf(':');
            if (colonIdx > 0)
            {
                var name = line[..colonIdx].Trim();
                var value = line[(colonIdx + 1)..].Trim();
                headers[name] = value;
            }
        }

        return (method, path, queryString, headers);
    }

    private static bool HeaderContains(Dictionary<string, string> headers, string name, string value)
    {
        if (!headers.TryGetValue(name, out var headerValue))
            return false;
        foreach (var token in headerValue.Split(','))
        {
            if (string.Equals(token.Trim(), value, StringComparison.OrdinalIgnoreCase))
                return true;
        }

        return false;
    }

    private static Dictionary<string, string> ParseCookies(string cookieHeader)
    {
        var cookies = new Dictionary<string, string>(StringComparer.Ordinal);
        foreach (var pair in cookieHeader.Split(';'))
        {
            var trimmed = pair.Trim();
            var eqIdx = trimmed.IndexOf('=');
            if (eqIdx > 0)
                cookies[trimmed[..eqIdx].Trim()] = trimmed[(eqIdx + 1)..].Trim();
        }

        return cookies;
    }
}

/// <summary>
/// Result of a WebSocket upgrade handshake attempt.
/// </summary>
public readonly record struct WsUpgradeResult(
    bool Success,
    bool Compress,
    bool Browser,
    bool NoCompFrag,
    bool MaskRead,
    bool MaskWrite,
    string? CookieJwt,
    string? CookieUsername,
    string? CookiePassword,
    string? CookieToken,
    string? ClientIp,
    WsClientKind Kind,
    WsDeflateParams? DeflateParams = null,
    string? Jwt = null)
{
    public static readonly WsUpgradeResult Failed = new(
        Success: false, Compress: false, Browser: false, NoCompFrag: false,
        MaskRead: true, MaskWrite: false, CookieJwt: null, CookieUsername: null,
        CookiePassword: null, CookieToken: null, ClientIp: null, Kind: WsClientKind.Client);

    public static readonly WsUpgradeResult Unauthorized = new(
        Success: false, Compress: false, Browser: false, NoCompFrag: false,
        MaskRead: true, MaskWrite: false, CookieJwt: null, CookieUsername: null,
        CookiePassword: null, CookieToken: null, ClientIp: null, Kind: WsClientKind.Client);
}
