package main

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/config"
)

// bedrock endpoint
const endpointStream = "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/invoke-with-response-stream"

const endpoint = "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/invoke"

// a simple prompt for anthropic.claude2 model
const prompt = `{"prompt": "\n\nHuman: how to cook chicken soup? \n\nAssistant: ", "max_tokens_to_sample": 2048}`

type ResponseStream struct {
	Bytes string `json:"bytes"`
}

func callBedRockBatch() {

	// create a hash of the prompt
	hash := sha256.Sum256([]byte(prompt))
	hashString := hex.EncodeToString(hash[:])

	fmt.Println(hashString)

	// create http request to bedrock endpoint
	request, error := http.NewRequest("POST", endpoint, bytes.NewBuffer([]byte(prompt)))

	if error != nil {
		fmt.Println(error)
	}

	// load aws credentials from config
	cfg, error := config.LoadDefaultConfig(context.TODO())

	if error != nil {
		fmt.Println(error)
	}

	credentials, error := cfg.Credentials.Retrieve(context.TODO())

	if error != nil {
		fmt.Println(error)
	}

	// create aws signature v4 signer to sign http requeset
	signer := v4.NewSigner()

	// sign the request
	signer.SignHTTP(
		context.TODO(),
		credentials,
		request,
		hashString,
		*aws.String("bedrock"),
		*aws.String("us-east-1"),
		time.Now(),
	)

	// create http client and send the request
	client := &http.Client{}
	response, error := client.Do(request)

	if error != nil {
		fmt.Println(error)
	}

	fmt.Println(response)

	// read response body
	data, error := io.ReadAll(response.Body)

	if error != nil {
		fmt.Println(error)
	}

	fmt.Println(string(data))

}

func callBedrockStream() {

	// create a hash of the prompt
	hash := sha256.Sum256([]byte(prompt))
	hashString := hex.EncodeToString(hash[:])

	fmt.Println(hashString)

	// create http request to bedrock endpoint
	request, error := http.NewRequest("POST", endpointStream, bytes.NewBuffer([]byte(prompt)))

	if error != nil {
		fmt.Println(error)
	}

	// load aws credentials from config
	cfg, error := config.LoadDefaultConfig(context.TODO())

	if error != nil {
		fmt.Println(error)
	}

	credentials, error := cfg.Credentials.Retrieve(context.TODO())

	if error != nil {
		fmt.Println(error)
	}

	// create aws signature v4 signer to sign http requeset
	signer := v4.NewSigner()

	// sign the request
	signer.SignHTTP(
		context.TODO(),
		credentials,
		request,
		hashString,
		*aws.String("bedrock"),
		*aws.String("us-east-1"),
		time.Now(),
	)

	fmt.Println(request.Header)

	// send http request
	client := &http.Client{}
	response, error := client.Do(request)

	if error != nil {
		fmt.Println(error)
	}

	fmt.Println(response)

	// process the streaming http response chunk by chunk
	var chunk ResponseStream
	buffer := make([]byte, 1024)
	r, _ := regexp.Compile(`\{([^}]+)\}`)

	for {
		n, err := response.Body.Read(buffer)

		if err != nil {
			break
		}

		// fmt.Println(string(buffer[:n]))

		// regex to capture completion
		matches := r.FindStringSubmatch(string(buffer[:n]))

		// fmt.Println(matches[1])

		json.Unmarshal([]byte("{"+matches[1]+"}"), &chunk)

		// fmt.Println(chunk.Bytes)

		// base64 string decoder
		decoded, error := base64.StdEncoding.DecodeString(chunk.Bytes)

		if error != nil {
			fmt.Println(error)
		}

		fmt.Println(string(decoded))

	}

}

func main() {
	fmt.Println("Hello World")

	// callBedrockStream()

	callBedRockBatch()
}