from pathlib import Path
from datetime import datetime
import shutil

path = Path('/home/admin/docker-projects/sub2api/backend/internal/service/openai_gateway_service.go')
trash_dir = Path('/home/admin/docker-projects/sub2api/.trash')
text = path.read_text()
orig = text

if 'func openAIResponseHasUsableOutput(body []byte) bool {' in text:
    print('already patched')
    raise SystemExit(0)

old_block = '''	if account.Type == AccountTypeOAuth {
		bodyLooksLikeSSE := bytes.Contains(body, []byte("data:")) || bytes.Contains(body, []byte("event:"))
		if isEventStreamResponse(resp.Header) || bodyLooksLikeSSE {
			return s.handleOAuthSSEToJSON(resp, c, body, originalModel, mappedModel)
		}
	}
'''
new_block = '''	bodyLooksLikeSSE := bytes.Contains(body, []byte("data:")) || bytes.Contains(body, []byte("event:"))
	if isEventStreamResponse(resp.Header) || bodyLooksLikeSSE {
		return s.handleOAuthSSEToJSON(resp, c, body, originalModel, mappedModel)
	}
'''
if old_block not in text:
    raise SystemExit('old nonstream SSE block not found')
text = text.replace(old_block, new_block, 1)

anchor = 'usage := &OpenAIUsage{}\n\tif ok {\n'
insert = '''usage := &OpenAIUsage{}
	if ok {
		if !openAIResponseHasUsableOutput(finalResponse) {
			fallbackText := extractOpenAISSEOutputText(bodyText)
			if strings.TrimSpace(fallbackText) == "" {
				return nil, fmt.Errorf("completed non-streaming response without output text")
			}
			nextResponse, injectErr := injectOpenAIResponseOutputText(finalResponse, fallbackText)
			if injectErr != nil {
				return nil, fmt.Errorf("reconstruct non-streaming response output: %w", injectErr)
			}
			finalResponse = nextResponse
		}
'''
if anchor not in text:
    raise SystemExit('handleOAuthSSEToJSON anchor not found')
text = text.replace(anchor, insert, 1)

helpers_anchor = 'func (s *OpenAIGatewayService) parseSSEUsageFromBody(body string) *OpenAIUsage {'
helpers = '''func openAIResponseHasUsableOutput(body []byte) bool {
	var payload struct {
		Output []json.RawMessage `json:"output"`
	}
	if err := json.Unmarshal(body, &payload); err != nil {
		return false
	}
	return len(payload.Output) > 0
}

func extractOpenAISSEOutputText(body string) string {
	lines := strings.Split(body, "\n")
	var deltaBuilder strings.Builder
	var doneParts []string
	sawDelta := false

	for _, line := range lines {
		data := openaiSSEDataRe.ReplaceAllString(line, "")
		if data == line || data == "" || data == "[DONE]" {
			continue
		}

		var event struct {
			Type  string `json:"type"`
			Delta string `json:"delta"`
			Text  string `json:"text"`
		}
		if err := json.Unmarshal([]byte(data), &event); err != nil {
			continue
		}

		switch strings.TrimSpace(event.Type) {
		case "response.output_text.delta":
			if event.Delta == "" {
				continue
			}
			sawDelta = true
			deltaBuilder.WriteString(event.Delta)
		case "response.output_text.done":
			if sawDelta {
				continue
			}
			if event.Text != "" {
				doneParts = append(doneParts, event.Text)
			}
		}
	}

	if sawDelta {
		return deltaBuilder.String()
	}
	return strings.Join(doneParts, "")
}

func injectOpenAIResponseOutputText(body []byte, text string) ([]byte, error) {
	var payload map[string]any
	if err := json.Unmarshal(body, &payload); err != nil {
		return nil, err
	}
	payload["output"] = []map[string]any{
		{
			"type": "message",
			"role": "assistant",
			"content": []map[string]any{
				{
					"type": "output_text",
					"text": text,
				},
			},
		},
	}
	if _, ok := payload["status"]; !ok {
		payload["status"] = "completed"
	}
	if _, ok := payload["object"]; !ok {
		payload["object"] = "response"
	}
	return json.Marshal(payload)
}

''' + helpers_anchor
if helpers_anchor not in text:
    raise SystemExit('helpers anchor not found')
text = text.replace(helpers_anchor, helpers, 1)

if text == orig:
    raise SystemExit('no changes applied')

ts = datetime.now().strftime('%Y%m%d_%H%M%S')
trash_dir.mkdir(parents=True, exist_ok=True)
backup = trash_dir / f'openai_gateway_service.go_{ts}'
shutil.copy2(path, backup)
path.write_text(text)
print(f'patched {path}')
print(f'backup {backup}')