package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"sync"

	"github.com/aws/aws-sdk-go-v2/aws"
	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/bedrock"
	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/document"
	brtypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
	smithy "github.com/aws/smithy-go"
)

// awsCredentialMode is how Bedrock selects AWS credentials.
type awsCredentialMode int

const (
	awsCredentialDefault awsCredentialMode = iota // SDK default chain
	awsCredentialStatic                           // access key + secret
	awsCredentialProfile                          // shared config profile
)

// awsCredentialModeFor applies Bedrock credential precedence:
// static keys, then a named profile, then the default AWS chain.
func awsCredentialModeFor(accessKey, secret, profile string) awsCredentialMode {
	if accessKey != "" && secret != "" {
		return awsCredentialStatic
	}
	if profile != "" {
		return awsCredentialProfile
	}
	return awsCredentialDefault
}

// BedrockProviderConfig holds configuration for creating a BedrockProvider.
type BedrockProviderConfig struct {
	Model              string
	Region             string
	AWSProfile         string
	AWSAccessKeyID     string
	AWSSecretAccessKey string
	AWSSessionToken    string
	ContextWindow      int
	MaxRetries         int
}

// bedrockConverse is the subset of the Bedrock Runtime client the provider
// uses. *bedrockruntime.Client implements it; tests substitute a fake.
type bedrockConverse interface {
	Converse(ctx context.Context, params *bedrockruntime.ConverseInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.ConverseOutput, error)
	ConverseStream(ctx context.Context, params *bedrockruntime.ConverseStreamInput, optFns ...func(*bedrockruntime.Options)) (*bedrockruntime.ConverseStreamOutput, error)
}

// BedrockProvider implements Provider against Amazon Bedrock's Converse API,
// which is one request shape for every model that supports it.
type BedrockProvider struct {
	model         string
	region        string
	profile       string
	accessKey     string
	secret        string
	sessionToken  string
	contextWindow int
	maxRetries    int

	mu       sync.Mutex
	converse bedrockConverse
}

// NewBedrockProvider creates a Bedrock provider. The constructor does not
// talk to AWS; credentials load on the first call.
func NewBedrockProvider(cfg BedrockProviderConfig) *BedrockProvider {
	if cfg.MaxRetries <= 0 {
		cfg.MaxRetries = 3
	}
	if cfg.ContextWindow <= 0 {
		cfg.ContextWindow = 200000
	}
	if cfg.Region == "" {
		cfg.Region = "us-east-1"
	}
	return &BedrockProvider{
		model:         cfg.Model,
		region:        cfg.Region,
		profile:       cfg.AWSProfile,
		accessKey:     cfg.AWSAccessKeyID,
		secret:        cfg.AWSSecretAccessKey,
		sessionToken:  cfg.AWSSessionToken,
		contextWindow: cfg.ContextWindow,
		maxRetries:    cfg.MaxRetries,
	}
}

func (b *BedrockProvider) ModelName() string     { return b.model }
func (b *BedrockProvider) ContextWindow() int    { return b.contextWindow }
func (b *BedrockProvider) SupportsToolUse() bool { return true }

func (b *BedrockProvider) loadAWSConfig(ctx context.Context) (aws.Config, error) {
	var opts []func(*awsconfig.LoadOptions) error
	if b.region != "" {
		opts = append(opts, awsconfig.WithRegion(b.region))
	}
	switch awsCredentialModeFor(b.accessKey, b.secret, b.profile) {
	case awsCredentialStatic:
		opts = append(opts, awsconfig.WithCredentialsProvider(
			credentials.NewStaticCredentialsProvider(b.accessKey, b.secret, b.sessionToken),
		))
	case awsCredentialProfile:
		opts = append(opts, awsconfig.WithSharedConfigProfile(b.profile))
	}
	return awsconfig.LoadDefaultConfig(ctx, opts...)
}

func (b *BedrockProvider) client(ctx context.Context) (bedrockConverse, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.converse != nil {
		return b.converse, nil
	}
	cfg, err := b.loadAWSConfig(ctx)
	if err != nil {
		return nil, fmt.Errorf("loading AWS config for Bedrock: %w", err)
	}
	b.converse = bedrockruntime.NewFromConfig(cfg)
	return b.converse, nil
}

// Complete sends a Converse request and returns the full response.
func (b *BedrockProvider) Complete(ctx context.Context, req CompletionRequest) (*CompletionResponse, error) {
	input, err := buildConverseInput(b.model, req)
	if err != nil {
		return nil, err
	}
	client, err := b.client(ctx)
	if err != nil {
		return nil, err
	}

	var lastErr error
	for attempt := 0; attempt <= b.maxRetries; attempt++ {
		out, err := client.Converse(ctx, input)
		if err == nil {
			return parseConverseOutput(out), nil
		}
		lastErr = err
		if !bedrockRetryable(err) || attempt == b.maxRetries {
			return nil, fmt.Errorf("bedrock converse: %w", err)
		}
		if err := openaiBackoff(ctx, attempt); err != nil {
			return nil, err
		}
	}
	return nil, fmt.Errorf("bedrock converse failed after %d retries: %w", b.maxRetries, lastErr)
}

// Stream sends a ConverseStream request and emits text and tool-call deltas.
func (b *BedrockProvider) Stream(ctx context.Context, req CompletionRequest) (<-chan StreamChunk, error) {
	input, err := buildConverseInput(b.model, req)
	if err != nil {
		return nil, err
	}
	client, err := b.client(ctx)
	if err != nil {
		return nil, err
	}
	out, err := client.ConverseStream(ctx, streamInputFrom(input))
	if err != nil {
		return nil, fmt.Errorf("bedrock converse stream: %w", err)
	}
	stream := out.GetStream()
	ch := make(chan StreamChunk, 100)
	go func() {
		defer close(ch)
		defer stream.Close()
		for event := range stream.Events() {
			switch e := event.(type) {
			case *brtypes.ConverseStreamOutputMemberContentBlockStart:
				if start, ok := e.Value.Start.(*brtypes.ContentBlockStartMemberToolUse); ok {
					ch <- StreamChunk{ToolCallDelta: &ToolCallDelta{
						ID:   aws.ToString(start.Value.ToolUseId),
						Name: aws.ToString(start.Value.Name),
					}}
				}
			case *brtypes.ConverseStreamOutputMemberContentBlockDelta:
				switch d := e.Value.Delta.(type) {
				case *brtypes.ContentBlockDeltaMemberText:
					if d.Value != "" {
						ch <- StreamChunk{Delta: d.Value}
					}
				case *brtypes.ContentBlockDeltaMemberToolUse:
					if d.Value.Input != nil && *d.Value.Input != "" {
						ch <- StreamChunk{ToolCallDelta: &ToolCallDelta{ArgumentsDelta: *d.Value.Input}}
					}
				}
			case *brtypes.ConverseStreamOutputMemberMessageStop:
				ch <- StreamChunk{Done: true}
				return
			}
		}
	}()
	return ch, nil
}

// HealthCheck verifies credentials and that the model id exists, without
// generating tokens. Foundation-model ids are checked first; inference
// profiles (us.anthropic…, global.…, …) fall through to GetInferenceProfile.
func (b *BedrockProvider) HealthCheck(ctx context.Context) error {
	cfg, err := b.loadAWSConfig(ctx)
	if err != nil {
		return fmt.Errorf("bedrock credentials: %w", err)
	}
	if cfg.Region == "" {
		return fmt.Errorf("bedrock region is empty — set orchestrator.region, --aws-region, or AWS_REGION")
	}
	client := bedrock.NewFromConfig(cfg)
	_, err = client.GetFoundationModel(ctx, &bedrock.GetFoundationModelInput{
		ModelIdentifier: aws.String(b.model),
	})
	if err == nil {
		return nil
	}
	if awsAuthError(err) {
		return fmt.Errorf("bedrock credentials rejected: %w", err)
	}
	_, ipErr := client.GetInferenceProfile(ctx, &bedrock.GetInferenceProfileInput{
		InferenceProfileIdentifier: aws.String(b.model),
	})
	if ipErr == nil {
		return nil
	}
	if awsAuthError(ipErr) {
		return fmt.Errorf("bedrock credentials rejected: %w", ipErr)
	}
	return fmt.Errorf("bedrock model %q was not found as a foundation model (%v) or an inference profile (%v)", b.model, err, ipErr)
}

func buildConverseInput(model string, req CompletionRequest) (*bedrockruntime.ConverseInput, error) {
	maxTokens := req.MaxTokens
	if maxTokens <= 0 {
		maxTokens = 8192
	}
	if maxTokens > math.MaxInt32 {
		maxTokens = math.MaxInt32
	}
	temp := float32(req.Temperature)

	messages, err := buildBedrockMessages(req)
	if err != nil {
		return nil, err
	}
	input := &bedrockruntime.ConverseInput{
		ModelId:  aws.String(model),
		Messages: messages,
		System:   bedrockSystem(req),
		InferenceConfig: &brtypes.InferenceConfiguration{
			MaxTokens:   aws.Int32(int32(maxTokens)),
			Temperature: aws.Float32(temp),
		},
	}
	if tools := bedrockTools(req.Tools); len(tools) > 0 {
		input.ToolConfig = &brtypes.ToolConfiguration{Tools: tools}
	}
	return input, nil
}

func streamInputFrom(in *bedrockruntime.ConverseInput) *bedrockruntime.ConverseStreamInput {
	return &bedrockruntime.ConverseStreamInput{
		ModelId:         in.ModelId,
		Messages:        in.Messages,
		System:          in.System,
		InferenceConfig: in.InferenceConfig,
		ToolConfig:      in.ToolConfig,
	}
}

func bedrockSystem(req CompletionRequest) []brtypes.SystemContentBlock {
	var blocks []brtypes.SystemContentBlock
	if req.SystemPrompt != "" {
		blocks = append(blocks, &brtypes.SystemContentBlockMemberText{Value: req.SystemPrompt})
	}
	for _, msg := range req.Messages {
		if msg.Role == "system" && msg.Content != "" {
			blocks = append(blocks, &brtypes.SystemContentBlockMemberText{Value: msg.Content})
		}
	}
	return blocks
}

func buildBedrockMessages(req CompletionRequest) ([]brtypes.Message, error) {
	var out []brtypes.Message
	for i := 0; i < len(req.Messages); i++ {
		msg := req.Messages[i]
		switch msg.Role {
		case "system":
			continue
		case "assistant":
			if msg.Content == "" {
				continue
			}
			out = appendBedrockBlocks(out, brtypes.ConversationRoleAssistant, []brtypes.ContentBlock{
				&brtypes.ContentBlockMemberText{Value: msg.Content},
			})
		case "tool":
			var blocks []brtypes.ContentBlock
			for ; i < len(req.Messages) && req.Messages[i].Role == "tool"; i++ {
				blocks = append(blocks, &brtypes.ContentBlockMemberToolResult{
					Value: brtypes.ToolResultBlock{
						ToolUseId: aws.String(req.Messages[i].ToolCallID),
						Content: []brtypes.ToolResultContentBlock{
							&brtypes.ToolResultContentBlockMemberText{Value: req.Messages[i].Content},
						},
					},
				})
			}
			i--
			out = appendBedrockBlocks(out, brtypes.ConversationRoleUser, blocks)
		default: // user and anything else
			if msg.Content == "" {
				continue
			}
			out = appendBedrockBlocks(out, brtypes.ConversationRoleUser, []brtypes.ContentBlock{
				&brtypes.ContentBlockMemberText{Value: msg.Content},
			})
		}
	}
	if len(out) == 0 {
		return nil, fmt.Errorf("bedrock converse requires at least one user or assistant message")
	}
	return out, nil
}

// appendBedrockBlocks adds blocks to the last message when the role matches,
// so consecutive user turns (including coalesced tool results) stay one message.
func appendBedrockBlocks(out []brtypes.Message, role brtypes.ConversationRole, blocks []brtypes.ContentBlock) []brtypes.Message {
	if len(blocks) == 0 {
		return out
	}
	if n := len(out); n > 0 && out[n-1].Role == role {
		out[n-1].Content = append(out[n-1].Content, blocks...)
		return out
	}
	return append(out, brtypes.Message{Role: role, Content: blocks})
}

func bedrockTools(tools []Tool) []brtypes.Tool {
	var out []brtypes.Tool
	for _, t := range tools {
		schema := any(map[string]any{"type": "object", "properties": map[string]any{}})
		if len(t.Parameters) > 0 {
			var parsed any
			if err := json.Unmarshal(t.Parameters, &parsed); err == nil && parsed != nil {
				schema = parsed
			}
		}
		spec := brtypes.ToolSpecification{
			Name:        aws.String(t.Name),
			InputSchema: &brtypes.ToolInputSchemaMemberJson{Value: document.NewLazyDocument(schema)},
		}
		if t.Description != "" {
			spec.Description = aws.String(t.Description)
		}
		out = append(out, &brtypes.ToolMemberToolSpec{Value: spec})
	}
	return out
}

func parseConverseOutput(out *bedrockruntime.ConverseOutput) *CompletionResponse {
	result := &CompletionResponse{}
	if out == nil {
		return result
	}
	result.StopReason = string(out.StopReason)
	if out.Usage != nil {
		result.Usage = Usage{
			InputTokens:  int(aws.ToInt32(out.Usage.InputTokens)),
			OutputTokens: int(aws.ToInt32(out.Usage.OutputTokens)),
		}
		if out.Usage.CacheReadInputTokens != nil {
			result.Usage.CacheReadInputTokens = int(*out.Usage.CacheReadInputTokens)
		}
		if out.Usage.CacheWriteInputTokens != nil {
			result.Usage.CacheCreationInputTokens = int(*out.Usage.CacheWriteInputTokens)
		}
	}
	msg, ok := out.Output.(*brtypes.ConverseOutputMemberMessage)
	if !ok {
		return result
	}
	for _, block := range msg.Value.Content {
		switch b := block.(type) {
		case *brtypes.ContentBlockMemberText:
			result.Content += b.Value
		case *brtypes.ContentBlockMemberToolUse:
			result.ToolCalls = append(result.ToolCalls, ToolCall{
				ID:        aws.ToString(b.Value.ToolUseId),
				Name:      aws.ToString(b.Value.Name),
				Arguments: documentJSON(b.Value.Input),
			})
		}
	}
	return result
}

func documentJSON(doc document.Interface) string {
	if doc == nil {
		return "{}"
	}
	// MarshalSmithyDocument works for both wire responses and
	// document.NewLazyDocument values. UnmarshalSmithyDocument on a
	// lazy document does not write the caller's destination.
	b, err := doc.MarshalSmithyDocument()
	if err != nil || len(b) == 0 || string(b) == "null" {
		return "{}"
	}
	return string(b)
}

func bedrockRetryable(err error) bool {
	var apiErr smithy.APIError
	if !errors.As(err, &apiErr) {
		return false
	}
	switch apiErr.ErrorCode() {
	case "ThrottlingException", "ServiceUnavailableException", "ModelTimeoutException",
		"InternalServerException", "ModelNotReadyException":
		return true
	default:
		return false
	}
}

func awsAuthError(err error) bool {
	var apiErr smithy.APIError
	if !errors.As(err, &apiErr) {
		return false
	}
	switch apiErr.ErrorCode() {
	case "AccessDeniedException", "UnrecognizedClientException", "ExpiredTokenException",
		"InvalidSignatureException", "IncompleteSignatureException", "SignatureDoesNotMatch",
		"AuthFailure":
		return true
	default:
		return false
	}
}