package usage

import (
	"bytes"
	"encoding/json"
	"strings"

	"aigw/internal/domain"
)

const maxCaptureBytes = 64 << 10

type Observer struct {
	protocol      domain.Protocol
	stream        bool
	buffer        []byte
	line          []byte
	usage         domain.Usage
	found         bool
	explicitTotal bool
}

func NewObserver(protocol domain.Protocol, stream bool) *Observer {
	return &Observer{protocol: protocol, stream: stream}
}

func (o *Observer) Write(p []byte) (int, error) {
	if o.stream {
		o.observeSSE(p)
	} else {
		o.captureTail(p)
	}
	return len(p), nil
}

func (o *Observer) Usage() domain.Usage {
	if o.stream {
		if len(o.line) > 0 {
			o.parseSSELine(o.line)
		}
		return o.usage
	}
	o.parseJSON(o.buffer)
	return o.usage
}

func (o *Observer) captureTail(p []byte) {
	if len(p) >= maxCaptureBytes {
		o.buffer = append(o.buffer[:0], p[len(p)-maxCaptureBytes:]...)
		return
	}
	if len(o.buffer)+len(p) > maxCaptureBytes {
		drop := len(o.buffer) + len(p) - maxCaptureBytes
		copy(o.buffer, o.buffer[drop:])
		o.buffer = o.buffer[:len(o.buffer)-drop]
	}
	o.buffer = append(o.buffer, p...)
}

func (o *Observer) observeSSE(p []byte) {
	o.line = append(o.line, p...)
	for {
		index := bytes.IndexByte(o.line, '\n')
		if index < 0 {
			if len(o.line) > maxCaptureBytes {
				o.line = append(o.line[:0], o.line[len(o.line)-maxCaptureBytes:]...)
			}
			return
		}
		line := bytes.TrimSpace(o.line[:index])
		o.parseSSELine(line)
		o.line = o.line[index+1:]
	}
}

func (o *Observer) parseSSELine(line []byte) {
	if !bytes.HasPrefix(line, []byte("data:")) || !bytes.Contains(line, []byte("\"usage\"")) {
		return
	}
	payload := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:")))
	if bytes.Equal(payload, []byte("[DONE]")) {
		return
	}
	o.parseJSON(payload)
}

type usageFields struct {
	PromptTokens             *int64 `json:"prompt_tokens"`
	CompletionTokens         *int64 `json:"completion_tokens"`
	TotalTokens              *int64 `json:"total_tokens"`
	InputTokens              *int64 `json:"input_tokens"`
	OutputTokens             *int64 `json:"output_tokens"`
	CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"`
	CacheReadInputTokens     *int64 `json:"cache_read_input_tokens"`
}

type responseEnvelope struct {
	Usage   *usageFields `json:"usage"`
	Message *struct {
		Usage *usageFields `json:"usage"`
	} `json:"message"`
}

func (o *Observer) parseJSON(payload []byte) {
	var envelope responseEnvelope
	if err := json.Unmarshal(payload, &envelope); err != nil {
		payload = extractUsageObject(payload)
		if len(payload) == 0 {
			return
		}
		var fields usageFields
		if json.Unmarshal(payload, &fields) == nil {
			o.apply(&fields)
		}
		return
	}
	if envelope.Usage != nil {
		o.apply(envelope.Usage)
	}
	if envelope.Message != nil && envelope.Message.Usage != nil {
		o.apply(envelope.Message.Usage)
	}
}

func (o *Observer) apply(fields *usageFields) {
	if fields.PromptTokens != nil {
		o.usage.InputTokens = *fields.PromptTokens
		o.found = true
	}
	if fields.InputTokens != nil {
		o.usage.InputTokens = *fields.InputTokens
		o.found = true
	}
	if fields.CompletionTokens != nil {
		o.usage.OutputTokens = *fields.CompletionTokens
		o.found = true
	}
	if fields.OutputTokens != nil {
		o.usage.OutputTokens = *fields.OutputTokens
		o.found = true
	}
	if fields.TotalTokens != nil {
		o.usage.TotalTokens = *fields.TotalTokens
		o.found = true
		o.explicitTotal = true
	}
	if fields.CacheCreationInputTokens != nil {
		o.usage.CacheCreationInputTokens = *fields.CacheCreationInputTokens
		o.found = true
	}
	if fields.CacheReadInputTokens != nil {
		o.usage.CacheReadInputTokens = *fields.CacheReadInputTokens
		o.found = true
	}
	if !o.explicitTotal && o.found {
		o.usage.TotalTokens = o.usage.InputTokens + o.usage.OutputTokens
	}
}

func extractUsageObject(payload []byte) []byte {
	index := strings.LastIndex(string(payload), `"usage"`)
	if index < 0 {
		return nil
	}
	rest := payload[index+len(`"usage"`):]
	start := bytes.IndexByte(rest, '{')
	if start < 0 {
		return nil
	}
	rest = rest[start:]
	depth := 0
	inString := false
	escaped := false
	for i, b := range rest {
		if inString {
			if escaped {
				escaped = false
				continue
			}
			if b == '\\' {
				escaped = true
			} else if b == '"' {
				inString = false
			}
			continue
		}
		switch b {
		case '"':
			inString = true
		case '{':
			depth++
		case '}':
			depth--
			if depth == 0 {
				return rest[:i+1]
			}
		}
	}
	return nil
}