package ai

import (
	"context"
	"encoding/json"
	"fmt"

	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
)

// BedrockClaudeClient AWS Bedrock Claude 客户端
type BedrockClaudeClient struct {
	client *bedrockruntime.Client
	model  string
}

// NewBedrockClaudeClient 创建 Bedrock Claude 客户端
func NewBedrockClaudeClient(region, accessKeyID, secretAccessKey, model string) (*BedrockClaudeClient, error) {
	if model == "" {
		model = "arn:aws:bedrock:us-east-1::inference-profile/global.anthropic.claude-opus-4-6-v1"
	}

	cfg, err := awsconfig.LoadDefaultConfig(context.Background(),
		awsconfig.WithRegion(region),
		awsconfig.WithCredentialsProvider(
			credentials.NewStaticCredentialsProvider(accessKeyID, secretAccessKey, ""),
		),
	)
	if err != nil {
		return nil, fmt.Errorf("load AWS config: %w", err)
	}

	client := bedrockruntime.NewFromConfig(cfg)
	return &BedrockClaudeClient{client: client, model: model}, nil
}

func (c *BedrockClaudeClient) Name() string { return "bedrock-claude" }

func (c *BedrockClaudeClient) Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
	body := bedrockRequest{
		AnthropicVersion: "bedrock-2023-05-31",
		MaxTokens:        req.MaxTokens,
		System:           req.SystemPrompt,
		Temperature:      req.Temperature,
		Messages: []bedrockMessage{
			{Role: "user", Content: req.UserMessage},
		},
	}

	jsonBody, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}

	output, err := c.client.InvokeModel(ctx, &bedrockruntime.InvokeModelInput{
		ModelId:     &c.model,
		ContentType: strPtr("application/json"),
		Accept:      strPtr("application/json"),
		Body:        jsonBody,
	})
	if err != nil {
		return nil, fmt.Errorf("invoke Bedrock model: %w", err)
	}

	var resp bedrockResponse
	if err := json.Unmarshal(output.Body, &resp); err != nil {
		return nil, fmt.Errorf("unmarshal response: %w", err)
	}

	var content string
	for _, block := range resp.Content {
		if block.Type == "text" {
			content += block.Text
		}
	}

	return &ChatResponse{
		Content:      content,
		Model:        resp.Model,
		InputTokens:  resp.Usage.InputTokens,
		OutputTokens: resp.Usage.OutputTokens,
	}, nil
}

func strPtr(s string) *string { return &s }

// Bedrock Anthropic Messages API 结构

type bedrockMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type bedrockRequest struct {
	AnthropicVersion string           `json:"anthropic_version"`
	MaxTokens        int              `json:"max_tokens"`
	System           string           `json:"system,omitempty"`
	Temperature      float64          `json:"temperature,omitempty"`
	Messages         []bedrockMessage `json:"messages"`
}

type bedrockResponse struct {
	Model   string `json:"model"`
	Content []struct {
		Type string `json:"type"`
		Text string `json:"text"`
	} `json:"content"`
	Usage struct {
		InputTokens  int `json:"input_tokens"`
		OutputTokens int `json:"output_tokens"`
	} `json:"usage"`
}