// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"time"

	"github.com/dashscope/dashscope-confidential-maas/sdk"
)

func main() {
	// Command line flags
	endpoint := flag.String("endpoint", "", "ModelStudio WebSocket endpoint (required)")
	model := flag.String("model", "", "Model name (required)")
	prompt := flag.String("prompt", "", "Prompt to send (required)")
	stream := flag.Bool("stream", false, "Enable streaming")

	flag.Parse()

	// Validate required parameters
	if *endpoint == "" {
		fmt.Fprintln(os.Stderr, "Error: --endpoint is required")
		flag.Usage()
		os.Exit(1)
	}
	if *model == "" {
		fmt.Fprintln(os.Stderr, "Error: --model is required")
		flag.Usage()
		os.Exit(1)
	}
	if *prompt == "" {
		fmt.Fprintln(os.Stderr, "Error: --prompt is required")
		flag.Usage()
		os.Exit(1)
	}

	// Read API key from environment variable
	apiKey := os.Getenv("DASHSCOPE_API_KEY")
	if apiKey == "" {
		fmt.Fprintln(os.Stderr, "Error: DASHSCOPE_API_KEY environment variable not set")
		os.Exit(1)
	}

	// Create client
	client, err := sdk.NewClient(&sdk.Options{
		Endpoint: *endpoint,
		APIKey:   apiKey,
		Model:    *model,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to create client: %v\n", err)
		os.Exit(1)
	}
	defer client.Close()

	// Create inference request
	req := &sdk.ChatCompletionRequest{
		Messages: []sdk.Message{
			{Role: "user", Content: *prompt},
		},
		Stream:      *stream,
		Temperature: sdk.Float(0.7),
		MaxTokens:   sdk.Int(2048),
		TopP:        sdk.Float(0.9),
	}

	ctx := context.Background()

	// Perform inference
	completion, err := client.Chat(ctx, req)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	// Print attestation result (from completion)
	printAttestationResult(completion.Attestation)

	// Handle response based on streaming mode
	if *stream {
		// Streaming mode: read from chunks channel
		fmt.Print("Assistant: ")
		for chunk := range completion.Chunks {
			fmt.Print(chunk)
		}
		fmt.Println()

		// Check for errors
		if completion.Error != nil {
			fmt.Fprintf(os.Stderr, "\nError: %v\n", completion.Error)
			os.Exit(1)
		}

		// Print usage statistics
		fmt.Printf("\nTokens: %d (prompt: %d, completion: %d)\n",
			completion.Usage.TotalTokens,
			completion.Usage.PromptTokens,
			completion.Usage.CompletionTokens)
	} else {
		// Non-streaming mode: use content directly
		fmt.Printf("Assistant: %s\n", completion.Content)
		fmt.Printf("Tokens: %d (prompt: %d, completion: %d)\n",
			completion.Usage.TotalTokens,
			completion.Usage.PromptTokens,
			completion.Usage.CompletionTokens)
	}
}

// printAttestationResult prints the attestation verification result
func printAttestationResult(result *sdk.AttestationResult) {
	if result == nil {
		return
	}

	fmt.Println("\n=== Attestation Result ===")
	fmt.Printf("Verified at: %s\n", result.VerifiedAt.Format(time.RFC3339))
	fmt.Printf("TEE Type: %s\n", result.Bundle.TEEType)

	detailsJSON, err := result.TEEDetailsJSON()
	if err != nil {
		fmt.Printf("  (no TEE details: %v)\n", err)
	} else {
		var buf bytes.Buffer
		if err := json.Indent(&buf, detailsJSON, "  ", "  "); err != nil {
			fmt.Printf("  Details: %s\n", string(detailsJSON))
		} else {
			fmt.Printf("  Details:\n%s\n", buf.String())
		}
	}
	fmt.Println()
}
