internal/runtime/executor/kiro_executor.go

package executor

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
	sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
	log "github.com/sirupsen/logrus"
	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"
)

// KiroExecutor implements ProviderExecutor for Kiro API (Amazon Q Developer / AWS CodeWhisperer).
// It translates OpenAI/Claude-compatible requests to the Kiro API format.
type KiroExecutor struct {
	cfg *config.Config
}

// NewKiroExecutor creates a new KiroExecutor instance.
func NewKiroExecutor(cfg *config.Config) *KiroExecutor {
	return &KiroExecutor{cfg: cfg}
}

// Identifier returns the provider key for this executor.
func (e *KiroExecutor) Identifier() string { return "kiro" }

// kiroAPIHost returns the Kiro API host for the given region.
func kiroAPIHost(region string) string {
	if region == "" {
		region = "us-east-1"
	}
	return fmt.Sprintf("https://codewhisperer.%s.amazonaws.com", region)
}

// kiroRefreshURL returns the Kiro Desktop Auth token refresh URL.
func kiroRefreshURL(region string) string {
	if region == "" {
		region = "us-east-1"
	}
	return fmt.Sprintf("https://prod.%s.auth.desktop.kiro.dev/refreshToken", region)
}

// kiroCreds extracts Kiro credentials from auth.
func kiroCreds(a *cliproxyauth.Auth) (accessToken, refreshToken, region, profileARN string) {
	if a == nil {
		return
	}
	if a.Attributes != nil {
		region = a.Attributes["region"]
		profileARN = a.Attributes["profile_arn"]
	}
	if a.Metadata != nil {
		if v, ok := a.Metadata["access_token"].(string); ok {
			accessToken = v
		}
		if v, ok := a.Metadata["refresh_token"].(string); ok {
			refreshToken = v
		}
		if v, ok := a.Metadata["region"].(string); ok && v != "" {
			region = v
		}
		if v, ok := a.Metadata["profile_arn"].(string); ok && v != "" {
			profileARN = v
		}
	}
	if region == "" {
		region = "us-east-1"
	}
	return
}

// Kiro model name mappings (display name -> internal Kiro ID)
var kiroModelMappings = map[string]string{
	"claude-sonnet-4":    "CLAUDE_SONNET_4_V1_0",
	"claude-sonnet-4.5":  "CLAUDE_SONNET_4_5_V1_0",
	"claude-haiku-4.5":   "CLAUDE_HAIKU_4_5_V1_0",
	"claude-opus-4.5":    "CLAUDE_OPUS_4_5_V1_0",
	"claude-3.7-sonnet":  "CLAUDE_3_7_SONNET_20250219_V1_0",
	"claude-3-7-sonnet":  "CLAUDE_3_7_SONNET_20250219_V1_0",
	"auto":               "auto",
}

// normalizeKiroModel normalizes model names for the Kiro API.
func normalizeKiroModel(model string) string {
	model = strings.TrimSpace(model)
	// Normalize dashes: claude-sonnet-4-5 -> claude-sonnet-4.5
	model = regexp.MustCompile(`(\d)-(\d)`).ReplaceAllString(model, "${1}.${2}")

	// Strip date suffixes like -20250929
	model = regexp.MustCompile(`-\d{8}$`).ReplaceAllString(model, "")

	if kiroID, ok := kiroModelMappings[model]; ok {
		return kiroID
	}

	// Pass through unknown models to let Kiro decide
	return model
}

// buildKiroPayload builds the Kiro API request payload from a Claude-format request.
func buildKiroPayload(body []byte, model string, profileARN string) ([]byte, error) {
	kiroModel := normalizeKiroModel(model)

	// Extract messages
	messages := gjson.GetBytes(body, "messages")
	if !messages.Exists() {
		return nil, fmt.Errorf("messages field is required")
	}

	// Build Kiro conversationState
	var userInputs []map[string]any
	var assistantResponses []map[string]any

	messages.ForEach(func(_, msg gjson.Result) bool {
		role := msg.Get("role").String()
		content := msg.Get("content")

		switch role {
		case "user":
			userMsg := map[string]any{
				"content": extractTextContent(content),
			}
			userInputs = append(userInputs, userMsg)
		case "assistant":
			assistantMsg := map[string]any{
				"content": extractTextContent(content),
			}
			assistantResponses = append(assistantResponses, assistantMsg)
		}
		return true
	})

	// Build system prompt
	var systemPrompt string
	if sys := gjson.GetBytes(body, "system"); sys.Exists() {
		if sys.IsArray() {
			var parts []string
			sys.ForEach(func(_, part gjson.Result) bool {
				if part.Get("type").String() == "text" {
					parts = append(parts, part.Get("text").String())
				}
				return true
			})
			systemPrompt = strings.Join(parts, "\n")
		} else {
			systemPrompt = sys.String()
		}
	}

	// Build conversation state
	conversationState := map[string]any{
		"currentMessage": map[string]any{
			"origin":          "USER",
			"userInputs":      userInputs,
			"userInputOrigin": "CONVERSATION",
		},
	}

	if len(assistantResponses) > 0 {
		conversationState["history"] = []map[string]any{
			{
				"turn": map[string]any{
					"userInputs":         userInputs[:len(userInputs)-1],
					"assistantResponses": assistantResponses,
				},
			},
		}
	}

	// Build Kiro request payload
	kiroPayload := map[string]any{
		"conversationState": conversationState,
		"additionalInstructions": systemPrompt,
		"profileArn": profileARN,
	}

	// Add model selection if not auto
	if kiroModel != "auto" {
		kiroPayload["modelRoutingConfiguration"] = map[string]any{
			"explicitModelId": kiroModel,
		}
	}

	return json.Marshal(kiroPayload)
}

// extractTextContent extracts text from message content.
func extractTextContent(content gjson.Result) string {
	if content.IsArray() {
		var parts []string
		content.ForEach(func(_, part gjson.Result) bool {
			if part.Get("type").String() == "text" {
				parts = append(parts, part.Get("text").String())
			}
			return true
		})
		return strings.Join(parts, "\n")
	}
	return content.String()
}

// PrepareRequest injects Kiro credentials into the outgoing HTTP request.
func (e *KiroExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error {
	if req == nil {
		return nil
	}
	accessToken, _, _, _ := kiroCreds(auth)
	if accessToken == "" {
		return nil
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	var attrs map[string]string
	if auth != nil {
		attrs = auth.Attributes
	}
	util.ApplyCustomHeadersFromAttrs(req, attrs)
	return nil
}

// HttpRequest injects Kiro credentials and executes the request.
func (e *KiroExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) {
	if req == nil {
		return nil, fmt.Errorf("kiro executor: request is nil")
	}
	if ctx == nil {
		ctx = req.Context()
	}
	httpReq := req.WithContext(ctx)
	if err := e.PrepareRequest(httpReq, auth); err != nil {
		return nil, err
	}
	httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
	return httpClient.Do(httpReq)
}

// Execute handles non-streaming Kiro API calls.
func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
	baseModel := thinking.ParseSuffix(req.Model).ModelName

	accessToken, _, region, profileARN := kiroCreds(auth)
	apiHost := kiroAPIHost(region)

	reporter := newUsageReporter(ctx, e.Identifier(), baseModel, auth)
	defer reporter.trackFailure(ctx, &err)

	from := opts.SourceFormat
	to := sdktranslator.FromString("claude")

	// Translate to Claude format first
	body := sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), false)
	body, _ = sjson.SetBytes(body, "model", baseModel)

	// Build Kiro payload
	kiroBody, err := buildKiroPayload(body, baseModel, profileARN)
	if err != nil {
		return resp, err
	}

	url := fmt.Sprintf("%s/v1/generateAssistantResponse", apiHost)
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroBody))
	if err != nil {
		return resp, err
	}

	applyKiroHeaders(httpReq, accessToken, profileARN, false)

	var authID, authLabel, authType, authValue string
	if auth != nil {
		authID = auth.ID
		authLabel = auth.Label
		authType, authValue = auth.AccountInfo()
	}
	recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
		URL:       url,
		Method:    http.MethodPost,
		Headers:   httpReq.Header.Clone(),
		Body:      kiroBody,
		Provider:  e.Identifier(),
		AuthID:    authID,
		AuthLabel: authLabel,
		AuthType:  authType,
		AuthValue: authValue,
	})

	httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
	httpResp, err := httpClient.Do(httpReq)
	if err != nil {
		recordAPIResponseError(ctx, e.cfg, err)
		return resp, err
	}
	recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())

	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
		b, _ := io.ReadAll(httpResp.Body)
		appendAPIResponseChunk(ctx, e.cfg, b)
		log.Debugf("kiro request error, status: %d, message: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
		err = statusErr{code: httpResp.StatusCode, msg: string(b)}
		if errClose := httpResp.Body.Close(); errClose != nil {
			log.Errorf("response body close error: %v", errClose)
		}
		return resp, err
	}

	defer func() {
		if errClose := httpResp.Body.Close(); errClose != nil {
			log.Errorf("response body close error: %v", errClose)
		}
	}()

	// Parse Kiro streaming response and convert to Claude format
	data, err := io.ReadAll(httpResp.Body)
	if err != nil {
		recordAPIResponseError(ctx, e.cfg, err)
		return resp, err
	}
	appendAPIResponseChunk(ctx, e.cfg, data)

	// Convert Kiro response to Claude format
	claudeResp := convertKiroToClaudeResponse(data, req.Model)

	var param any
	out := sdktranslator.TranslateNonStream(
		ctx,
		to,
		from,
		req.Model,
		bytes.Clone(opts.OriginalRequest),
		body,
		claudeResp,
		&param,
	)
	resp = cliproxyexecutor.Response{Payload: []byte(out)}
	return resp, nil
}

// ExecuteStream handles streaming Kiro API calls.
func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
	baseModel := thinking.ParseSuffix(req.Model).ModelName

	accessToken, _, region, profileARN := kiroCreds(auth)
	apiHost := kiroAPIHost(region)

	reporter := newUsageReporter(ctx, e.Identifier(), baseModel, auth)
	defer reporter.trackFailure(ctx, &err)

	from := opts.SourceFormat
	to := sdktranslator.FromString("claude")

	// Translate to Claude format first
	body := sdktranslator.TranslateRequest(from, to, baseModel, bytes.Clone(req.Payload), true)
	body, _ = sjson.SetBytes(body, "model", baseModel)

	// Build Kiro payload
	kiroBody, err := buildKiroPayload(body, baseModel, profileARN)
	if err != nil {
		return nil, err
	}

	url := fmt.Sprintf("%s/v1/generateAssistantResponse", apiHost)
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(kiroBody))
	if err != nil {
		return nil, err
	}

	applyKiroHeaders(httpReq, accessToken, profileARN, true)

	var authID, authLabel, authType, authValue string
	if auth != nil {
		authID = auth.ID
		authLabel = auth.Label
		authType, authValue = auth.AccountInfo()
	}
	recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
		URL:       url,
		Method:    http.MethodPost,
		Headers:   httpReq.Header.Clone(),
		Body:      kiroBody,
		Provider:  e.Identifier(),
		AuthID:    authID,
		AuthLabel: authLabel,
		AuthType:  authType,
		AuthValue: authValue,
	})

	httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
	httpResp, err := httpClient.Do(httpReq)
	if err != nil {
		recordAPIResponseError(ctx, e.cfg, err)
		return nil, err
	}
	recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())

	if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
		b, _ := io.ReadAll(httpResp.Body)
		appendAPIResponseChunk(ctx, e.cfg, b)
		log.Debugf("kiro request error, status: %d, message: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
		if errClose := httpResp.Body.Close(); errClose != nil {
			log.Errorf("response body close error: %v", errClose)
		}
		err = statusErr{code: httpResp.StatusCode, msg: string(b)}
		return nil, err
	}

	out := make(chan cliproxyexecutor.StreamChunk)
	stream = out

	go func() {
		defer close(out)
		defer func() {
			if errClose := httpResp.Body.Close(); errClose != nil {
				log.Errorf("response body close error: %v", errClose)
			}
		}()

		scanner := bufio.NewScanner(httpResp.Body)
		scanner.Buffer(nil, 52_428_800) // 50MB

		var contentBuilder strings.Builder
		var param any

		for scanner.Scan() {
			line := scanner.Bytes()
			appendAPIResponseChunk(ctx, e.cfg, line)

			// Parse Kiro SSE format and convert to Claude format
			chunk := parseKiroStreamLine(line, req.Model, &contentBuilder)
			if chunk == nil {
				continue
			}

			// Translate to target format
			chunks := sdktranslator.TranslateStream(
				ctx,
				to,
				from,
				req.Model,
				bytes.Clone(opts.OriginalRequest),
				body,
				chunk,
				&param,
			)
			for i := range chunks {
				out <- cliproxyexecutor.StreamChunk{Payload: []byte(chunks[i])}
			}
		}

		if errScan := scanner.Err(); errScan != nil {
			recordAPIResponseError(ctx, e.cfg, errScan)
			reporter.publishFailure(ctx)
			out <- cliproxyexecutor.StreamChunk{Err: errScan}
		}
	}()

	return stream, nil
}

// CountTokens returns the token count for the given request.
func (e *KiroExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
	// Kiro doesn't have a dedicated token counting endpoint, return an estimate
	from := opts.SourceFormat
	to := sdktranslator.FromString("claude")

	body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)

	// Rough estimate based on content length
	contentLen := len(body)
	estimatedTokens := contentLen / 4 // rough approximation

	result := map[string]any{
		"input_tokens": estimatedTokens,
	}
	data, _ := json.Marshal(result)

	out := sdktranslator.TranslateTokenCount(ctx, to, from, int64(estimatedTokens), data)
	return cliproxyexecutor.Response{Payload: []byte(out)}, nil
}

// Refresh refreshes Kiro credentials.
func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
	log.Debugf("kiro executor: refresh called")
	if auth == nil {
		return nil, fmt.Errorf("kiro executor: auth is nil")
	}

	_, refreshToken, region, _ := kiroCreds(auth)
	if refreshToken == "" {
		return auth, nil
	}

	refreshURL := kiroRefreshURL(region)

	payload := map[string]string{
		"refreshToken": refreshToken,
	}
	payloadBytes, _ := json.Marshal(payload)

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, bytes.NewReader(payloadBytes))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 30*time.Second)
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer func() {
		if errClose := resp.Body.Close(); errClose != nil {
			log.Errorf("response body close error: %v", errClose)
		}
	}()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("kiro refresh failed: %d %s", resp.StatusCode, string(b))
	}

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var result struct {
		AccessToken  string `json:"accessToken"`
		RefreshToken string `json:"refreshToken,omitempty"`
		ExpiresIn    int64  `json:"expiresIn,omitempty"`
		ProfileArn   string `json:"profileArn,omitempty"`
	}
	if err := json.Unmarshal(data, &result); err != nil {
		return nil, err
	}

	if auth.Metadata == nil {
		auth.Metadata = make(map[string]any)
	}
	auth.Metadata["access_token"] = result.AccessToken
	if result.RefreshToken != "" {
		auth.Metadata["refresh_token"] = result.RefreshToken
	}
	if result.ExpiresIn > 0 {
		auth.Metadata["expired"] = time.Now().Add(time.Duration(result.ExpiresIn) * time.Second).Format(time.RFC3339)
	}
	// Profile ARN is returned by the refresh API - store it for future requests
	if result.ProfileArn != "" {
		auth.Metadata["profile_arn"] = result.ProfileArn
	}
	auth.Metadata["type"] = "kiro"
	auth.Metadata["last_refresh"] = time.Now().Format(time.RFC3339)

	return auth, nil
}

// applyKiroHeaders sets the required headers for Kiro API requests.
func applyKiroHeaders(r *http.Request, accessToken, profileARN string, stream bool) {
	r.Header.Set("Authorization", "Bearer "+accessToken)
	r.Header.Set("Content-Type", "application/json")
	r.Header.Set("Accept", "application/vnd.amazon.eventstream")
	r.Header.Set("X-Amz-Target", "AmazonQDeveloperStreamingService.GenerateAssistantResponse")
	r.Header.Set("X-Amzn-Codewhisperer-Profilearn", profileARN)
	r.Header.Set("Connection", "keep-alive")

	if stream {
		r.Header.Set("Accept", "application/vnd.amazon.eventstream")
	} else {
		r.Header.Set("Accept", "application/json")
	}
}

// convertKiroToClaudeResponse converts Kiro response to Claude format.
func convertKiroToClaudeResponse(data []byte, model string) []byte {
	// Parse Kiro streaming response chunks
	var contentBuilder strings.Builder

	lines := bytes.Split(data, []byte("\n"))
	for _, line := range lines {
		line = bytes.TrimSpace(line)
		if len(line) == 0 {
			continue
		}

		// Parse AWS event stream format
		if bytes.HasPrefix(line, []byte(":event-type")) || bytes.HasPrefix(line, []byte(":message-type")) {
			continue
		}

		// Extract text content from Kiro response
		text := gjson.GetBytes(line, "assistantResponseEvent.content").String()
		if text != "" {
			contentBuilder.WriteString(text)
		}
	}

	// Build Claude-format response
	claudeResp := map[string]any{
		"id":           fmt.Sprintf("msg_%d", time.Now().UnixNano()),
		"type":         "message",
		"role":         "assistant",
		"model":        model,
		"content": []map[string]any{
			{
				"type": "text",
				"text": contentBuilder.String(),
			},
		},
		"stop_reason": "end_turn",
		"usage": map[string]any{
			"input_tokens":  0,
			"output_tokens": len(contentBuilder.String()) / 4,
		},
	}

	result, _ := json.Marshal(claudeResp)
	return result
}

// parseKiroStreamLine parses a Kiro SSE line and converts to Claude format.
func parseKiroStreamLine(line []byte, model string, contentBuilder *strings.Builder) []byte {
	line = bytes.TrimSpace(line)
	if len(line) == 0 {
		return nil
	}

	// Skip AWS event stream metadata lines
	if bytes.HasPrefix(line, []byte(":")) {
		return nil
	}

	// Parse data payload
	if bytes.HasPrefix(line, []byte("data:")) {
		line = bytes.TrimPrefix(line, []byte("data:"))
		line = bytes.TrimSpace(line)
	}

	if len(line) == 0 || !gjson.ValidBytes(line) {
		return nil
	}

	// Extract text content
	text := gjson.GetBytes(line, "assistantResponseEvent.content").String()
	if text == "" {
		return nil
	}

	contentBuilder.WriteString(text)

	// Build Claude streaming format
	event := map[string]any{
		"type": "content_block_delta",
		"index": 0,
		"delta": map[string]any{
			"type": "text_delta",
			"text": text,
		},
	}

	data, _ := json.Marshal(event)
	return append([]byte("data: "), data...)
}