package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync/atomic"
	"time"
)

func initTrace(logDir string, cfg TraceConfig) (*traceRecorder, func()) {
	if !cfg.Enabled {
		return &traceRecorder{enabled: false}, func() {}
	}
	if logDir == "" {
		logDir = defaultLogDir
	}
	if err := os.MkdirAll(logDir, 0755); err != nil {
		logEvent("Trace", "error", "trace 目录创建失败，已关闭 trace: %v", err)
		return &traceRecorder{enabled: false}, func() {}
	}

	path := filepath.Join(logDir, "douyu_trace_"+time.Now().In(cst).Format("20060102_150405")+".log")
	file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
	if err != nil {
		logEvent("Trace", "error", "trace 文件创建失败，已关闭 trace: %v", err)
		return &traceRecorder{enabled: false}, func() {}
	}

	recorder := &traceRecorder{
		enabled:      true,
		maxBodyBytes: cfg.MaxBodyBytes,
		writer:       newAsyncLogWriter(file, 2048),
		filePath:     path,
	}
	if cfg.MaxBodyBytes > 0 {
		logEvent("Trace", "warn", "trace 已启用: %s，单段正文最多保留 %d 字节，不做脱敏", path, cfg.MaxBodyBytes)
	} else {
		logEvent("Trace", "warn", "trace 已启用: %s，全量保留请求和响应，不做脱敏", path)
	}
	return recorder, func() {
		_ = recorder.Close()
		_ = file.Close()
	}
}

func (t *traceRecorder) Close() error {
	if t == nil || t.writer == nil {
		return nil
	}
	return t.writer.Close()
}

func (t *traceRecorder) record(label string, req *http.Request, status int, cost time.Duration, reqBody, respBody []byte, errText string, apiError int, apiMsg string) {
	if t == nil || !t.enabled || req == nil || req.URL == nil {
		return
	}
	if label == "服务时间" && status == http.StatusOK && errText == "" && apiError == 0 {
		return
	}
	entry := traceEntry{
		Seq:          atomic.AddUint64(&t.seq, 1),
		Time:         time.Now().In(cst).Format("2006-01-02 15:04:05.000"),
		Label:        label,
		Method:       req.Method,
		URL:          req.URL.Redacted(),
		Path:         req.URL.Path,
		Status:       status,
		Cost:         cost,
		ReqHeader:    sanitizeHeader(req.Header),
		ReqBody:      sanitizeBodyForTrace(reqBody, req.Header.Get("Content-Type"), t.maxBodyBytes),
		RespBody:     sanitizeResponseBody(respBody, t.maxBodyBytes),
		APIError:     apiError,
		APIMessage:   apiMsg,
		ErrorText:    errText,
		MaxBodyBytes: t.maxBodyBytes,
	}
	_, _ = t.writer.Write([]byte(formatTraceEntry(entry)))
}

func traceRequestBody(t *traceRecorder, req *http.Request) []byte {
	if t == nil || !t.enabled || req == nil || req.GetBody == nil {
		return nil
	}
	body, err := req.GetBody()
	if err != nil {
		return nil
	}
	defer body.Close()
	if t.maxBodyBytes > 0 {
		data, _ := io.ReadAll(io.LimitReader(body, int64(t.maxBodyBytes)+1))
		return data
	}
	data, _ := io.ReadAll(body)
	return data
}

func sanitizeHeader(header http.Header) map[string][]string {
	out := make(map[string][]string, len(header))
	for key, values := range header {
		canonical := http.CanonicalHeaderKey(key)
		copied := make([]string, len(values))
		copy(copied, values)
		out[canonical] = copied
	}
	return out
}

func sanitizeBodyForTrace(body []byte, contentType string, max int) string {
	if len(body) == 0 {
		return ""
	}
	mediaType, params, _ := mime.ParseMediaType(contentType)
	if strings.HasPrefix(mediaType, "multipart/") {
		boundary := params["boundary"]
		if boundary != "" {
			return sanitizeMultipartBody(body, boundary, max)
		}
	}
	return sanitizeTextBody(body, max)
}

func sanitizeResponseBody(body []byte, max int) string {
	return sanitizeTextBody(body, max)
}

func sanitizeTextBody(body []byte, max int) string {
	raw := string(body)
	if pretty, ok := prettyJSON(raw); ok {
		return truncateTraceText(pretty, max)
	}
	return truncateTraceText(raw, max)
}

func sanitizeMultipartBody(body []byte, boundary string, max int) string {
	reader := multipartReader(bytes.NewReader(body), boundary)
	fields := make(map[string][]string)
	for {
		part, err := reader.NextPart()
		if err == io.EOF {
			break
		}
		if err != nil {
			return truncateTraceText(string(body), max)
		}
		name := part.FormName()
		if name == "" {
			continue
		}
		var valueBytes []byte
		if max > 0 {
			valueBytes, _ = io.ReadAll(io.LimitReader(part, int64(max)+1))
		} else {
			valueBytes, _ = io.ReadAll(part)
		}
		value := truncateTraceText(string(valueBytes), max)
		fields[name] = append(fields[name], value)
	}
	encoded, err := json.Marshal(fields)
	if err != nil {
		return ""
	}
	return string(encoded)
}

func multipartReader(r io.Reader, boundary string) *multipart.Reader {
	return multipart.NewReader(r, boundary)
}

func truncateTraceText(s string, max int) string {
	if max <= 0 || len(s) <= max {
		return s
	}
	return s[:max] + fmt.Sprintf("...<truncated %d bytes>", len(s)-max)
}

type traceEntry struct {
	Seq          uint64
	Time         string
	Label        string
	Method       string
	URL          string
	Path         string
	Status       int
	Cost         time.Duration
	ReqHeader    map[string][]string
	ReqBody      string
	RespBody     string
	APIError     int
	APIMessage   string
	ErrorText    string
	MaxBodyBytes int
}

func formatTraceEntry(e traceEntry) string {
	var b strings.Builder
	statusText := "成功"
	if e.Status < 200 || e.Status >= 300 || e.APIError != 0 || e.ErrorText != "" {
		statusText = "异常"
	}

	b.WriteString(strings.Repeat("=", 96))
	b.WriteByte('\n')
	fmt.Fprintf(&b, "#%06d %s [%s] %s %s\n", e.Seq, e.Time, statusText, e.Method, e.URL)
	fmt.Fprintf(&b, "接口: %s\n", emptyDefault(e.Label, "-"))
	fmt.Fprintf(&b, "路径: %s\n", emptyDefault(e.Path, "-"))
	fmt.Fprintf(&b, "结果: HTTP %d | 耗时 %s | apiError=%d", e.Status, roundDuration(e.Cost), e.APIError)
	if e.APIMessage != "" {
		fmt.Fprintf(&b, " | msg=%s", e.APIMessage)
	}
	b.WriteByte('\n')
	if e.ErrorText != "" {
		fmt.Fprintf(&b, "错误: %s\n", e.ErrorText)
	}
	b.WriteString("\n[请求头]\n")
	b.WriteString(formatTraceHeaders(e.ReqHeader))
	b.WriteString("\n[请求体]\n")
	b.WriteString(formatTraceBody(e.ReqBody))
	b.WriteString("\n[响应体]\n")
	b.WriteString(formatTraceBody(e.RespBody))
	if e.MaxBodyBytes > 0 {
		fmt.Fprintf(&b, "\n提示: 单段正文最多保留 %d 字节；trace 不做脱敏。\n", e.MaxBodyBytes)
	} else {
		b.WriteString("\n提示: trace 全量保留请求和响应，不做脱敏。\n")
	}
	b.WriteByte('\n')
	return b.String()
}

func formatTraceHeaders(headers map[string][]string) string {
	if len(headers) == 0 {
		return "<empty>\n"
	}
	keys := make([]string, 0, len(headers))
	for key := range headers {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	var b strings.Builder
	for _, key := range keys {
		fmt.Fprintf(&b, "%s: %s\n", key, strings.Join(headers[key], ", "))
	}
	return b.String()
}

func formatTraceBody(body string) string {
	body = strings.TrimSpace(body)
	if body == "" {
		return "<empty>\n"
	}
	if pretty, ok := prettyJSON(body); ok {
		return pretty + "\n"
	}
	return body + "\n"
}

func prettyJSON(s string) (string, bool) {
	var value any
	if err := json.Unmarshal([]byte(s), &value); err != nil {
		return "", false
	}
	out, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		return "", false
	}
	return string(out), true
}