package gateway

import (
	"bytes"
	"io"
	"strings"
)

// outputRedactReadCloser 对上游模型响应应用输出侧脱敏(隐私信息拦截替换):
//   - 非流式(application/json):首次 Read 前缓冲整个响应,处理后再输出;
//   - 流式(text/event-stream):逐 data 行处理,不改变事件边界。
type outputRedactReadCloser struct {
	io.ReadCloser
	engine interface {
		OutputRedact([]byte) ([]byte, bool)
	}
	sse      bool
	buffered []byte // 已处理待输出的字节
	done     bool   // 非流式已完成缓冲与处理
	pending  []byte // 流式:未完成的行
}

func newOutputRedactReadCloser(body io.ReadCloser, contentType string, engine interface {
	OutputRedact([]byte) ([]byte, bool)
}) io.ReadCloser {
	if engine == nil {
		return body
	}
	return &outputRedactReadCloser{
		ReadCloser: body, engine: engine,
		sse: strings.Contains(strings.ToLower(contentType), "text/event-stream"),
	}
}

func (r *outputRedactReadCloser) Read(buffer []byte) (int, error) {
	if !r.sse {
		// 非流式:首次 Read 时一次性缓冲+处理。
		if !r.done {
			r.done = true
			raw, err := io.ReadAll(r.ReadCloser)
			_ = err
			if replaced, changed := r.engine.OutputRedact(raw); changed {
				r.buffered = replaced
			} else {
				r.buffered = raw
			}
		}
		if len(r.buffered) == 0 {
			return 0, io.EOF
		}
		n := copy(buffer, r.buffered)
		r.buffered = r.buffered[n:]
		if len(r.buffered) == 0 {
			return n, io.EOF
		}
		return n, nil
	}
	// 流式:先输出已处理的行,再读上游。
	if len(r.buffered) > 0 {
		n := copy(buffer, r.buffered)
		r.buffered = r.buffered[n:]
		return n, nil
	}
	chunk := make([]byte, 32<<10)
	n, err := r.ReadCloser.Read(chunk)
	if n > 0 {
		r.pending = append(r.pending, chunk[:n]...)
		r.processLines()
	}
	if err == io.EOF && len(r.pending) > 0 {
		// 流结束:剩余不完整行原样输出(不处理,避免破坏事件边界)。
		r.buffered = append(r.buffered, r.pending...)
		r.pending = nil
	}
	if len(r.buffered) > 0 {
		n2 := copy(buffer, r.buffered)
		r.buffered = r.buffered[n2:]
		return n2, err
	}
	return n, err
}

// processLines 把 pending 中的完整行处理进 buffered。
func (r *outputRedactReadCloser) processLines() {
	for {
		index := bytes.IndexByte(r.pending, '\n')
		if index < 0 {
			return
		}
		line := r.pending[:index]
		r.pending = r.pending[index+1:]
		trimmed := strings.TrimSpace(string(line))
		if !strings.HasPrefix(trimmed, "data:") {
			r.buffered = append(r.buffered, line...)
			r.buffered = append(r.buffered, '\n')
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
		if payload == "" || payload == "[DONE]" {
			r.buffered = append(r.buffered, line...)
			r.buffered = append(r.buffered, '\n')
			continue
		}
		if replaced, changed := r.engine.OutputRedact([]byte(payload)); changed {
			r.buffered = append(r.buffered, []byte("data: "+string(replaced)+"\n")...)
		} else {
			r.buffered = append(r.buffered, line...)
			r.buffered = append(r.buffered, '\n')
		}
	}
}