module veb

import compress.gzip
import compress.zstd
import net.http

pub type MiddlewareHandler[T] = fn (mut T) bool

interface MiddlewareApp {
    get_handlers_for_route(route_path string) []RouteMiddleware
    get_handlers_for_route_after(route_path string) []RouteMiddleware
    get_global_handlers() []voidptr
    get_global_handlers_after() []voidptr
}

struct RouteMiddleware {
    url_parts []string
    methods   []http.Method
    handler   voidptr
}

pub struct Middleware[T] {
mut:
    global_handlers       []voidptr
    global_handlers_after []voidptr
    route_handlers        []RouteMiddleware
    route_handlers_after  []RouteMiddleware
}

@[params]
pub struct MiddlewareOptions[T] {
pub:
    handler fn (mut ctx T) bool @[required]
    after   bool
    methods []http.Method
}

// string representation of Middleware
pub fn (m &Middleware[T]) str() string {
    return 'veb.Middleware[${T.name}]{
        global_handlers: [${m.global_handlers.len}]
        global_handlers_after: [${m.global_handlers_after.len}]
        route_handlers: [${m.route_handlers.len}]
        route_handlers_after: [${m.route_handlers_after.len}]
    }'
}

// use registers a global middleware handler
pub fn (mut m Middleware[T]) use(options MiddlewareOptions[T]) {
    if options.after {
        m.global_handlers_after << voidptr(options.handler)
    } else {
        m.global_handlers << voidptr(options.handler)
    }
}

// route_use registers a middleware handler for a specific route(s)
pub fn (mut m Middleware[T]) route_use(route string, options MiddlewareOptions[T]) {
    middleware := RouteMiddleware{
        url_parts: route.split('/').filter(it != '')
        methods:   options.methods.clone()
        handler:   voidptr(options.handler)
    }

    if options.after {
        m.route_handlers_after << middleware
    } else {
        m.route_handlers << middleware
    }
}

fn (m &Middleware[T]) get_handlers_for_route(route_path string) []RouteMiddleware {
    mut handlers := []RouteMiddleware{}
    route_parts := route_path.split('/').filter(it != '')

    for handler in m.route_handlers {
        if _ := route_matches(route_parts, handler.url_parts) {
            handlers << handler
        } else if handler.url_parts.len == 0 && route_path == '/index' {
            handlers << handler
        }
    }

    return handlers
}

fn (m &Middleware[T]) get_handlers_for_route_after(route_path string) []RouteMiddleware {
    mut handlers := []RouteMiddleware{}
    route_parts := route_path.split('/').filter(it != '')

    for handler in m.route_handlers_after {
        if _ := route_matches(route_parts, handler.url_parts) {
            handlers << handler
        } else if handler.url_parts.len == 0 && route_path == '/index' {
            handlers << handler
        }
    }

    return handlers
}

fn (m &Middleware[T]) get_global_handlers() []voidptr {
    return m.global_handlers
}

fn (m &Middleware[T]) get_global_handlers_after() []voidptr {
    return m.global_handlers_after
}

fn app_route_handlers[A](app &A, route_path string) []RouteMiddleware {
    $if A is $struct {
        $for field in A.fields {
            $if field.is_embed {
                $if field.name == 'Middleware' {
                    return app.$(field.name).get_handlers_for_route(route_path)
                } $else $if field.typ is $struct {
                    handlers := app_route_handlers(app.$(field.name), route_path)
                    if handlers.len > 0 {
                        return handlers
                    }
                }
            }
        }
    }
    return []RouteMiddleware{}
}

fn app_route_handlers_after[A](app &A, route_path string) []RouteMiddleware {
    $if A is $struct {
        $for field in A.fields {
            $if field.is_embed {
                $if field.name == 'Middleware' {
                    return app.$(field.name).get_handlers_for_route_after(route_path)
                } $else $if field.typ is $struct {
                    handlers := app_route_handlers_after(app.$(field.name), route_path)
                    if handlers.len > 0 {
                        return handlers
                    }
                }
            }
        }
    }
    return []RouteMiddleware{}
}

fn app_global_handlers[A](app &A) []voidptr {
    $if A is $struct {
        $for field in A.fields {
            $if field.is_embed {
                $if field.name == 'Middleware' {
                    return app.$(field.name).get_global_handlers()
                } $else $if field.typ is $struct {
                    handlers := app_global_handlers(app.$(field.name))
                    if handlers.len > 0 {
                        return handlers
                    }
                }
            }
        }
    }
    return []voidptr{}
}

fn app_global_handlers_after[A](app &A) []voidptr {
    $if A is $struct {
        $for field in A.fields {
            $if field.is_embed {
                $if field.name == 'Middleware' {
                    return app.$(field.name).get_global_handlers_after()
                } $else $if field.typ is $struct {
                    handlers := app_global_handlers_after(app.$(field.name))
                    if handlers.len > 0 {
                        return handlers
                    }
                }
            }
        }
    }
    return []voidptr{}
}

fn validate_middleware[T](mut ctx T, raw_handlers []voidptr) bool {
    for handler in raw_handlers {
        func := MiddlewareHandler[T](handler)
        if func(mut ctx) == false {
            return false
        }
    }

    return true
}

fn route_middleware_matches_method(route_middleware RouteMiddleware, request_method http.Method) bool {
    return route_middleware.methods.len == 0 || request_method in route_middleware.methods
}

fn get_handlers_for_method(route_middlewares []RouteMiddleware, request_method http.Method) []voidptr {
    mut handlers := []voidptr{}
    for route_middleware in route_middlewares {
        if route_middleware_matches_method(route_middleware, request_method) {
            handlers << route_middleware.handler
        }
    }
    return handlers
}

// Compression encoding types for HTTP responses
enum ContentEncoding {
    gzip
    zstd
}

// send_compressed_response compresses the response body and updates the response.
// Returns true if compression should be skipped, false if compression was applied.
fn send_compressed_response(mut ctx Context, encoding ContentEncoding) bool {
    compressed, encoding_name := match encoding {
        .zstd {
            data := zstd.compress(ctx.res.body.bytes()) or {
                eprintln('[veb] error while compressing with zstd: ${err.msg()}')
                return true
            }
            data, 'zstd'
        }
        .gzip {
            data := gzip.compress(ctx.res.body.bytes()) or {
                eprintln('[veb] error while compressing with gzip: ${err.msg()}')
                return true
            }
            data, 'gzip'
        }
    }

    // Set HTTP headers for compressed content
    ctx.res.header.add(.content_encoding, encoding_name)
    ctx.res.header.set(.vary, 'Accept-Encoding')

    // Replace the response body with the compressed data and update Content-Length.
    // The normal response path will handle sending it.
    ctx.res.body = compressed.bytestr()
    ctx.res.header.set(.content_length, compressed.len.str())
    ctx.already_compressed = true

    return false
}

// should_skip_compression checks if compression should be skipped for this context.
fn should_skip_compression(ctx Context) bool {
    // Skip if already compressed (optimization for static files compressed in send_file)
    if ctx.already_compressed {
        return true
    }
    // Skip compression for files in streaming mode (no takeover)
    // Files in takeover mode (small files loaded in memory) are compressed
    if ctx.return_type == .file && ctx.takeover_mode == .none {
        return true
    }
    return false
}

// encode_gzip adds gzip encoding to the HTTP Response body.
// This middleware compresses dynamic routes and static files loaded in memory (takeover mode).
// Static files in streaming mode are compressed by send_file() when static compression is enabled,
// and this middleware skips them to avoid double compression (via the already_compressed flag).
// Register this middleware as last!
// Usage example: app.use(veb.encode_gzip[Context]())
pub fn encode_gzip[T]() MiddlewareOptions[T] {
    return MiddlewareOptions[T]{
        after:   true
        handler: fn [T](mut ctx T) bool {
            if should_skip_compression(ctx.Context) {
                return true
            }
            return send_compressed_response(mut ctx.Context, .gzip)
        }
    }
}

// encode_zstd adds zstd encoding to the HTTP Response body.
// This middleware compresses dynamic routes and static files loaded in memory (takeover mode).
// Static files in streaming mode are compressed by send_file() when static compression is enabled,
// and this middleware skips them to avoid double compression (via the already_compressed flag).
// Register this middleware as last!
// Usage example: app.route_use('/api', veb.encode_zstd[Context]())
pub fn encode_zstd[T]() MiddlewareOptions[T] {
    return MiddlewareOptions[T]{
        after:   true
        handler: fn [T](mut ctx T) bool {
            if should_skip_compression(ctx.Context) {
                return true
            }
            return send_compressed_response(mut ctx.Context, .zstd)
        }
    }
}

// encode_auto adds automatic content encoding (zstd or gzip) based on the client's Accept-Encoding header.
// This middleware checks the Accept-Encoding header and compresses with zstd if supported, otherwise gzip.
// Static files in streaming mode are compressed by send_file() when static compression is enabled,
// and this middleware skips them to avoid double compression (via the already_compressed flag).
// Register this middleware as last!
// Usage example: app.use(veb.encode_auto[Context]())
pub fn encode_auto[T]() MiddlewareOptions[T] {
    return MiddlewareOptions[T]{
        after:   true
        handler: fn [T](mut ctx T) bool {
            if should_skip_compression(ctx.Context) {
                return true
            }

            // Check Accept-Encoding header to determine best compression
            accept_encoding := ctx.req.header.get(.accept_encoding) or { '' }
            supports_zstd := accept_encoding.contains('zstd')
            supports_gzip := accept_encoding.contains('gzip')

            // Try zstd first (better compression ratio), fallback to gzip
            if supports_zstd {
                return send_compressed_response(mut ctx.Context, .zstd)
            }
            if supports_gzip {
                return send_compressed_response(mut ctx.Context, .gzip)
            }

            // No supported compression
            return true
        }
    }
}

// decode_gzip decodes the body of a gzip'ed HTTP request.
// Register this middleware before you do anything with the request body!
// Usage example: app.use(veb.decode_gzip[Context]())
pub fn decode_gzip[T]() MiddlewareOptions[T] {
    return MiddlewareOptions[T]{
        handler: fn [T](mut ctx T) bool {
            if encoding := ctx.req.header.get(.content_encoding) {
                if encoding == 'gzip' {
                    decompressed := gzip.decompress(ctx.req.data.bytes()) or {
                        ctx.request_error('invalid gzip encoding')
                        return false
                    }
                    ctx.req.data = decompressed.bytestr()
                }
            }
            return true
        }
    }
}

// decode_zstd decodes the body of a zstd-compressed HTTP request.
// Register this middleware before you do anything with the request body!
// Usage example: app.use(veb.decode_zstd[Context]())
pub fn decode_zstd[T]() MiddlewareOptions[T] {
    return MiddlewareOptions[T]{
        handler: fn [T](mut ctx T) bool {
            if encoding := ctx.req.header.get(.content_encoding) {
                if encoding == 'zstd' {
                    decompressed := zstd.decompress(ctx.req.data.bytes()) or {
                        ctx.request_error('invalid zstd encoding')
                        return false
                    }
                    ctx.req.data = decompressed.bytestr()
                }
            }
            return true
        }
    }
}

pub const cors_safelisted_response_headers = [http.CommonHeader.cache_control, .content_language,
    .content_length, .content_type, .expires, .last_modified, .pragma].map(it.str()).join(',')

// CorsOptions is used to set CORS response headers.
// See https://developer.mozilla.org/en-US/docs/Web/HTTP/CORS#the_http_response_headers
@[params]
pub struct CorsOptions {
pub:
    // from which origin(s) can cross-origin requests be made; `Access-Control-Allow-Origin`
    origins []string @[required]
    // indicate whether the server allows credentials, e.g. cookies, in cross-origin requests.
    // ;`Access-Control-Allow-Credentials`
    allow_credentials bool
    // allowed HTTP headers for a cross-origin request; `Access-Control-Allow-Headers`
    allowed_headers []string
    // allowed HTTP methods for a cross-origin request; `Access-Control-Allow-Methods`
    allowed_methods []http.Method
    // indicate if clients are able to access other headers than the "CORS-safelisted"
    // response headers; `Access-Control-Expose-Headers`
    expose_headers []string
    // how long the results of a preflight request can be cached, value is in seconds
    // ; `Access-Control-Max-Age`
    max_age ?int
}

// set_headers adds the CORS headers on the response
pub fn (options &CorsOptions) set_headers(mut ctx Context) {
    // A browser will reject a CORS request when the Access-Control-Allow-Origin header
    // is not present. By not setting the CORS headers when an invalid origin is supplied
    // we force the browser to reject the preflight and the actual request.
    origin := ctx.req.header.get(.origin) or { return }
    if options.origins != ['*'] && origin !in options.origins {
        return
    }

    ctx.set_header(.access_control_allow_origin, origin)
    ctx.set_header(.vary, 'Origin, Access-Control-Request-Headers')

    // dont' set the value of `Access-Control-Allow-Credentials` to 'false', but
    // omit the header if the value is `false`
    if options.allow_credentials {
        ctx.set_header(.access_control_allow_credentials, 'true')
    }

    if options.allowed_headers.len > 0 {
        ctx.set_header(.access_control_allow_headers, options.allowed_headers.join(','))
    } else if _ := ctx.req.header.get(.access_control_request_headers) {
        // a server must respond with `Access-Control-Allow-Headers` if
        // `Access-Control-Request-Headers` is present in a preflight request
        ctx.set_header(.access_control_allow_headers, cors_safelisted_response_headers)
    }

    if options.allowed_methods.len > 0 {
        method_str := options.allowed_methods.str().trim('[]')
        ctx.set_header(.access_control_allow_methods, method_str)
    }

    if options.expose_headers.len > 0 {
        ctx.set_header(.access_control_expose_headers, options.expose_headers.join(','))
    }

    if max_age := options.max_age {
        ctx.set_header(.access_control_max_age, max_age.str())
    }
}

// validate_request checks if a cross-origin request is made and verifies the CORS
// headers. If a cross-origin request is invalid this method will send a response
// using `ctx`.
pub fn (options &CorsOptions) validate_request(mut ctx Context) bool {
    origin := ctx.req.header.get(.origin) or { return true }
    if options.origins != ['*'] && origin !in options.origins {
        ctx.res.set_status(.forbidden)
        ctx.text('invalid CORS origin')

        $if veb_trace_cors ? {
            eprintln('[veb]: rejected CORS request from "${origin}". Reason: invalid origin')
        }
        return false
    }

    ctx.set_header(.access_control_allow_origin, origin)
    ctx.set_header(.vary, 'Origin, Access-Control-Request-Headers')

    if options.allow_credentials {
        ctx.set_header(.access_control_allow_credentials, 'true')
    }

    // validate request method
    if ctx.req.method !in options.allowed_methods {
        ctx.res.set_status(.method_not_allowed)
        ctx.text('${ctx.req.method} requests are not allowed')

        $if veb_trace_cors ? {
            eprintln('[veb]: rejected CORS request from "${origin}". Reason: invalid request method: ${ctx.req.method}')
        }
        return false
    }

    if options.allowed_headers.len > 0 && options.allowed_headers != ['*'] {
        // validate request headers
        for header in ctx.req.header.keys() {
            if header !in options.allowed_headers {
                ctx.res.set_status(.forbidden)
                ctx.text('invalid Header "${header}"')

                $if veb_trace_cors ? {
                    eprintln('[veb]: rejected CORS request from "${origin}". Reason: invalid header "${header}"')
                }
                return false
            }
        }
    }

    $if veb_trace_cors ? {
        eprintln('[veb]: received CORS request from "${origin}": HTTP ${ctx.req.method} ${ctx.req.url}')
    }

    return true
}

// cors handles cross-origin requests by adding Access-Control-* headers to a
// preflight request and validating the headers of a cross-origin request.
// Usage example:
// ```v
// app.use(veb.cors[Context](veb.CorsOptions{
//     origins: ['*']
//     allowed_methods: [.get, .head, .patch, .put, .post, .delete]
// }))
// ```
pub fn cors[T](options CorsOptions) MiddlewareOptions[T] {
    return MiddlewareOptions[T]{
        handler: fn [options] [T](mut ctx T) bool {
            if ctx.req.method == .options { // preflight
                options.set_headers(mut ctx.Context)
                ctx.text('ok')
                return false
            } else {
                return options.validate_request(mut ctx.Context)
            }
        }
    }
}