package util

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import (
	"net/http"
)

// Interceptor implements http.ResponseWriter.
// It intercepts writes to w, and tracks the HTTP code and the count of bytes written, while still writing to w.
type Interceptor struct {
	W         http.ResponseWriter
	Code      int
	ByteCount int
}

// WriteHeader implements http.ResponseWriter.
// It does the real write to Interceptor's internal ResponseWriter, while keeping track of the code.
func (i *Interceptor) WriteHeader(rc int) {
	i.W.WriteHeader(rc)
	i.Code = rc
}

// Write implements http.ResponseWriter.
// It does the real write to Interceptor's internal ResponseWriter, while keeping track of the count of bytes written.
// It also sets Interceptor's tracked code to 200 if WriteHeader wasn't called (which is what the real http.ResponseWriter will write to the client).
func (i *Interceptor) Write(b []byte) (int, error) {
	wi, werr := i.W.Write(b)
	i.ByteCount += wi
	if i.Code == 0 {
		i.Code = 200
	}
	return wi, werr
}

// Header implements http.ResponseWriter.
// It returns Interceptor's internal ResponseWriter.Header, without modification or tracking.
func (i *Interceptor) Header() http.Header {
	return i.W.Header()
}

// BodyInterceptor fulfills the Writer interface, but records the body and doesn't actually write. This allows performing operations on the entire body written by a handler, for example, compressing or hashing. To actually write, call `RealWrite()`. Note this means `len(b)` and `nil` are always returned by `Write()`, any real write errors will be returned by `RealWrite()`.
type BodyInterceptor struct {
	W         http.ResponseWriter
	BodyBytes []byte
}

// WriteHeader implements http.ResponseWriter.
// It does the real write to Interceptor's internal ResponseWriter, without modification or tracking.
func (i *BodyInterceptor) WriteHeader(rc int) {
	i.W.WriteHeader(rc)
}

// Write implements http.ResponseWriter.
// It writes the given bytes to BodyInterceptor's internal tracking bytes, and does not write them to the internal ResponseWriter.
// To write the internal bytes, call BodyInterceptor.RealWrite.
func (i *BodyInterceptor) Write(b []byte) (int, error) {
	i.BodyBytes = append(i.BodyBytes, b...)
	return len(b), nil
}

// Header implements http.ResponseWriter.
// It returns BodyInterceptor's internal ResponseWriter.Header, without modification or tracking.
func (i *BodyInterceptor) Header() http.Header {
	return i.W.Header()
}

// RealWrite writes BodyInterceptor's internal bytes, which were stored by calls to Write.
func (i *BodyInterceptor) RealWrite(b []byte) (int, error) {
	wi, werr := i.W.Write(i.BodyBytes)
	return wi, werr
}

// Body returns the internal bytes stored by calls to Write.
func (i *BodyInterceptor) Body() []byte {
	return i.BodyBytes
}