package main

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

// This program demonstrates a simple streaming client that posts a chat/completions
// request with `stream: true` and prints chunks as they arrive. It avoids using the
// SDK so it's easier to adapt to providers that expose an OpenAI-compatible
// streaming chat/completions endpoint (for example: https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions).
func main() {
	apiKey := os.Getenv("DASHSCOPE_API_KEY")
	if apiKey == "" {
		fmt.Fprintln(os.Stderr, "DASHSCOPE_API_KEY is not set")
		os.Exit(2)
	}

	// Build request body for chat/completions streaming.
	// Adjust model/messages according to the provider's expectations.
	body := map[string]any{
		"model": "qwen-plus",
		"messages": []map[string]string{
			{"role": "user", "content": "你是谁"},
		},
		"stream": true,
	}

	b, _ := json.Marshal(body)

	// Use the chat/completions endpoint which commonly supports streaming.
	url := "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"

	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, url, bytes.NewReader(b))
	if err != nil {
		panic(err)
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: 0}

	resp, err := client.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		// Read body for debugging (limit size to avoid huge dumps)
		data, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		panic(fmt.Errorf("request failed: %s - %s", resp.Status, string(data)))
	}

	// Read streaming response progressively. Many OpenAI-compatible servers
	// use Server-Sent Events (SSE) style streaming where each event line is
	// prefixed with "data: " and terminated by a blank line.

	reader := bufio.NewReader(resp.Body)

	// Set a small context deadline so a stalled stream doesn't hang forever.
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()

	// Read loop: handle lines like "data: {json}\n" or raw JSON chunks.
	for {
		select {
		case <-ctx.Done():
			fmt.Fprintln(os.Stderr, "stream timeout")
			return
		default:
		}

		line, err := reader.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				return
			}
			panic(err)
		}

		line = strings.TrimSpace(line)
		if line == "" {
			// ignore empty lines
			continue
		}

		// SSE-style: "data: [DONE]" or "data: { ... }"
		if strings.HasPrefix(line, "data:") {
			payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if payload == "[DONE]" {
				// stream finished
				return
			}

			// Try to decode JSON and extract text fields commonly used in streams.
			var obj map[string]any
			if err := json.Unmarshal([]byte(payload), &obj); err == nil {
				// Try common paths: choices[0].delta.content or choices[0].text
				if choices, ok := obj["choices"].([]any); ok && len(choices) > 0 {
					if choiceMap, ok := choices[0].(map[string]any); ok {
						// delta.content
						if delta, ok := choiceMap["delta"].(map[string]any); ok {
							if content, ok := delta["content"].(string); ok {
								fmt.Print(content)
								continue
							}
						}
						// text
						if text, ok := choiceMap["text"].(string); ok {
							fmt.Print(text)
							continue
						}
					}
				}

				// Fallback: print the raw JSON payload
				fmt.Println(payload)
			} else {
				// Not JSON, print raw payload
				fmt.Println(payload)
			}
			continue
		}

		// Some providers may stream raw JSON chunks (not prefixed with data:)
		// Try to handle those as well: attempt to parse the line as JSON and
		// print fields if present, otherwise print the raw line.
		var obj map[string]any
		if err := json.Unmarshal([]byte(line), &obj); err == nil {
			if choices, ok := obj["choices"].([]any); ok && len(choices) > 0 {
				if choiceMap, ok := choices[0].(map[string]any); ok {
					if delta, ok := choiceMap["delta"].(map[string]any); ok {
						if content, ok := delta["content"].(string); ok {
							fmt.Print(content)
							continue
						}
					}
					if text, ok := choiceMap["text"].(string); ok {
						fmt.Print(text)
						continue
					}
				}
			}
			// fallback print full JSON
			fmt.Println(line)
		} else {
			fmt.Println(line)
		}
	}
}