package middleware

import (
	"encoding/json"
	"net/http"
	"server/internal/model/common"
)

type ResponseWriter struct {
	http.ResponseWriter
	statusCode int
	bytes      int
	errorMsg   string
}

func NewResponseRecorder(w http.ResponseWriter) *ResponseWriter {
	return &ResponseWriter{
		ResponseWriter: w,
		statusCode:     http.StatusOK,
	}
}

func (rw *ResponseWriter) WriteHeader(code int) {
	rw.statusCode = code
	rw.ResponseWriter.WriteHeader(code)
}

// extractErrorMessage 从响应体中提取错误信息
func (rw *ResponseWriter) extractErrorMessage(body []byte) string {
	var resp common.Response
	if err := json.Unmarshal(body, &resp); err == nil && resp.Message != "" {
		return resp.Message
	}
	return string(body)
}

func (rw *ResponseWriter) Write(b []byte) (int, error) {
	n, err := rw.ResponseWriter.Write(b)
	rw.bytes += n

	// 只在错误状态码且未记录错误时处理
	if rw.statusCode >= 400 && rw.errorMsg == "" && n > 0 {
		rw.errorMsg = rw.extractErrorMessage(b[:n])
	}

	return n, err
}